Repository: yarpc/yab
Pull request #327 (2021-03-09 17:10:21)

Add `stream-interval` stream request option
Stream interval option ensures consecutive stream message sends always have a gap of at-least `stream-interval` specified time duration, basically stream requests are dispatched slowly. This is useful in stream benchmarking/testing where users want to observe stream behaviour under real world conditions where each message send would have some gap between them. This works with ad-hoc and benchmark streaming requests.

```
yab --timeout=1m --service "test-server" --method uber.yarpc.encoding.protobuf.Test/Stream -p grpc://localhost:55555 -r '{"value":"req-1"}{"value":"req-2"}' --stream-interval=500ms

{
  "value": "req-1"
}

// Wait for `stream-interval`=500ms

{
  "value": "req-1req-2"
}
```

codecov[bot]: # [Codecov](https://codecov.io/gh/yarpc/yab/pull/327?src=pr&el=h1) Report
> Merging [#327](https://codecov.io/gh/yarpc/yab/pull/327?src=pr&el=desc) (13ad957) into [dev](https://codecov.io/gh/yarpc/yab/commit/e80962df3b25580799432a2fbba0953de26c1ca8?el=desc) (e80962d) will **increase** coverage by `0.02%`.
> The diff coverage is `95.23%`.

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

```diff
@@            Coverage Diff             @@
##              dev     #327      +/-   ##
==========================================
+ Coverage   90.31%   90.34%   +0.02%     
==========================================
  Files          54       54              
  Lines        2644     2661      +17     
==========================================
+ Hits         2388     2404      +16     
- Misses        180      181       +1     
  Partials       76       76              
```


| [Impacted Files](https://codecov.io/gh/yarpc/yab/pull/327?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [options.go](https://codecov.io/gh/yarpc/yab/pull/327/diff?src=pr&el=tree#diff-b3B0aW9ucy5nbw==) | `93.54% <ø> (ø)` | |
| [handler.go](https://codecov.io/gh/yarpc/yab/pull/327/diff?src=pr&el=tree#diff-aGFuZGxlci5nbw==) | `81.21% <95.00%> (+1.33%)` | :arrow_up: |
| [benchmark\_stream.go](https://codecov.io/gh/yarpc/yab/pull/327/diff?src=pr&el=tree#diff-YmVuY2htYXJrX3N0cmVhbS5nbw==) | `85.71% <100.00%> (ø)` | |

------

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


Patch:
diff --git a/benchmark_stream.go b/benchmark_stream.go
index 0f2df68..62d731f 100644
--- a/benchmark_stream.go
+++ b/benchmark_stream.go
@@ -33,6 +33,7 @@ type benchmarkStreamMethod struct {
 	serializer            encoding.Serializer
 	streamRequest         *transport.StreamRequest
 	streamRequestMessages [][]byte
+	opts                  StreamRequestOptions
 }
 
 // Call dispatches stream request on the provided transport.
@@ -40,7 +41,7 @@ func (m benchmarkStreamMethod) Call(t transport.Transport) (benchmarkCallReporte
 	streamIO := newStreamIOBenchmark(m.streamRequestMessages)
 
 	start := time.Now()
-	err := makeStreamRequest(t, m.streamRequest, m.serializer, streamIO)
+	err := makeStreamRequest(t, m.streamRequest, m.serializer, streamIO, m.opts)
 	callReport := newBenchmarkStreamCallReport(time.Since(start), streamIO.streamMessagesReceived(), streamIO.streamMessagesSent())
 
 	if err != nil {
diff --git a/handler.go b/handler.go
index 494ab85..5b285d4 100644
--- a/handler.go
+++ b/handler.go
@@ -8,6 +8,7 @@ import (
 	"io"
 	"io/ioutil"
 	"sync"
+	"time"
 
 	"github.com/uber/tchannel-go"
 	"github.com/yarpc/yab/encoding"
@@ -95,7 +96,7 @@ func (r requestHandler) handleStreamRequest() {
 	streamIO := newStreamIOInitializer(r.out, r.serializer, streamMsgReader)
 
 	if r.shouldMakeInitialRequest() {
-		if err = makeStreamRequest(r.transport, streamReq, r.serializer, streamIO); err != nil {
+		if err = makeStreamRequest(r.transport, streamReq, r.serializer, streamIO, r.opts.ROpts.StreamRequestOptions); err != nil {
 			r.out.Fatalf("%v\n", err)
 		}
 	}
@@ -113,6 +114,7 @@ func (r requestHandler) handleStreamRequest() {
 		serializer:            r.serializer,
 		streamRequest:         streamReq,
 		streamRequestMessages: streamRequests,
+		opts:                  r.opts.ROpts.StreamRequestOptions,
 	})
 }
 
@@ -130,7 +132,7 @@ func (r requestHandler) isStreamingMethod() bool {
 // 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, streamIO StreamIO) error {
+func makeStreamRequest(t transport.Transport, streamReq *transport.StreamRequest, serializer encoding.Serializer, streamIO StreamIO, opts StreamRequestOptions) error {
 	streamTransport, ok := t.(transport.StreamTransport)
 	if !ok {
 		return fmt.Errorf("Transport does not support stream calls: %q", t.Protocol())
@@ -149,9 +151,9 @@ func makeStreamRequest(t transport.Transport, streamReq *transport.StreamRequest
 
 	switch serializer.MethodType() {
 	case encoding.BidirectionalStream:
-		return makeBidiStream(ctx, stream, streamIO)
+		return makeBidiStream(ctx, stream, streamIO, opts)
 	case encoding.ClientStream:
-		return makeClientStream(ctx, stream, streamIO)
+		return makeClientStream(ctx, stream, streamIO, opts)
 	default:
 		return makeServerStream(ctx, stream, streamIO)
 	}
@@ -194,14 +196,18 @@ func makeServerStream(ctx context.Context, stream *yarpctransport.ClientStream,
 }
 
 // makeClientStream starts client-side streaming rpc
-func makeClientStream(ctx context.Context, stream *yarpctransport.ClientStream, streamIO StreamIO) error {
+func makeClientStream(ctx context.Context, stream *yarpctransport.ClientStream, streamIO StreamIO, opts StreamRequestOptions) error {
 	var err error
+	reqWaiter := newIntervalWaiter(opts.Interval.Duration())
+
 	for err == nil {
 		var reqBody []byte
 		reqBody, err = streamIO.NextRequest()
 		if err != nil {
 			break
 		}
+
+		reqWaiter.wait(ctx)
 		err = sendStreamMessage(ctx, stream, reqBody)
 	}
 
@@ -220,12 +226,13 @@ func makeClientStream(ctx context.Context, stream *yarpctransport.ClientStream,
 }
 
 // makeBidiStream starts bi-directional streaming rpc
-func makeBidiStream(ctx context.Context, stream *yarpctransport.ClientStream, streamIO StreamIO) error {
+func makeBidiStream(ctx context.Context, stream *yarpctransport.ClientStream, streamIO StreamIO, opts StreamRequestOptions) error {
 	var wg sync.WaitGroup
 	var sendErr error
 
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
+	reqWaiter := newIntervalWaiter(opts.Interval.Duration())
 
 	wg.Add(1)
 	// Start go routine to concurrently send stream messages.
@@ -247,6 +254,7 @@ func makeBidiStream(ctx context.Context, stream *yarpctransport.ClientStream, st
 				break
 			}
 
+			reqWaiter.wait(ctx)
 			err = sendStreamMessage(ctx, stream, reqBody)
 		}
 
@@ -395,3 +403,37 @@ func newStreamIOInitializer(out output, serializer encoding.Serializer, streamMs
 		streamMsgReader: streamMsgReader,
 	}
 }
+
+// intervalWaiter provides `wait` method to maintain time gap of `interval` between
+// consecutive `wait` calls.
+type intervalWaiter struct {
+	interval    time.Duration // required time gap between calls.
+	lastAllowed time.Time
+}
+
+func newIntervalWaiter(interval time.Duration) *intervalWaiter {
+	return &intervalWaiter{
+		interval: interval,
+	}
+}
+
+// wait method waits until the time gap between previous call and current call is
+// more than interval provided. It returns early if the provided context is done.
+func (s *intervalWaiter) wait(ctx context.Context) {
+	now := time.Now()
+	diff := now.Sub(s.lastAllowed)
+
+	// gap between this and previous call is more than required interval.
+	if diff >= s.interval {
+		s.lastAllowed = now
+		return
+	}
+
+	// wait for the remaining duration or context completion.
+	select {
+	case <-ctx.Done():
+	case <-time.After(s.interval - diff):
+	}
+
+	s.lastAllowed = time.Now()
+}
diff --git a/options.go b/options.go
index fc37239..8513ad3 100644
--- a/options.go
+++ b/options.go
@@ -70,6 +70,13 @@ type RequestOptions struct {
 		JSON     bool        `long:"json" hidden:"true"`
 		Raw      bool        `long:"raw" hidden:"true"`
 	}
+
+	StreamRequestOptions StreamRequestOptions
+}
+
+// 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."`
 }
 
 // TransportOptions are transport related options.


Test patch:
diff --git a/handler_test.go b/handler_test.go
index 8e63f99..79240ad 100644
--- a/handler_test.go
+++ b/handler_test.go
@@ -22,9 +22,11 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"io"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -127,3 +129,26 @@ func TestStreamRequestRecorder(t *testing.T) {
 		assert.Equal(t, requests, gotRequests)
 	})
 }
+
+func TestIntervalWaiter(t *testing.T) {
+	t.Run("initial call must not wait", func(t *testing.T) {
+		waiter := newIntervalWaiter(time.Millisecond * 100)
+		now := time.Now()
+		waiter.wait(context.Background())
+		timeTaken := time.Now().Sub(now)
+		assert.LessOrEqual(t, int64(timeTaken), int64(time.Millisecond), "Unexpected time taken on wait: %v", timeTaken)
+	})
+
+	t.Run("interval gap must be 100ms between consecutive calls", func(t *testing.T) {
+		waiter := newIntervalWaiter(time.Millisecond * 100)
+		waiter.wait(context.Background())
+
+		for i := 0; i < 2; i++ {
+			now := time.Now()
+			waiter.wait(context.Background())
+			timeTaken := time.Now().Sub(now)
+			assert.GreaterOrEqual(t, int64(timeTaken), int64(time.Millisecond*100), "Unexpected time taken on wait: %v", timeTaken)
+			assert.LessOrEqual(t, int64(timeTaken), int64(time.Millisecond*200), "Unexpected time taken on wait: %v", timeTaken)
+		}
+	})
+}
diff --git a/integration_test.go b/integration_test.go
index 28330c8..f4793bf 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -983,6 +983,74 @@ test: 1
 	}
 }
 
+func TestGRPCStreamWithStreamIntervalOption(t *testing.T) {
+	t.Run("client stream", func(t *testing.T) {
+		addr, server := setupGRPCServer(t, &simpleService{
+			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),
+				},
+			},
+			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),
+				},
+			},
+			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)
+	})
+}
+
 func TestGRPCReflectionSource(t *testing.T) {
 	addr, server := setupGRPCServer(t, &simpleService{})
 	defer server.GracefulStop()