Repository: yarpc/yab
Pull request #336 (2021-04-21 08:15:03)

Delay client send stream closure
Adds option `delay-close-send-stream` which accepts time duration to delay the closure of the send stream in the client streaming, and bidirectional streaming methods. 

Currently, Yab closes the send stream once all the client requests have been dispatched which causes certain bidirectional streaming servers to stop the stream entirely. With this option, once all the client requests have been dispatched we can delay the closure of send stream. 

This is particularly useful for Adhoc requests where servers stop sending the messages once send stream is closed. 

For example, the client opens a bidirectional stream to watch the user activity of two users. But the server closes the stream if the client-side send stream is stopped.

Earlier:
```zsh
yab --timeout=1m  --service "test-service" --method uber.yarpc.encoding.protobuf.Test/Stream -p grpc://localhost:5555   -r '{"watch": "user-1"} {"watch": "user-2"}'

// server doesn't send any message as the client has ended the send stream after sending the above two messages.
```

Now:
```zsh
yab --timeout=1m  --service "test-service" --method uber.yarpc.encoding.protobuf.Test/Stream -p grpc://localhost:5555   -r '{"watch": "user-1"} {"watch": "user-2"}' --delay-close-send-stream=1m

{"user": "user-1", "activity": "LOGIN"}
.
.
.
// server continues to stream responses until send stream is open
```



