Repository: yarpc/yab
Pull request #320 (2021-02-24 01:37:51)

[gRPC Streams] Add stream-io interface 
* Refactor existing request-supplier and response-handler functions into a separate interface for fetching stream requests and handling stream responses. 


codecov[bot]: # [Codecov](https://codecov.io/gh/yarpc/yab/pull/320?src=pr&el=h1) Report
> Merging [#320](https://codecov.io/gh/yarpc/yab/pull/320?src=pr&el=desc) (3a5f28c) into [dev](https://codecov.io/gh/yarpc/yab/commit/0b7ba29a9a83faf60927e75d363e3793b7bc3754?el=desc) (0b7ba29) will **decrease** coverage by `0.42%`.
> The diff coverage is `81.32%`.

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

```diff
@@            Coverage Diff             @@
##              dev     #320      +/-   ##
==========================================
- Coverage   90.70%   90.27%   -0.43%     
==========================================
  Files          51       52       +1     
  Lines        2441     2592     +151     
==========================================
+ Hits         2214     2340     +126     
- Misses        165      178      +13     
- Partials       62       74      +12     
```


| [Impacted Files](https://codecov.io/gh/yarpc/yab/pull/320?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [handler.go](https://codecov.io/gh/yarpc/yab/pull/320/diff?src=pr&el=tree#diff-aGFuZGxlci5nbw==) | `79.87% <79.87%> (ø)` | |
| [main.go](https://codecov.io/gh/yarpc/yab/pull/320/diff?src=pr&el=tree#diff-bWFpbi5nbw==) | `92.92% <100.00%> (+1.65%)` | :arrow_up: |
| [encoding/thrift\_request.go](https://codecov.io/gh/yarpc/yab/pull/320/diff?src=pr&el=tree#diff-ZW5jb2RpbmcvdGhyaWZ0X3JlcXVlc3QuZ28=) | `100.00% <0.00%> (+1.75%)` | :arrow_up: |
| [encoding/protobuf\_health.go](https://codecov.io/gh/yarpc/yab/pull/320/diff?src=pr&el=tree#diff-ZW5jb2RpbmcvcHJvdG9idWZfaGVhbHRoLmdv) | `83.33% <0.00%> (+4.16%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/yarpc/yab/pull/320?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/320?src=pr&el=footer). Last update [c7bb2f4...3a5f28c](https://codecov.io/gh/yarpc/yab/pull/320?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/handler.go b/handler.go
index 4ca4557..1624e70 100644
--- a/handler.go
+++ b/handler.go
@@ -16,12 +16,16 @@ import (
 	"go.uber.org/zap"
 )
 
-// streamRequestSupplierFn is a function that returns the stream message body
-// if supplier has no more messages it must return io.EOF
-type streamRequestSupplierFn func() (requestBody []byte, err error)
-
-// streamResponseHandlerFn is a function that receives the stream response body
-type streamResponseHandlerFn func(responseBody []byte) error
+// StreamIO defines the interface for fetching stream requests and handling
+// stream responses.
+type StreamIO interface {
+	// NextRequest returns the next stream message body
+	// If there are no more messages it will return EOF.
+	NextRequest() (request []byte, err error)
+
+	// HandleResponse handles the received stream response message.
+	HandleResponse(responseBody []byte) error
+}
 
 type requestHandler struct {
 	out        output
@@ -88,8 +92,10 @@ func (r requestHandler) handleStreamRequest() {
 		r.out.Fatalf("Failed while preparing the request: %v\n", err)
 	}
 
+	streamIO := newStreamIOInitializer(r.out, r.serializer, streamMsgReader)
+
 	if r.shouldMakeInitialRequest() {
-		if err = makeStreamRequest(r.transport, streamReq, r.serializer, streamRequestSupplier(streamMsgReader), r.responseHandler); err != nil {
+		if err = makeStreamRequest(r.transport, streamReq, r.serializer, streamIO); err != nil {
 			r.out.Fatalf("%v\n", err)
 		}
 	}
@@ -102,51 +108,16 @@ func (r requestHandler) shouldMakeInitialRequest() bool {
 	return !(r.opts.BOpts.enabled() && r.opts.BOpts.WarmupRequests == 0)
 }
 
-// responseHandler validates the response bytes and prints indented JSON body
-func (r requestHandler) responseHandler(body []byte) error {
-	res, err := r.serializer.Response(&transport.Response{Body: body})
-	if err != nil {
-		return fmt.Errorf("Failed while serializing stream response: %v", err)
-	}
-
-	bs, err := json.MarshalIndent(res, "", "  ")
-	if err != nil {
-		return fmt.Errorf("Failed to convert map to JSON: %v\nMap: %+v", err, res)
-	}
-
-	r.out.Printf("%s\n\n", bs)
-	return nil
-}
-
 // isStreamingMethod returns true if RPC is streaming type
 func (r requestHandler) isStreamingMethod() bool {
 	return r.serializer.MethodType() != encoding.Unary
 }
 
-// streamRequestSupplier returns a stream request supplier function which uses
-// provided stream request reader
-func streamRequestSupplier(streamMsgReader encoding.StreamRequestReader) streamRequestSupplierFn {
-	nextBodyFn := func() ([]byte, error) {
-		msg, err := streamMsgReader.NextBody()
-		if err == io.EOF {
-			return nil, err
-		}
-		if err != nil {
-			return nil, fmt.Errorf("Failed while reading stream input: %v", err)
-		}
-
-		return msg, nil
-	}
-	return nextBodyFn
-}
-
 // makeStreamRequest opens a stream rpc from the given transport and stream request
 // it then delegates to handler based on rpc type to handle request and response of the stream
 // nextBodyFn is called to get the next stream message body
 // responseHandlerFn is called with the response of the stream
-func makeStreamRequest(t transport.Transport, streamReq *transport.StreamRequest, serializer encoding.Serializer,
-	nextBodyFn streamRequestSupplierFn, responseHandlerFn streamResponseHandlerFn) error {
-
+func makeStreamRequest(t transport.Transport, streamReq *transport.StreamRequest, serializer encoding.Serializer, streamIO StreamIO) error {
 	streamTransport, ok := t.(transport.StreamTransport)
 	if !ok {
 		return fmt.Errorf("Transport does not support stream calls: %q", t.Protocol())
@@ -165,19 +136,18 @@ func makeStreamRequest(t transport.Transport, streamReq *transport.StreamRequest
 
 	switch serializer.MethodType() {
 	case encoding.BidirectionalStream:
-		return makeBidiStream(ctx, stream, nextBodyFn, responseHandlerFn)
+		return makeBidiStream(ctx, stream, streamIO)
 	case encoding.ClientStream:
-		return makeClientStream(ctx, stream, nextBodyFn, responseHandlerFn)
+		return makeClientStream(ctx, stream, streamIO)
 	default:
-		return makeServerStream(ctx, stream, nextBodyFn, responseHandlerFn)
+		return makeServerStream(ctx, stream, streamIO)
 	}
 }
 
 // makeServerStream starts server-side streaming rpc
-func makeServerStream(ctx context.Context, stream *yarpctransport.ClientStream,
-	nextBodyFn streamRequestSupplierFn, responseHandlerFn streamResponseHandlerFn) error {
+func makeServerStream(ctx context.Context, stream *yarpctransport.ClientStream, streamIO StreamIO) error {
 
-	req, err := nextBodyFn()
+	req, err := streamIO.NextRequest()
 	// Use nil body if no initial request input is empty, since request
 	// is mandatory in server streaming rpc.
 	if err != nil && err != io.EOF {
@@ -185,7 +155,7 @@ func makeServerStream(ctx context.Context, stream *yarpctransport.ClientStream,
 	}
 	if err == nil {
 		// Verify there is no second request.
-		if _, err = nextBodyFn(); err == nil {
+		if _, err = streamIO.NextRequest(); err == nil {
 			return fmt.Errorf("Request data contains more than 1 message for server-streaming RPC")
 		} else if err != io.EOF {
 			return err
@@ -202,7 +172,7 @@ func makeServerStream(ctx context.Context, stream *yarpctransport.ClientStream,
 			break
 		}
 
-		err = responseHandlerFn(resBody)
+		err = streamIO.HandleResponse(resBody)
 	}
 
 	if err == io.EOF {
@@ -212,13 +182,12 @@ func makeServerStream(ctx context.Context, stream *yarpctransport.ClientStream,
 }
 
 // makeClientStream starts client-side streaming rpc
-func makeClientStream(ctx context.Context, stream *yarpctransport.ClientStream,
-	nextBodyFn streamRequestSupplierFn, responseHandlerFn streamResponseHandlerFn) error {
+func makeClientStream(ctx context.Context, stream *yarpctransport.ClientStream, streamIO StreamIO) error {
 
 	var err error
 	for err == nil {
 		var reqBody []byte
-		reqBody, err = nextBodyFn()
+		reqBody, err = streamIO.NextRequest()
 		if err != nil {
 			break
 		}
@@ -236,12 +205,11 @@ func makeClientStream(ctx context.Context, stream *yarpctransport.ClientStream,
 	if err != nil {
 		return err
 	}
-	return responseHandlerFn(res)
+	return streamIO.HandleResponse(res)
 }
 
 // makeBidiStream starts bi-directional streaming rpc
-func makeBidiStream(ctx context.Context, stream *yarpctransport.ClientStream,
-	nextBodyFn streamRequestSupplierFn, responseHandlerFn streamResponseHandlerFn) error {
+func makeBidiStream(ctx context.Context, stream *yarpctransport.ClientStream, streamIO StreamIO) error {
 
 	var wg sync.WaitGroup
 	var sendErr error
@@ -257,7 +225,7 @@ func makeBidiStream(ctx context.Context, stream *yarpctransport.ClientStream,
 		var err error
 		for err == nil {
 			var reqBody []byte
-			reqBody, err = nextBodyFn()
+			reqBody, err = streamIO.NextRequest()
 			if err == io.EOF {
 				err = closeSendStream(ctx, stream)
 				break
@@ -285,7 +253,7 @@ func makeBidiStream(ctx context.Context, stream *yarpctransport.ClientStream,
 			break
 		}
 
-		receiveErr = responseHandlerFn(resBody)
+		receiveErr = streamIO.HandleResponse(resBody)
 	}
 
 	cancel()
@@ -342,3 +310,59 @@ func closeSendStream(ctx context.Context, stream *yarpctransport.ClientStream) e
 	}
 	return nil
 }
+
+// streamIOInitializer uses provided stream message reader to provide stream
+// IO methods to get next requests and handle responses.
+type streamIOInitializer struct {
+	eofReached bool // flag indicates if streamMsgReader returned EOF
+
+	out             output
+	serializer      encoding.Serializer
+	streamMsgReader encoding.StreamRequestReader
+}
+
+// NextRequest returns the next stream request body from the given stream
+// message reader. It returns EOF when there are no more messages to be read
+// from the provided stream message reader.
+func (s *streamIOInitializer) NextRequest() ([]byte, error) {
+	if s.eofReached {
+		return nil, io.EOF
+	}
+
+	msg, err := s.streamMsgReader.NextBody()
+	if err == io.EOF {
+		s.eofReached = true
+
+		return nil, err
+	}
+	if err != nil {
+		return nil, fmt.Errorf("Failed while reading stream input: %v", err)
+	}
+
+	return msg, nil
+}
+
+// HandleResponse deserializes the given response bytes and prints indented JSON body.
+func (s *streamIOInitializer) HandleResponse(body []byte) error {
+	res, err := s.serializer.Response(&transport.Response{Body: body})
+	if err != nil {
+		return fmt.Errorf("Failed while serializing stream response: %v", err)
+	}
+
+	bs, err := json.MarshalIndent(res, "", "  ")
+	if err != nil {
+		return fmt.Errorf("Failed to convert map to JSON: %v\nMap: %+v", err, res)
+	}
+
+	s.out.Printf("%s\n\n", bs)
+	return nil
+}
+
+// newStreamIOInitializer returns streamIO which also records requests.
+func newStreamIOInitializer(out output, serializer encoding.Serializer, streamMsgReader encoding.StreamRequestReader) *streamIOInitializer {
+	return &streamIOInitializer{
+		out:             out,
+		serializer:      serializer,
+		streamMsgReader: streamMsgReader,
+	}
+}


Test patch:
diff --git a/handler_test.go b/handler_test.go
new file mode 100644
index 0000000..aeab700
--- /dev/null
+++ b/handler_test.go
@@ -0,0 +1,110 @@
+// Copyright (c) 2021 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package main
+
+import (
+	"bytes"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/yarpc/yab/encoding"
+)
+
+type mockStreamReader struct {
+	idx      int
+	requests [][]byte
+
+	returnErr error
+}
+
+func (m *mockStreamReader) NextBody() ([]byte, error) {
+	if m.returnErr != nil {
+		return nil, m.returnErr
+	}
+	if m.idx == len(m.requests) {
+		return nil, io.EOF
+	}
+
+	req := m.requests[m.idx]
+	m.idx++
+
+	return req, nil
+}
+
+func TestStreamRequestRecorder(t *testing.T) {
+	t.Run("next request body with eof", func(t *testing.T) {
+		requests := [][]byte{
+			[]byte("request-a"),
+			[]byte("request-b"),
+		}
+
+		streamIO := streamIOInitializer{streamMsgReader: &mockStreamReader{requests: requests}}
+
+		for i, expectedReq := range requests {
+			req, err := streamIO.NextRequest()
+			require.NoError(t, err)
+			assert.Equal(t, expectedReq, req, "requests[%d] is not equal", i)
+		}
+
+		_, err := streamIO.NextRequest()
+		assert.EqualError(t, err, io.EOF.Error())
+		assert.True(t, streamIO.eofReached)
+
+		// NextRequest must return EOF for every call once EOF was returned earlier
+		_, err = streamIO.NextRequest()
+		assert.EqualError(t, err, io.EOF.Error())
+	})
+
+	t.Run("eof reached true", func(t *testing.T) {
+		streamIO := streamIOInitializer{eofReached: true}
+
+		_, err := streamIO.NextRequest()
+		assert.EqualError(t, err, io.EOF.Error())
+	})
+
+	t.Run("handle response success", func(t *testing.T) {
+		var outBuf bytes.Buffer
+		out := testOutput{
+			Buffer: &outBuf,
+		}
+
+		body := `{
+  "test": "1"
+}
+
+`
+		streamIO := streamIOInitializer{out: out, serializer: encoding.NewJSON("test")}
+		require.NoError(t, streamIO.HandleResponse([]byte(body)))
+		assert.Equal(t, body, out.String())
+	})
+
+	t.Run("handle response failure", func(t *testing.T) {
+		var outBuf bytes.Buffer
+		out := testOutput{
+			Buffer: &outBuf,
+		}
+
+		streamIO := streamIOInitializer{out: out, serializer: encoding.NewJSON("test")}
+		require.Error(t, streamIO.HandleResponse([]byte("a:b")))
+	})
+}