Repository: yarpc/yab
Pull request #312 (2020-12-11 18:01:10)

gRPC: Add max-response-size option to support large responses
This change addresses #311, adds new option `max-response-size` (applicable only for gRPC requests)

codecov[bot]: # [Codecov](https://codecov.io/gh/yarpc/yab/pull/312?src=pr&el=h1) Report
> Merging [#312](https://codecov.io/gh/yarpc/yab/pull/312?src=pr&el=desc) (173725e) into [dev](https://codecov.io/gh/yarpc/yab/commit/a7ef873b4c9525a3a62e70d6816ac24f408576ad?el=desc) (a7ef873) will **increase** coverage by `0.02%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/yarpc/yab/pull/312/graphs/tree.svg?width=650&height=150&src=pr&token=yEgrqUTCuh)](https://codecov.io/gh/yarpc/yab/pull/312?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##              dev     #312      +/-   ##
==========================================
+ Coverage   79.86%   79.89%   +0.02%     
==========================================
  Files          59       59              
  Lines        2856     2860       +4     
==========================================
+ Hits         2281     2285       +4     
  Misses        476      476              
  Partials       99       99              
```


| [Impacted Files](https://codecov.io/gh/yarpc/yab/pull/312?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [options.go](https://codecov.io/gh/yarpc/yab/pull/312/diff?src=pr&el=tree#diff-b3B0aW9ucy5nbw==) | `93.54% <ø> (ø)` | |
| [transport.go](https://codecov.io/gh/yarpc/yab/pull/312/diff?src=pr&el=tree#diff-dHJhbnNwb3J0Lmdv) | `97.46% <100.00%> (+0.03%)` | :arrow_up: |
| [transport/grpc.go](https://codecov.io/gh/yarpc/yab/pull/312/diff?src=pr&el=tree#diff-dHJhbnNwb3J0L2dycGMuZ28=) | `77.94% <100.00%> (+1.01%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/yarpc/yab/pull/312?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/yarpc/yab/pull/312?src=pr&el=footer). Last update [a7ef873...ad2f45e](https://codecov.io/gh/yarpc/yab/pull/312?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

rabbbit: Thanks @jronak, I owe you a beer;)

Patch:
diff --git a/.codecov.yml b/.codecov.yml
index 09c3fc3..a4b9d43 100644
--- a/.codecov.yml
+++ b/.codecov.yml
@@ -13,3 +13,5 @@ coverage:
         if_not_found: success  # if parent is not found report status as success, error, or failure
         if_ci_failed: error    # if ci fails report status as success, error, or failure
 
+ignore:
+ - /testdata/**/*
diff --git a/Makefile b/Makefile
index 16b73c9..a646e85 100644
--- a/Makefile
+++ b/Makefile
@@ -28,5 +28,6 @@ docs:
 
 .PHONY: test_ci
 test_ci: install build
-	go test -coverprofile=cover.out -coverpkg=./... ./...
+	go test -race -coverprofile=cover_temp.out -coverpkg=./... ./...
+	grep -v "^github.com/yarpc/yab/testdata/*" cover_temp.out > cover.out
 	go tool cover -html=cover.out -o cover.html
diff --git a/options.go b/options.go
index 9d9eea2..fc37239 100644
--- a/options.go
+++ b/options.go
@@ -74,16 +74,17 @@ type RequestOptions struct {
 
 // TransportOptions are transport related options.
 type TransportOptions struct {
-	ServiceName      string            `short:"s" long:"service" description:"The TChannel/Hyperbahn service name"`
-	Peers            []string          `short:"p" long:"peer" description:"The host:port of the service to call"`
-	PeerList         string            `short:"P" long:"peer-list" description:"Path or URL of a JSON, YAML, or flat file containing a list of host:ports. -P? for supported protocols."`
-	CallerName       string            `long:"caller" description:"Caller will override the default caller name (which is yab-$USER)."`
-	RoutingKey       string            `long:"rk" description:"The routing key overrides the service name traffic group for proxies."`
-	RoutingDelegate  string            `long:"rd" description:"The routing delegate overrides the routing key traffic group for proxies."`
-	ShardKey         string            `long:"sk" description:"The shard key is a transport header that clues where to send a request within a clustered traffic group."`
-	Jaeger           bool              `long:"jaeger" description:"Use the Jaeger tracing client to send Uber style traces and baggage headers"`
-	TransportHeaders map[string]string `short:"T" long:"topt" description:"Transport options for TChannel, protocol headers for HTTP"`
-	HTTPMethod       string            `long:"http-method" description:"The HTTP method to use"`
+	ServiceName         string            `short:"s" long:"service" description:"The TChannel/Hyperbahn service name"`
+	Peers               []string          `short:"p" long:"peer" description:"The host:port of the service to call"`
+	PeerList            string            `short:"P" long:"peer-list" description:"Path or URL of a JSON, YAML, or flat file containing a list of host:ports. -P? for supported protocols."`
+	CallerName          string            `long:"caller" description:"Caller will override the default caller name (which is yab-$USER)."`
+	RoutingKey          string            `long:"rk" description:"The routing key overrides the service name traffic group for proxies."`
+	RoutingDelegate     string            `long:"rd" description:"The routing delegate overrides the routing key traffic group for proxies."`
+	ShardKey            string            `long:"sk" description:"The shard key is a transport header that clues where to send a request within a clustered traffic group."`
+	Jaeger              bool              `long:"jaeger" description:"Use the Jaeger tracing client to send Uber style traces and baggage headers"`
+	TransportHeaders    map[string]string `short:"T" long:"topt" description:"Transport options for TChannel, protocol headers for HTTP"`
+	HTTPMethod          string            `long:"http-method" description:"The HTTP method to use"`
+	GRPCMaxResponseSize int               `long:"grpc-max-response-size" description:"Maximum response size for gRPC requests. Default value is 4MB"`
 
 	// This is a hack to work around go-flags not allowing disabling flags:
 	// https://github.com/jessevdk/go-flags/issues/191
diff --git a/transport.go b/transport.go
index 95d8bd6..cef4f17 100644
--- a/transport.go
+++ b/transport.go
@@ -171,6 +171,7 @@ func getTransport(opts TransportOptions, resolved resolvedProtocolEncoding, trac
 			Encoding:        resolved.enc.String(),
 			RoutingKey:      opts.RoutingKey,
 			RoutingDelegate: opts.RoutingDelegate,
+			MaxResponseSize: opts.GRPCMaxResponseSize,
 		})
 	}
 
diff --git a/transport/grpc.go b/transport/grpc.go
index 2be55ce..5a2c247 100644
--- a/transport/grpc.go
+++ b/transport/grpc.go
@@ -53,6 +53,7 @@ type GRPCOptions struct {
 	Encoding        string
 	RoutingKey      string
 	RoutingDelegate string
+	MaxResponseSize int
 }
 
 // NewGRPC returns a transport that calls a GRPC service.
@@ -81,7 +82,12 @@ func newGRPC(options GRPCOptions) (*grpcTransport, error) {
 		return nil, errGRPCNoCaller
 	}
 
-	transport := grpc.NewTransport(grpc.Tracer(options.Tracer))
+	transportOptions := []grpc.TransportOption{grpc.Tracer(options.Tracer)}
+	if options.MaxResponseSize > 0 {
+		transportOptions = append(transportOptions, grpc.ClientMaxRecvMsgSize(options.MaxResponseSize))
+	}
+
+	transport := grpc.NewTransport(transportOptions...)
 	outbound := transport.NewOutbound(peer.Bind(roundrobin.New(transport), peer.BindPeers(peersToIdentifiers(options.Addresses))))
 	if err := transport.Start(); err != nil {
 		return nil, err


Test patch:
diff --git a/transport/grpc_test.go b/transport/grpc_test.go
index 65087ab..9ecdf4f 100644
--- a/transport/grpc_test.go
+++ b/transport/grpc_test.go
@@ -26,6 +26,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"strings"
 	"testing"
 
 	"github.com/opentracing/opentracing-go"
@@ -104,7 +105,7 @@ func TestGRPCSuccess(t *testing.T) {
 		testBarResponse := &testBarResponse{}
 		require.NoError(t, json.Unmarshal(response.Body, testBarResponse))
 		require.Equal(t, "hello", testBarResponse.One)
-	})
+	}, 0)
 }
 
 func TestGRPCError(t *testing.T) {
@@ -115,12 +116,36 @@ func TestGRPCError(t *testing.T) {
 		require.NoError(t, err)
 		_, err = grpcTestEnv.Transport.Call(context.Background(), request)
 		require.Equal(t, yarpcerrors.UnknownErrorf("hello"), err)
+	}, 0)
+}
+
+func TestGRPCMaxResponseSize(t *testing.T) {
+	t.Run("With default max response size", func(t *testing.T) {
+		doWithGRPCTestEnv(t, "example-caller", 1, []transport.Procedure{
+			newTestJSONProcedure("example", "Foo::Bar", testLargeResponse),
+		}, func(t *testing.T, grpcTestEnv *grpcTestEnv) {
+			request, err := newTestJSONRequest("example", "Foo::Bar", &testBarRequest{One: "hello", Size: 1024 * 1024 * 4})
+			require.NoError(t, err)
+			_, err = grpcTestEnv.Transport.Call(context.Background(), request)
+			require.EqualError(t, err, "code:resource-exhausted message:grpc: received message larger than max (4194315 vs. 4194304)")
+		}, 0)
+	})
+	t.Run("With custom max response size", func(t *testing.T) {
+		doWithGRPCTestEnv(t, "example-caller", 1, []transport.Procedure{
+			newTestJSONProcedure("example", "Foo::Bar", testLargeResponse),
+		}, func(t *testing.T, grpcTestEnv *grpcTestEnv) {
+			request, err := newTestJSONRequest("example", "Foo::Bar", &testBarRequest{One: "hello", Size: 1024 * 1024 * 4})
+			require.NoError(t, err)
+			_, err = grpcTestEnv.Transport.Call(context.Background(), request)
+			require.NoError(t, err)
+		}, 1024*1025*4)
 	})
 }
 
 type testBarRequest struct {
 	One   string
 	Error string
+	Size  int
 }
 
 type testBarResponse struct {
@@ -139,6 +164,13 @@ func testBar(ctx context.Context, request *testBarRequest) (*testBarResponse, er
 	}, nil
 }
 
