Repository: yarpc/yab
Pull request #322 (2021-02-26 05:07:21)

Refactor benchmark to accept benchmarker interface 
Add `benchmarker` interface in benchmark and warmup methods. This change sets up clean interface to add streaming benchmarker from #319. Decouples transport warmup methods from unary benchmark which enables them to be used by streaming benchmarks as well. 


Patch:
diff --git a/bench_method.go b/bench_method.go
index b94a568..afda89e 100644
--- a/bench_method.go
+++ b/bench_method.go
@@ -25,7 +25,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/yarpc/yab/encoding"
 	"github.com/yarpc/yab/transport"
 
 	"github.com/opentracing/opentracing-go"
@@ -36,21 +35,22 @@ type peerTransport struct {
 	peerID int
 }
 
-type benchmarkMethod struct {
-	serializer encoding.Serializer
-	req        *transport.Request
+// 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)
 }
 
-// WarmTransport warms up a transport and returns it. The transport is warmed
+// warmTransport warms up a transport and returns it. The transport is warmed
 // up by making some number of requests through it.
-func (m benchmarkMethod) WarmTransport(opts TransportOptions, resolved resolvedProtocolEncoding, warmupRequests int) (transport.Transport, error) {
+func warmTransport(b benchmarkCaller, opts TransportOptions, resolved resolvedProtocolEncoding, warmupRequests int) (transport.Transport, error) {
 	transport, err := getTransport(opts, resolved, opentracing.NoopTracer{})
 	if err != nil {
 		return nil, err
 	}
 
 	for i := 0; i < warmupRequests; i++ {
-		_, err := makeRequest(transport, m.req)
+		_, err := b.Call(transport)
 		if err != nil {
 			return nil, err
 		}
@@ -59,21 +59,6 @@ func (m benchmarkMethod) WarmTransport(opts TransportOptions, resolved resolvedP
 	return transport, nil
 }
 
-func (m benchmarkMethod) call(t transport.Transport) (time.Duration, error) {
-	start := time.Now()
-	res, err := makeRequest(t, m.req)
-	duration := time.Since(start)
-
-	if err == nil {
-		err = m.serializer.CheckSuccess(res)
-	}
-	return duration, err
-}
-
-func (m benchmarkMethod) Method() string {
-	return m.req.Method
-}
-
 func peerBalancer(peers []string) func(i int) (string, int) {
 	numPeers := len(peers)
 	startOffset := rand.Intn(numPeers)
@@ -83,9 +68,9 @@ func peerBalancer(peers []string) func(i int) (string, int) {
 	}
 }
 
-// WarmTransports returns n transports that have been warmed up.
+// warmTransports returns n transports that have been warmed up.
 // No requests may fail during the warmup period.
-func (m benchmarkMethod) WarmTransports(n int, tOpts TransportOptions, resolved resolvedProtocolEncoding, warmupRequests int) ([]peerTransport, error) {
+func warmTransports(b benchmarkCaller, n int, tOpts TransportOptions, resolved resolvedProtocolEncoding, warmupRequests int) ([]peerTransport, error) {
 	peerFor := peerBalancer(tOpts.Peers)
 	transports := make([]peerTransport, n)
 	errs := make([]error, n)
@@ -99,7 +84,7 @@ func (m benchmarkMethod) WarmTransports(n int, tOpts TransportOptions, resolved
 			peerHostPort, peerIndex := peerFor(i)
 			tOpts.Peers = []string{peerHostPort}
 
-			tp, err := m.WarmTransport(tOpts, resolved, warmupRequests)
+			tp, err := warmTransport(b, tOpts, resolved, warmupRequests)
 			transports[i] = peerTransport{tp, peerIndex}
 			errs[i] = err
 		}(i, tOpts)
diff --git a/benchmark.go b/benchmark.go
index c49481b..ee4b638 100644
--- a/benchmark.go
+++ b/benchmark.go
@@ -109,9 +109,9 @@ func (o BenchmarkOptions) enabled() bool {
 	return o.MaxDuration != 0 || o.MaxRequests != 0
 }
 
-func runWorker(t transport.Transport, m benchmarkMethod, s *benchmarkState, run *limiter.Run, logger *zap.Logger) {
+func runWorker(t transport.Transport, b benchmarkCaller, s *benchmarkState, run *limiter.Run, logger *zap.Logger) {
 	for cur := run; cur.More(); {
-		latency, err := m.call(t)
+		latency, err := b.Call(t)
 		if err != nil {
 			s.recordError(err)
 			// TODO: Add information about which peer specifically failed.
@@ -123,7 +123,7 @@ func runWorker(t transport.Transport, m benchmarkMethod, s *benchmarkState, run
 	}
 }
 
-func runBenchmark(out output, logger *zap.Logger, allOpts Options, resolved resolvedProtocolEncoding, m benchmarkMethod) {
+func runBenchmark(out output, logger *zap.Logger, allOpts Options, resolved resolvedProtocolEncoding, methodName string, b benchmarkCaller) {
 	opts := allOpts.BOpts
 
 	if err := opts.validate(); err != nil {
@@ -167,12 +167,12 @@ func runBenchmark(out output, logger *zap.Logger, allOpts Options, resolved reso
 
 	// Warm up number of connections.
 	logger.Debug("Warming up connections.", zap.Int("numConns", numConns))
-	connections, err := m.WarmTransports(numConns, allOpts.TOpts, resolved, opts.WarmupRequests)
+	connections, err := warmTransports(b, numConns, allOpts.TOpts, resolved, opts.WarmupRequests)
 	if err != nil {
 		out.Fatalf("Failed to warmup connections for benchmark: %v", err)
 	}
 
-	globalStatter, err := statsd.NewClient(logger, opts.StatsdHostPort, allOpts.TOpts.ServiceName, m.Method())
+	globalStatter, err := statsd.NewClient(logger, opts.StatsdHostPort, allOpts.TOpts.ServiceName, methodName)
 	if err != nil {
 		out.Fatalf("Failed to create statsd client for benchmark: %v", err)
 	}
@@ -210,7 +210,7 @@ func runBenchmark(out output, logger *zap.Logger, allOpts Options, resolved reso
 			wg.Add(1)
 			go func(c transport.Transport) {
 				defer wg.Done()
-				runWorker(c, m, state, run, logger)
+				runWorker(c, b, state, run, logger)
 			}(c)
 		}
 	}
diff --git a/benchmark_unary.go b/benchmark_unary.go
new file mode 100644
index 0000000..d952b7c
--- /dev/null
+++ b/benchmark_unary.go
@@ -0,0 +1,46 @@
+// 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 (
+	"time"
+
+	"github.com/yarpc/yab/encoding"
+	"github.com/yarpc/yab/transport"
+)
+
+// benchmarkUnaryMethod benchmarks unary requests.
+type benchmarkUnaryMethod struct {
+	serializer encoding.Serializer
+	req        *transport.Request
+}
+
+// Call dispatches unary request on the provided transport.
+func (m benchmarkUnaryMethod) Call(t transport.Transport) (time.Duration, error) {
+	start := time.Now()
+	res, err := makeRequest(t, m.req)
+	duration := time.Since(start)
+
+	if err == nil {
+		err = m.serializer.CheckSuccess(res)
+	}
+	return duration, err
+}
diff --git a/handler.go b/handler.go
index 1624e70..2c83a79 100644
--- a/handler.go
+++ b/handler.go
@@ -69,7 +69,7 @@ func (r requestHandler) handleUnaryRequest() {
 		makeInitialRequest(r.out, r.transport, r.serializer, req)
 	}
 
-	runBenchmark(r.out, r.logger, r.opts, r.resolved, benchmarkMethod{
+	runBenchmark(r.out, r.logger, r.opts, r.resolved, req.Method, benchmarkUnaryMethod{
 		serializer: r.serializer,
 		req:        req,
 	})


Test patch:
diff --git a/benchmark_test.go b/benchmark_test.go
index 669dc3b..798b570 100644
--- a/benchmark_test.go
+++ b/benchmark_test.go
@@ -90,7 +90,7 @@ func TestBenchmark(t *testing.T) {
 				Concurrency: 2,
 			},
 			TOpts: s.transportOpts(),
-		}, _resolvedTChannelThrift, m)
+		}, _resolvedTChannelThrift, fooMethod, m)
 
 		bufStr := buf.String()
 		assert.Contains(t, bufStr, "Max RPS")
@@ -146,7 +146,7 @@ func TestRunBenchmarkErrors(t *testing.T) {
 		// need to run the benchmark in a separate goroutine.
 		go func() {
 			defer wg.Done()
-			runBenchmark(out, _testLogger, opts, _resolvedTChannelThrift, m)
+			runBenchmark(out, _testLogger, opts, _resolvedTChannelThrift, fooMethod, m)
 		}()
 
 		wg.Wait()
@@ -198,7 +198,7 @@ func TestBenchmarkStatsPerPeer(t *testing.T) {
 		ROpts: RequestOptions{
 			Procedure: fooMethod,
 		},
-	}, _resolvedTChannelThrift, m)
+	}, _resolvedTChannelThrift, fooMethod, m)
 
 	// Ensure one backend gets more calls than the other
 	want := map[string]int{
@@ -284,7 +284,7 @@ func TestBenchmarkOutput(t *testing.T) {
 					TOpts: s.transportOpts(),
 				}
 
-				runBenchmark(out, _testLogger, opts, _resolvedTChannelThrift, m)
+				runBenchmark(out, _testLogger, opts, _resolvedTChannelThrift, fooMethod, m)
 				bufStr := buf.String()
 				bufWarnStr := bufWarn.String()
 
diff --git a/bench_method_test.go b/benchmark_unary_test.go
similarity index 92%
rename from bench_method_test.go
rename to benchmark_unary_test.go
index 0e3b79f..672ab95 100644
--- a/bench_method_test.go
+++ b/benchmark_unary_test.go
@@ -36,7 +36,7 @@ import (
 	"go.uber.org/atomic"
 )
 
-func benchmarkMethodForTest(t *testing.T, procedure string, p transport.Protocol) benchmarkMethod {
+func benchmarkMethodForTest(t *testing.T, procedure string, p transport.Protocol) benchmarkUnaryMethod {
 	return benchmarkMethodForROpts(t, RequestOptions{
 		Encoding:   encoding.Thrift,
 		ThriftFile: validThrift,
@@ -44,7 +44,7 @@ func benchmarkMethodForTest(t *testing.T, procedure string, p transport.Protocol
 	}, p)
 }
 
-func benchmarkMethodForROpts(t *testing.T, rOpts RequestOptions, p transport.Protocol) benchmarkMethod {
+func benchmarkMethodForROpts(t *testing.T, rOpts RequestOptions, p transport.Protocol) benchmarkUnaryMethod {
 	serializer, err := NewSerializer(Options{ROpts: rOpts}, resolvedProtocolEncoding{
 		protocol: p,
 		enc:      encoding.Thrift,
@@ -55,7 +55,7 @@ func benchmarkMethodForROpts(t *testing.T, rOpts RequestOptions, p transport.Pro
 	require.NoError(t, err, "Failed to serialize Thrift body")
 
 	req.Timeout = time.Second
-	return benchmarkMethod{serializer, req}
+	return benchmarkUnaryMethod{serializer, req}
 }
 
 func TestBenchmarkMethodWarmTransport(t *testing.T) {
@@ -96,7 +96,7 @@ func TestBenchmarkMethodWarmTransport(t *testing.T) {
 			Peers:       []string{tt.peer},
 		}
 
-		transport, err := m.WarmTransport(tOpts, resolvedProtocolEncoding{
+		transport, err := warmTransport(m, tOpts, resolvedProtocolEncoding{
 			protocol: transport.TChannel,
 			enc:      encoding.JSON,
 		}, 1 /* warmupRequests */)
@@ -159,7 +159,7 @@ func TestBenchmarkMethodCall(t *testing.T) {
 			m.req.Method = tt.reqMethod
 		}
 
-		d, err := m.call(tp)
+		d, 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")
@@ -233,7 +233,7 @@ func TestBenchmarkMethodWarmTransportsSuccess(t *testing.T) {
 		ServiceName: "foo",
 		Peers:       serverHPs,
 	}
-	transports, err := m.WarmTransports(numServers, tOpts, _resolvedTChannelThrift, 1 /* warmupRequests */)
+	transports, err := warmTransports(m, numServers, tOpts, _resolvedTChannelThrift, 1 /* warmupRequests */)
 	assert.NoError(t, err, "WarmTransports should not fail")
 	assert.Equal(t, numServers, len(transports), "Got unexpected number of transports")
 	for i, transport := range transports {
@@ -292,7 +292,7 @@ func TestBenchmarkMethodWarmTransportsError(t *testing.T) {
 			ServiceName: "foo",
 			Peers:       []string{s.hostPort()},
 		}
-		_, err := m.WarmTransports(10, tOpts, _resolvedTChannelThrift, tt.warmup)
+		_, err := warmTransports(m, 10, tOpts, _resolvedTChannelThrift, tt.warmup)
 		if tt.wantErr {
 			assert.Error(t, err, "%v: WarmTransports should fail", msg)
 		} else {
@@ -300,11 +300,3 @@ func TestBenchmarkMethodWarmTransportsError(t *testing.T) {
 		}
 	}
 }
-
-func TestBenchmarkMethodHealth(t *testing.T) {
-	m := benchmarkMethodForROpts(t, RequestOptions{
-		Encoding: encoding.Thrift,
-		Health:   true,
-	}, transport.TChannel)
-	assert.Equal(t, "Meta::health", m.Method())
-}