Repository: yarpc/yab
Pull request #324 (2021-03-01 17:06:14)

Add benchmark caller result interface
This change intends to enable `Call` to return richer call result details apart from Latency for streaming benchmark calls such as stream messages sent and received during the benchmark call.

Follow-up PR: #325 


Patch:
diff --git a/bench_method.go b/bench_method.go
index afda89e..07df6c6 100644
--- a/bench_method.go
+++ b/bench_method.go
@@ -38,7 +38,19 @@ type peerTransport struct {
 // benchmarkCaller exposes method to dispatch requests for benchmark.
 type benchmarkCaller interface {
 	// Call dispatches a request using the provided transport.
-	Call(transport.Transport) (time.Duration, error)
+	Call(transport.Transport) (benchmarkCallLatencyResult, error)
+}
+
+type benchmarkCallLatencyResult struct {
+	latency time.Duration
+}
+
+func newBenchmarkCallLatencyResult(latency time.Duration) benchmarkCallLatencyResult {
+	return benchmarkCallLatencyResult{latency}
+}
+
+func (r benchmarkCallLatencyResult) Latency() time.Duration {
+	return r.latency
 }
 
 // warmTransport warms up a transport and returns it. The transport is warmed
diff --git a/benchmark.go b/benchmark.go
index fe746c7..2e51d0b 100644
--- a/benchmark.go
+++ b/benchmark.go
@@ -111,7 +111,7 @@ func (o BenchmarkOptions) enabled() bool {
 
 func runWorker(t transport.Transport, b benchmarkCaller, s *benchmarkState, run *limiter.Run, logger *zap.Logger) {
 	for cur := run; cur.More(); {
-		latency, err := b.Call(t)
+		callResult, err := b.Call(t)
 		if err != nil {
 			s.recordError(err)
 			// TODO: Add information about which peer specifically failed.
@@ -119,7 +119,7 @@ func runWorker(t transport.Transport, b benchmarkCaller, s *benchmarkState, run
 			continue
 		}
 
-		s.recordLatency(latency)
+		s.recordLatency(callResult.Latency())
 	}
 }
 
diff --git a/benchmark_stream.go b/benchmark_stream.go
index ea63ee5..2c319b4 100644
--- a/benchmark_stream.go
+++ b/benchmark_stream.go
@@ -36,26 +36,26 @@ type benchmarkStreamMethod struct {
 }
 
 // Call dispatches stream request on the provided transport.
-func (m benchmarkStreamMethod) Call(t transport.Transport) (time.Duration, error) {
+func (m benchmarkStreamMethod) Call(t transport.Transport) (benchmarkCallLatencyResult, error) {
 	streamIO := newStreamIOBenchmark(m.streamRequestMessages)
 
 	start := time.Now()
 	err := makeStreamRequest(t, m.streamRequest, m.serializer, streamIO)
-	duration := time.Since(start)
+	callResult := newBenchmarkCallLatencyResult(time.Since(start))
 
 	if err != nil {
-		return duration, err
+		return callResult, err
 	}
 
 	// response validation is delayed to avoid consuming benchmark time for
 	// deserializing the responses.
 	for _, res := range streamIO.streamResponses {
 		if err = m.serializer.CheckSuccess(&transport.Response{Body: res}); err != nil {
-			return duration, err
+			return callResult, err
 		}
 	}
 
-	return duration, err
+	return callResult, err
 }
 
 // streamIOBenchmark provides stream IO methods using the provided stream requests
diff --git a/benchmark_unary.go b/benchmark_unary.go
index d952b7c..f2bf835 100644
--- a/benchmark_unary.go
+++ b/benchmark_unary.go
@@ -34,13 +34,13 @@ type benchmarkUnaryMethod struct {
 }
 
 // Call dispatches unary request on the provided transport.
-func (m benchmarkUnaryMethod) Call(t transport.Transport) (time.Duration, error) {
+func (m benchmarkUnaryMethod) Call(t transport.Transport) (benchmarkCallLatencyResult, error) {
 	start := time.Now()
 	res, err := makeRequest(t, m.req)
-	duration := time.Since(start)
+	latency := time.Since(start)
 
 	if err == nil {
 		err = m.serializer.CheckSuccess(res)
 	}
-	return duration, err
+	return newBenchmarkCallLatencyResult(latency), err
 }


Test patch:
diff --git a/benchmark_unary_test.go b/benchmark_unary_test.go
index 5211205..4f1d359 100644
--- a/benchmark_unary_test.go
+++ b/benchmark_unary_test.go
@@ -159,7 +159,7 @@ func TestBenchmarkMethodCall(t *testing.T) {
 			m.req.Method = tt.reqMethod
 		}
 
-		d, err := m.Call(tp)
+		res, err := m.Call(tp)
 		if tt.wantErr != "" {
 			if assert.Error(t, err, "call should fail") {
 				assert.Contains(t, err.Error(), tt.wantErr, "call should return 0 duration")
@@ -168,7 +168,7 @@ func TestBenchmarkMethodCall(t *testing.T) {
 		}
 
 		assert.NoError(t, err, "call should not fail")
-		assert.True(t, d > time.Microsecond, "duration was too short, got %v", d)
+		assert.True(t, res.Latency() > time.Microsecond, "duration was too short, got %v", res.Latency())
 	}
 }