+func testLargeResponse(ctx context.Context, request *testBarRequest) (*testBarResponse, error) {
+	return &testBarResponse{
+		// filling One with non-zero value to avoid html-escape in json encoder
+		One: strings.Repeat("a", request.Size),
+	}, nil
+}
+
 func newTestJSONRequest(service string, method string, request interface{}) (*Request, error) {
 	body, err := json.Marshal(request)
 	if err != nil {
@@ -163,8 +195,9 @@ func doWithGRPCTestEnv(
 	numInbounds int,
 	procedures []transport.Procedure,
 	f func(*testing.T, *grpcTestEnv),
+	maxResponseSize int,
 ) {
-	grpcTestEnv, err := newGRPCTestEnv(caller, numInbounds, procedures)
+	grpcTestEnv, err := newGRPCTestEnv(caller, numInbounds, procedures, maxResponseSize)
 	require.NoError(t, err)
 	defer func() {
 		assert.NoError(t, grpcTestEnv.Close())
@@ -183,8 +216,10 @@ func newGRPCTestEnv(
 	caller string,
 	numInbounds int,
 	procedures []transport.Procedure,
+	maxResponseSize int,
 ) (_ *grpcTestEnv, err error) {
-	yarpcTransport := grpc.NewTransport()
+	options := []grpc.TransportOption{grpc.ServerMaxSendMsgSize(1024 * 1024 * 10)}
+	yarpcTransport := grpc.NewTransport(options...)
 	if err := yarpcTransport.Start(); err != nil {
 		return nil, err
 	}
@@ -216,10 +251,11 @@ func newGRPCTestEnv(
 	}
 
 	transport, err := NewGRPC(GRPCOptions{
-		Addresses: addresses,
-		Tracer:    opentracing.NoopTracer{},
-		Caller:    caller,
-		Encoding:  "json",
+		Addresses:       addresses,
+		Tracer:          opentracing.NoopTracer{},
+		Caller:          caller,
+		Encoding:        "json",
+		MaxResponseSize: maxResponseSize,
 	})
 	if err != nil {
 		return nil, err