Repository: yarpc/yab
Pull request #316 (2021-01-09 17:31:25)

gRPC streams support - 3
## Aim
Support gRPC streams in Yab

## Changes
* Add new interface `StreamTransport` which exposes method to call streaming requests
* Add `CallStream` method in gRPC transport 

Next Stack: #317


Patch:
diff --git a/transport/grpc.go b/transport/grpc.go
index 5a2c247..18572cb 100644
--- a/transport/grpc.go
+++ b/transport/grpc.go
@@ -64,6 +64,7 @@ func NewGRPC(options GRPCOptions) (TransportCloser, error) {
 type grpcTransport struct {
 	Transport       transport.Transport
 	Outbound        transport.UnaryOutbound
+	StreamOutbound  transport.StreamOutbound
 	Caller          string
 	Encoding        string
 	RoutingKey      string
@@ -100,6 +101,7 @@ func newGRPC(options GRPCOptions) (*grpcTransport, error) {
 	return &grpcTransport{
 		Transport:       transport,
 		Outbound:        outbound,
+		StreamOutbound:  outbound,
 		Caller:          options.Caller,
 		Encoding:        options.Encoding,
 		RoutingKey:      options.RoutingKey,
@@ -133,10 +135,33 @@ func (t *grpcTransport) Call(ctx context.Context, request *Request) (*Response,
 	return yarpcResponseToResponse(transportResponse)
 }
 
+func (t *grpcTransport) CallStream(ctx context.Context, request *StreamRequest) (*transport.ClientStream, error) {
+	return t.StreamOutbound.CallStream(ctx, t.requestToYARPCStreamRequest(request))
+}
+
 func (t *grpcTransport) Close() error {
 	return multierr.Combine(t.Transport.Stop(), t.Outbound.Stop())
 }
 
+func (t *grpcTransport) requestToYARPCStreamRequest(streamRequest *StreamRequest) *transport.StreamRequest {
+	if t == nil || streamRequest == nil || streamRequest.Request == nil {
+		return nil
+	}
+
+	return &transport.StreamRequest{
+		Meta: &transport.RequestMeta{
+			Caller:          t.Caller,
+			Service:         streamRequest.Request.TargetService,
+			Encoding:        transport.Encoding(t.Encoding),
+			Procedure:       streamRequest.Request.Method,
+			Headers:         transport.HeadersFromMap(streamRequest.Request.Headers),
+			ShardKey:        streamRequest.Request.ShardKey,
+			RoutingKey:      t.RoutingKey,
+			RoutingDelegate: t.RoutingDelegate,
+		},
+	}
+}
+
 func (t *grpcTransport) requestToYARPCRequest(request *Request) *transport.Request {
 	return &transport.Request{
 		Caller:          t.Caller,
diff --git a/transport/interface.go b/transport/interface.go
index ed8eb24..23cc520 100644
--- a/transport/interface.go
+++ b/transport/interface.go
@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/opentracing/opentracing-go"
+	"go.uber.org/yarpc/api/transport"
 	"golang.org/x/net/context"
 )
 
@@ -66,13 +67,19 @@ const (
 )
 
 // Transport defines the interface for the underlying transport over which
-// calls are made.
+// unary calls are made.
 type Transport interface {
 	Call(ctx context.Context, request *Request) (*Response, error)
 	Protocol() Protocol
 	Tracer() opentracing.Tracer
 }
 
+// StreamTransport defines the interface for the underlying transport which
+// supports streaming
+type StreamTransport interface {
+	CallStream(ctx context.Context, request *StreamRequest) (*transport.ClientStream, error)
+}
+
 // TransportCloser is a Transport that can be closed.
 type TransportCloser interface {
 	Transport


Test patch:
diff --git a/transport/grpc_test.go b/transport/grpc_test.go
index 9ecdf4f..ccee3ad 100644
--- a/transport/grpc_test.go
+++ b/transport/grpc_test.go
@@ -21,17 +21,25 @@
 package transport
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net"
 	"strings"
 	"testing"
+	"time"
 
+	"github.com/golang/protobuf/proto"
 	"github.com/opentracing/opentracing-go"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/yarpc/yab/testdata/protobuf/simple"
+	googlegrpc "google.golang.org/grpc"
+
 	"go.uber.org/multierr"
 	"go.uber.org/yarpc/api/transport"
 	yarpcjson "go.uber.org/yarpc/encoding/json"
@@ -95,17 +103,93 @@ func TestGRPCConstructor(t *testing.T) {
 
 func TestGRPCSuccess(t *testing.T) {
 	doWithGRPCTestEnv(t, "example-caller", 5, []transport.Procedure{
-		newTestJSONProcedure("example", "Foo::Bar", testBar),
-	}, func(t *testing.T, grpcTestEnv *grpcTestEnv) {
-		request, err := newTestJSONRequest("example", "Foo::Bar", &testBarRequest{One: "hello"})
-		require.NoError(t, err)
-		response, err := grpcTestEnv.Transport.Call(context.Background(), request)
-		require.NoError(t, err)
-		require.NotNil(t, response)
-		testBarResponse := &testBarResponse{}
-		require.NoError(t, json.Unmarshal(response.Body, testBarResponse))
-		require.Equal(t, "hello", testBarResponse.One)
-	}, 0)
+		newTestJSONProcedure("example", "Foo::Bar", testBar)},
+		func(t *testing.T, grpcTestEnv *grpcTestEnv) {
+			request, err := newTestJSONRequest("example", "Foo::Bar", &testBarRequest{One: "hello"})
+			require.NoError(t, err)
+			response, err := grpcTestEnv.Transport.Call(context.Background(), request)
+			require.NoError(t, err)
+			require.NotNil(t, response)
+			testBarResponse := &testBarResponse{}
+			require.NoError(t, json.Unmarshal(response.Body, testBarResponse))
+			require.Equal(t, "hello", testBarResponse.One)
+		}, 0)
+}
+
+type simpleSvc struct {
+	streamsOpened int
+}
+
+func (s *simpleSvc) Baz(c context.Context, in *simple.Foo) (*simple.Foo, error) {
+	return in, nil
+}
+
+func (s *simpleSvc) BidiStream(stream simple.Bar_BidiStreamServer) error {
+	s.streamsOpened++
+	for {
+		msg, err := stream.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return err
+		}
+		stream.Send(msg)
+	}
+	return nil
+}
+
+func (*simpleSvc) ClientStream(simple.Bar_ClientStreamServer) error {
+	return nil
+}
+
+func (*simpleSvc) ServerStream(*simple.Foo, simple.Bar_ServerStreamServer) error {
+	return nil
+}
+
+func TestGRPCStream(t *testing.T) {
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	server := googlegrpc.NewServer()
+	svc := &simpleSvc{}
+	simple.RegisterBarServer(server, svc)
+	go func() {
+		require.NoError(t, server.Serve(lis))
+	}()
+	defer server.Stop()
+
+	client, err := NewGRPC(GRPCOptions{
+		Addresses: []string{lis.Addr().String()},
+		Tracer:    opentracing.NoopTracer{},
+		Caller:    "test",
+		Encoding:  "proto",
+	})
+	require.NoError(t, err)
+
+	streamClient, ok := client.(StreamTransport)
+	require.True(t, ok)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	stream, err := streamClient.CallStream(ctx, &StreamRequest{
+		Request: &Request{
+			TargetService: "Bar",
+			Method:        "Bar::BidiStream",
+		},
+	})
+	require.NoError(t, err)
+	req, err := proto.Marshal(&simple.Foo{})
+	require.NoError(t, err)
+	err = stream.SendMessage(ctx, &transport.StreamMessage{
+		Body: ioutil.NopCloser(bytes.NewReader(req)),
+	})
+	require.NoError(t, err)
+	msg, err := stream.ReceiveMessage(ctx)
+	require.NoError(t, err)
+	require.NotNil(t, msg)
+	require.NoError(t, stream.Close(ctx))
+	assert.Equal(t, 1, svc.streamsOpened)
 }
 
 func TestGRPCError(t *testing.T) {
diff --git a/transport/tchannel_test.go b/transport/tchannel_test.go
index e5be654..f9f6f5d 100644
--- a/transport/tchannel_test.go
+++ b/transport/tchannel_test.go
@@ -80,7 +80,7 @@ func setupServerAndTransport(t *testing.T, changeOpts ...func(*TChannelOptions))
 	opts := TChannelOptions{
 		SourceService: "yab",
 		TargetService: svr.ServiceName(),
-		Peers:     []string{svr.PeerInfo().HostPort},
+		Peers:         []string{svr.PeerInfo().HostPort},
 		Encoding:      "raw",
 	}