codecov[bot]: # [Codecov](https://codecov.io/gh/yarpc/yab/pull/336?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=yarpc) Report
> Merging [#336](https://codecov.io/gh/yarpc/yab/pull/336?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=yarpc) (4cb4e7d) into [dev](https://codecov.io/gh/yarpc/yab/commit/fd09a634d96ee3985e3ea65aecb8f2f9f99ad225?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=yarpc) (fd09a63) will **decrease** coverage by `0.09%`.
> The diff coverage is `83.33%`.

[![Impacted file tree graph](https://codecov.io/gh/yarpc/yab/pull/336/graphs/tree.svg?width=650&height=150&src=pr&token=yEgrqUTCuh&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=yarpc)](https://codecov.io/gh/yarpc/yab/pull/336?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=yarpc)

```diff
@@            Coverage Diff             @@
##              dev     #336      +/-   ##
==========================================
- Coverage   90.49%   90.39%   -0.10%     
==========================================
  Files          54       54              
  Lines        2671     2675       +4     
==========================================
+ Hits         2417     2418       +1     
- Misses        179      181       +2     
- Partials       75       76       +1     
```


| [Impacted Files](https://codecov.io/gh/yarpc/yab/pull/336?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=yarpc) | Coverage Δ | |
|---|---|---|
| [options.go](https://codecov.io/gh/yarpc/yab/pull/336/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=yarpc#diff-b3B0aW9ucy5nbw==) | `93.54% <ø> (ø)` | |
| [handler.go](https://codecov.io/gh/yarpc/yab/pull/336/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=yarpc#diff-aGFuZGxlci5nbw==) | `82.41% <83.33%> (-1.30%)` | :arrow_down: |

------

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


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b641f04..a721714 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@ Changelog
 =========
 
 # Unreleased
+* Add `stream-delay-close-send` option which delays client send stream closure.
 * New: gRPC details are now printed along with the error if there are any.
 
 # 0.19.1 (2021-04-02)
diff --git a/README.md b/README.md
index 8eedbff..f955925 100644
--- a/README.md
+++ b/README.md
@@ -73,7 +73,11 @@ Request Options:
       --multiplexed-thrift       Enables the Thrift TMultiplexedProtocol used
                                  by services that host multiple Thrift services
                                  on a single endpoint.
-      --stream-interval=         Interval between consecutive stream message sends, applicable separately to every stream request opened on a connection.
+      --stream-interval=         Interval between consecutive stream message sends,
+                                 applicable separately to every stream request
+                                 opened on a connection.
+      --stream-delay-close-send= Delay the closure of send stream once all the
+                                 request messages have been sent.
 
 Transport Options:
   -s, --service=                 The TChannel/Hyperbahn service name
diff --git a/handler.go b/handler.go
index 6888023..61cdd56 100644
--- a/handler.go
+++ b/handler.go
@@ -212,7 +212,7 @@ func makeClientStream(ctx context.Context, stream *yarpctransport.ClientStream,
 	}
 
 	if err == io.EOF {
-		err = closeSendStream(ctx, stream)
+		err = closeSendStream(ctx, stream, opts.DelayCloseSendStream.Duration())
 	}
 	if err != nil {
 		return err
@@ -241,7 +241,7 @@ func makeBidiStream(ctx context.Context, cancel context.CancelFunc, stream *yarp
 			var reqBody []byte
 			reqBody, err = streamIO.NextRequest()
 			if err == io.EOF {
-				err = closeSendStream(ctx, stream)
+				err = closeSendStream(ctx, stream, opts.DelayCloseSendStream.Duration())
 				break
 			}
 			if err != nil {
@@ -317,8 +317,16 @@ func receiveStreamMessage(ctx context.Context, stream *yarpctransport.ClientStre
 }
 
 // closeSendStream closes the stream from the client side while
-// stream can continue to receive messages from server
-func closeSendStream(ctx context.Context, stream *yarpctransport.ClientStream) error {
+// stream can continue to receive messages from server. If non-zero delay is
+// passed, stream is closed after the delay.
+func closeSendStream(ctx context.Context, stream *yarpctransport.ClientStream, delayCloseSendStream time.Duration) error {
+	if delayCloseSendStream != 0 {
+		select {
+		case <-time.After(delayCloseSendStream):
+		case <-ctx.Done():
+		}
+	}
+
 	// YARPC stream.Close method internally invokes closeSend on gRPC clientStream.
 	if err := stream.Close(ctx); err != nil {
 		return fmt.Errorf("Failed to close send stream: %v", err)
diff --git a/options.go b/options.go
index 8513ad3..bb03194 100644
--- a/options.go
+++ b/options.go
@@ -76,7 +76,8 @@ type RequestOptions struct {
 
 // StreamRequestOptions are stream request related options
 type StreamRequestOptions struct {
-	Interval timeMillisFlag `long:"stream-interval" description:"Interval between consecutive stream message sends, applicable separately to every stream request opened on a connection."`
+	Interval             timeMillisFlag `long:"stream-interval" description:"Interval between consecutive stream message sends, applicable separately to every stream request opened on a connection."`
+	DelayCloseSendStream timeMillisFlag `long:"stream-delay-close-send" description:"Delay the closure of send stream once all the request messages have been sent."`
 }
 
 // TransportOptions are transport related options.
diff --git a/protobuf/source_reflection.go b/protobuf/source_reflection.go
index ea02445..dda4685 100644
--- a/protobuf/source_reflection.go
+++ b/protobuf/source_reflection.go
@@ -64,14 +64,17 @@ func NewDescriptorProviderReflection(args ReflectionArgs) (DescriptorProvider, e
 		routingHeaders.Append(ygrpc.RoutingKeyHeader, args.RoutingKey)
 	}
 
-	metadataContext := metadata.NewOutgoingContext(context.Background(), routingHeaders)
+	ctx, cancel := context.WithTimeout(context.Background(), args.Timeout)
+	metadataContext := metadata.NewOutgoingContext(ctx, routingHeaders)
 	return &grpcreflectSource{
-		client: grpcreflect.NewClient(metadataContext, pbClient),
+		client:     grpcreflect.NewClient(metadataContext, pbClient),
+		cancelFunc: cancel,
 	}, nil
 }
 
 type grpcreflectSource struct {
-	client *grpcreflect.Client
+	client     *grpcreflect.Client
+	cancelFunc context.CancelFunc
 }
 
 func (s *grpcreflectSource) FindMessage(messageType string) (*desc.MessageDescriptor, error) {
@@ -83,19 +86,23 @@ func (s *grpcreflectSource) FindMessage(messageType string) (*desc.MessageDescri
 		return nil, nil
 	}
 
+	if err != nil {
+		return nil, wrapReflectionError(err)
+	}
+
 	return msg, err
 }
 
 func (s *grpcreflectSource) FindService(fullyQualifiedName string) (*desc.ServiceDescriptor, error) {
 	service, err := s.client.ResolveService(fullyQualifiedName)
 	if err != nil {
-		available, availableErr := s.client.ListServices()
-		if availableErr != nil {
-			err = availableErr
+		if !grpcreflect.IsElementNotFoundError(err) {
+			return nil, wrapReflectionError(err)
 		}
 
-		if !grpcreflect.IsElementNotFoundError(err) {
-			return nil, err
+		available, availableErr := s.client.ListServices()
+		if availableErr != nil && !grpcreflect.IsElementNotFoundError(availableErr) {
+			return nil, wrapReflectionError(availableErr)
 		}
 
 		return nil, encodingerror.NotFound{
@@ -111,5 +118,10 @@ func (s *grpcreflectSource) FindService(fullyQualifiedName string) (*desc.Servic
 }
 
 func (s *grpcreflectSource) Close() {
+	s.cancelFunc()
 	s.client.Reset()
 }
+
+func wrapReflectionError(err error) error {
+	return fmt.Errorf("error in protobuf reflection: %v", err)
+}


Test patch:
diff --git a/handler_test.go b/handler_test.go
index 79240ad..ed146ce 100644
--- a/handler_test.go
+++ b/handler_test.go
@@ -31,6 +31,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/yarpc/yab/encoding"
+	"go.uber.org/yarpc/api/transport"
 )
 
 type mockStreamReader struct {
@@ -152,3 +153,70 @@ func TestIntervalWaiter(t *testing.T) {
 		}
 	})
 }
+
+type mockStreamCloser struct {
+	closeErr error
+}
+
+func (m mockStreamCloser) Close(context.Context) error { return m.closeErr }
+
+func (mockStreamCloser) Context() context.Context { return context.Background() }
+
+func (mockStreamCloser) Request() *transport.StreamRequest { return nil }
+
+func (mockStreamCloser) SendMessage(context.Context, *transport.StreamMessage) error { return nil }
+
+func (mockStreamCloser) ReceiveMessage(context.Context) (*transport.StreamMessage, error) {
+	return nil, nil
+}
+
+func TestCloseSendStream(t *testing.T) {
+	tests := []struct {
+		description string
+		delay       time.Duration
+		err         error
+	}{
+		{
+			description: "close without delay",
+		},
+		{
+			description: "close with error",
+			err:         errors.New("test error"),
+		},
+		{
+			description: "close with delay",
+			delay:       time.Millisecond * 50,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.description, func(t *testing.T) {
+			stream, err := transport.NewClientStream(mockStreamCloser{closeErr: test.err})
+			require.NoError(t, err, "unexpected client stream error")
+
+			start := time.Now()
+			err = closeSendStream(context.Background(), stream, test.delay)
+			end := time.Now()
+
+			if test.err != nil {
+				require.Error(t, err, "unexpected nil error from close")
+				assert.Contains(t, err.Error(), test.err.Error(), "unexpected error message from close")
+			} else {
+				assert.NoError(t, err, "unexpected error from close")
+			}
+
+			assert.True(t, end.After(start.Add(test.delay)), "unexpected execution time")
+		})
+	}
+
+	t.Run("must not delay when context is done", func(t *testing.T) {
+		stream, err := transport.NewClientStream(mockStreamCloser{})
+		require.NoError(t, err, "unexpected client stream error")
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		start := time.Now()
+		require.NoError(t, closeSendStream(ctx, stream, time.Second), "unexpected error from close")
+		assert.True(t, time.Now().Before(start.Add(time.Second)), "expected close to execute immediately")
+	})
+}
diff --git a/integration_test.go b/integration_test.go
index f33605c..2ce1dba 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -55,6 +55,7 @@ import (
 	ytchan "go.uber.org/yarpc/transport/tchannel"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
+	rpb "google.golang.org/grpc/reflection/grpc_reflection_v1alpha"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -1015,72 +1016,113 @@ test: 1
 	}
 }
 
-func TestGRPCStreamWithStreamIntervalOption(t *testing.T) {
-	t.Run("client stream", func(t *testing.T) {
-		addr, server := setupGRPCServer(t, &simpleService{
+func TestGRPCStreamWithBoundedExecutionTime(t *testing.T) {
+	tests := []struct {
+		desc         string
+		opts         Options
+		wantExecTime time.Duration
+
+		returnOutput  []simple.Foo
+		expectedInput []simple.Foo
+	}{
+		{
+			desc: "client side streaming with stream interval option",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ClientStream",
+					Timeout:           timeMillisFlag(time.Second * 5),
+					RequestJSON:       `{"test":1}{"test":2}{"test":-1}`,
+					StreamRequestOptions: StreamRequestOptions{
+						Interval: timeMillisFlag(time.Millisecond * 100),
+					},
+				},
+			},
 			returnOutput:  []simple.Foo{{Test: 4}},
 			expectedInput: []simple.Foo{{Test: 1}, {Test: 2}, {Test: -1}},
-		})
-		defer server.Stop()
-
-		opts := Options{
-			ROpts: RequestOptions{
-				FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
-				Procedure:         "Bar/ClientStream",
-				Timeout:           timeMillisFlag(time.Second * 5),
-				RequestJSON:       `{"test":1}{"test":2}{"test":-1}`,
-				StreamRequestOptions: StreamRequestOptions{
-					Interval: timeMillisFlag(time.Millisecond * 1000),
+			// Test must run for more than 200ms as there are three requests and
+			// it must take 200ms totally between consecutive messages.
+			wantExecTime: time.Millisecond * 200,
+		},
+		{
+			desc: "bidirectional streaming with stream interval option",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/BidiStream",
+					Timeout:           timeMillisFlag(time.Second * 5),
+					RequestJSON:       `{"test":1}{"test":2}{"test":-1}`,
+					StreamRequestOptions: StreamRequestOptions{
+						Interval: timeMillisFlag(time.Millisecond * 100),
+					},
 				},
 			},
-			TOpts: TransportOptions{
-				ServiceName: "foo",
-				Peers:       []string{"grpc://" + addr.String()},
-			},
-		}
-
-		now := time.Now()
-		_, gotErr := runTestWithOpts(opts)
-		duration := time.Now().Sub(now)
-
-		assert.Empty(t, gotErr)
-		// test must run for more than 2 seconds as there are three requests and
-		// it must take 2 seconds totally between consecutive messages
-		assert.True(t, duration > (time.Second*2), "Unexpected time taken on wait: %v", duration)
-	})
-
-	t.Run("bidirectional stream", func(t *testing.T) {
-		addr, server := setupGRPCServer(t, &simpleService{
 			returnOutput:  []simple.Foo{{Test: 4}, {Test: 9}},
 			expectedInput: []simple.Foo{{Test: 1}, {Test: 2}, {Test: -1}},
-		})
-		defer server.Stop()
-
-		opts := Options{
-			ROpts: RequestOptions{
-				FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
-				Procedure:         "Bar/BidiStream",
-				Timeout:           timeMillisFlag(time.Second * 5),
-				RequestJSON:       `{"test":1}{"test":2}{"test":-1}`,
-				StreamRequestOptions: StreamRequestOptions{
-					Interval: timeMillisFlag(time.Millisecond * 1000),
+			// Test must run for more than 200ms as there are three requests and
+			// it must take 200ms totally between consecutive messages.
+			wantExecTime: time.Millisecond * 200,
+		},
+		{
+			desc: "client side streaming with close send stream delay",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/ClientStream",
+					Timeout:           timeMillisFlag(time.Second * 5),
+					RequestJSON:       `{"test":1}{"test":2}{"test":-1}`,
+					StreamRequestOptions: StreamRequestOptions{
+						DelayCloseSendStream: timeMillisFlag(time.Millisecond * 100),
+					},
 				},
 			},
-			TOpts: TransportOptions{
-				ServiceName: "foo",
-				Peers:       []string{"grpc://" + addr.String()},
+			returnOutput:  []simple.Foo{{Test: 4}},
+			expectedInput: []simple.Foo{{Test: 1}, {Test: 2}, {Test: -1}},
+			// Test must run for more than 100ms as send stream closure has a delay
+			// of 100ms.
+			wantExecTime: time.Millisecond * 100,
+		},
+		{
+			desc: "bidirectional streaming with close send stream delay",
+			opts: Options{
+				ROpts: RequestOptions{
+					FileDescriptorSet: []string{"testdata/protobuf/simple/simple.proto.bin"},
+					Procedure:         "Bar/BidiStream",
+					Timeout:           timeMillisFlag(time.Second * 5),
+					RequestJSON:       `{"test":1}{"test":2}`,
+					StreamRequestOptions: StreamRequestOptions{
+						DelayCloseSendStream: timeMillisFlag(time.Millisecond * 100),
+					},
+				},
 			},
-		}
+			returnOutput:  []simple.Foo{{Test: 4}, {Test: 9}},
+			expectedInput: []simple.Foo{{Test: 1}, {Test: 2}, {Test: -1}},
+			// Test must run for more than 100ms as send stream closure has a delay
+			// of 100ms.
+			wantExecTime: time.Millisecond * 100,
+		},
+	}
 
-		now := time.Now()
-		_, gotErr := runTestWithOpts(opts)
-		duration := time.Now().Sub(now)
+	for _, tt := range tests {
+		t.Run(tt.desc, func(t *testing.T) {
+			addr, server := setupGRPCServer(t, &simpleService{
+				expectedInput: tt.expectedInput,
+				returnOutput:  tt.returnOutput,
+			})
+			defer server.Stop()
 
-		assert.Empty(t, gotErr)
-		// test must run for more than 2 seconds as there are three requests and
-		// it must take 2 seconds totally between consecutive messages
-		assert.True(t, duration > (time.Second*2), "Unexpected time taken on wait: %v", duration)
-	})
+			tt.opts.TOpts = TransportOptions{
+				ServiceName: "foo",
+				Peers:       []string{"grpc://" + addr.String()},
+			}
+			now := time.Now()
+			_, gotErr := runTestWithOpts(tt.opts)
+			duration := time.Now().Sub(now)
+
+			assert.Empty(t, gotErr)
+			assert.True(t, duration > tt.wantExecTime, "Unexpected time taken: %v", duration)
+		})
+	}
 }
 
 func TestGRPCReflectionSource(t *testing.T) {
@@ -1185,3 +1227,46 @@ func TestGRPCReflectionSource(t *testing.T) {
 		})
 	}
 }
+
+type protoReflectService struct {
+	waitChan chan struct{} // If non-nil, handler waits on the channel
+}
+
+func (p protoReflectService) ServerReflectionInfo(s rpb.ServerReflection_ServerReflectionInfoServer) error {
+	if p.waitChan != nil {
+		// Wait on the given channel or until stream context completes.
+		select {
+		case <-s.Context().Done():
+		case <-p.waitChan:
+		}
+	}
+	return nil
+}
+
+func TestGPCReflectionTimeout(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	waitChan := make(chan struct{})
+	reflectSvc := protoReflectService{waitChan: waitChan}
+	server := grpc.NewServer()
+	rpb.RegisterServerReflectionServer(server, reflectSvc)
+	go server.Serve(ln)
+
+	defer server.Stop()
+	defer close(waitChan)
+
+	gotOut, gotErr := runTestWithOpts(Options{
+		ROpts: RequestOptions{
+			Procedure:   "Bar/Baz",
+			Timeout:     timeMillisFlag(time.Second * 1),
+			RequestJSON: `{"test":0}`,
+		},
+		TOpts: TransportOptions{
+			ServiceName: "foo",
+			Peers:       []string{ln.Addr().String()},
+		},
+	})
+	assert.Empty(t, gotOut, "Expected empty stdout")
+	assert.Contains(t, gotErr, "error in protobuf reflection: rpc error: code = DeadlineExceeded desc = context deadline exceeded", "Expected deadline exceeded error")
+}
diff --git a/protobuf/source_reflection_test.go b/protobuf/source_reflection_test.go
index a33097c..141be5d 100644
--- a/protobuf/source_reflection_test.go
+++ b/protobuf/source_reflection_test.go
@@ -1,6 +1,7 @@
 package protobuf
 
 import (
+	"errors"
 	"net"
 	"testing"
 	"time"
@@ -195,8 +196,42 @@ func TestReflectionRoutingHeaders(t *testing.T) {
 	}
 }
 
+func TestE2eErrors(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer ln.Close()
+
+	s := grpc.NewServer()
+	drs := &dummyReflectionServer{returnErr: errors.New("test error")}
+	rpb.RegisterServerReflectionServer(s, drs)
+	go s.Serve(ln)
+
+	// Ensure that all streams are closed by the end of the test.
+	defer s.GracefulStop()
+
+	source, err := NewDescriptorProviderReflection(ReflectionArgs{
+		Timeout: time.Second,
+		Peers:   []string{ln.Addr().String()},
+	})
+	require.NoError(t, err)
+	defer source.Close()
+
+	t.Run("test error handling in find method", func(t *testing.T) {
+		desc, err := source.FindMessage("test")
+		require.Nil(t, desc, "unexpected message descriptor")
+		assert.EqualError(t, err, "error in protobuf reflection: rpc error: code = Unknown desc = test error")
+	})
+
+	t.Run("test error handling in find service", func(t *testing.T) {
+		desc, err := source.FindService("test")
+		require.Nil(t, desc, "unexpected message descriptor")
+		assert.EqualError(t, err, "error in protobuf reflection: rpc error: code = Unknown desc = test error")
+	})
+}
+
 type dummyReflectionServer struct {
-	md metadata.MD
+	md        metadata.MD
+	returnErr error
 }
 
 func (s *dummyReflectionServer) Reset() {
@@ -204,6 +239,10 @@ func (s *dummyReflectionServer) Reset() {
 }
 
 func (s *dummyReflectionServer) ServerReflectionInfo(r rpb.ServerReflection_ServerReflectionInfoServer) error {
+	if s.returnErr != nil {
+		return s.returnErr
+	}
+
 	if md, ok := metadata.FromIncomingContext(r.Context()); ok {
 		s.md = md
 	}