diff --git a/benchmark.go b/benchmark.go index a3e016b8..8e4303fe 100644 --- a/benchmark.go +++ b/benchmark.go @@ -22,6 +22,7 @@ package main import ( "errors" + "io/ioutil" "os" "os/signal" "runtime" @@ -31,9 +32,14 @@ import ( "github.com/yarpc/yab/limiter" "github.com/yarpc/yab/statsd" "github.com/yarpc/yab/transport" + + pb "gopkg.in/cheggaaa/pb.v1" ) var ( + progressMarker int + progressUnit pb.Units + errNegativeDuration = errors.New("duration cannot be negative") errNegativeMaxReqs = errors.New("max requests cannot be negative") ) @@ -75,8 +81,10 @@ func (o BenchmarkOptions) enabled() bool { return o.MaxDuration != 0 || o.MaxRequests != 0 } -func runWorker(t transport.Transport, m benchmarkMethod, s *benchmarkState, run *limiter.Run) { +func runWorker(t transport.Transport, m benchmarkMethod, s *benchmarkState, run *limiter.Run, pb *pb.ProgressBar) { for cur := run; cur.More(); { + pb.Increment() + latency, err := m.call(t) if err != nil { s.recordError(err) @@ -99,8 +107,10 @@ func runBenchmark(out output, allOpts Options, m benchmarkMethod) { if opts.RPS > 0 && opts.MaxDuration > 0 { // The RPS * duration in seconds may cap opts.MaxRequests. + // This int cast rounds down 1.0 * 999ms (0.999) = 0.999 = int(0.999) = 0 rpsMax := int(float64(opts.RPS) * opts.MaxDuration.Seconds()) - if rpsMax < opts.MaxRequests || opts.MaxRequests == 0 { + + if (rpsMax > 0 && rpsMax < opts.MaxRequests) || opts.MaxRequests == 0 { opts.MaxRequests = rpsMax } } @@ -132,6 +142,15 @@ func runBenchmark(out output, allOpts Options, m benchmarkMethod) { states[i] = newBenchmarkState(statter) } + progressMarker, progressUnit = progressBarSetup(&opts) + progressBar := pb.New(progressMarker) + progressBar.SetUnits(progressUnit) + progressBar.Output = ioutil.Discard + if opts.ProgressBar { + progressBar.Output = out + } + progressBar.Start() + run := limiter.New(opts.MaxRequests, opts.RPS, opts.MaxDuration) stopOnInterrupt(out, run) @@ -143,7 +162,7 @@ func runBenchmark(out output, allOpts Options, m benchmarkMethod) { wg.Add(1) go func(c transport.Transport) { defer wg.Done() - runWorker(c, m, state, run) + runWorker(c, m, state, run, progressBar) }(c) } } @@ -153,6 +172,9 @@ func runBenchmark(out output, allOpts Options, m benchmarkMethod) { // Wait for all the worker goroutines to end. wg.Wait() total := time.Since(start) + progressBar.Finish() + + // progressBar.FinishPrint("Benchmark finished") // Merge all the states into 0 overall := states[0] @@ -180,3 +202,10 @@ func stopOnInterrupt(out output, r *limiter.Run) { r.Stop() }() } + +func progressBarSetup(opts *BenchmarkOptions) (int, pb.Units) { + if opts.MaxRequests > 0 { + return opts.MaxRequests, pb.U_NO + } + return int(opts.MaxDuration), pb.U_DURATION +} diff --git a/benchmark_test.go b/benchmark_test.go index 698a86f2..857abd5e 100644 --- a/benchmark_test.go +++ b/benchmark_test.go @@ -31,6 +31,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/uber/tchannel-go/testutils" "go.uber.org/atomic" + "gopkg.in/cheggaaa/pb.v1" ) func TestBenchmark(t *testing.T) { @@ -41,12 +42,17 @@ func TestBenchmark(t *testing.T) { rps int want int wantDuration time.Duration + + wantMarker int + wantUnit pb.Units }{ { - msg: "Capped by max requests", - n: 100, - d: 100 * time.Second, - want: 100, + msg: "Capped by max requests", + n: 100, + d: 100 * time.Second, + want: 100, + wantMarker: 100, + wantUnit: pb.U_NO, }, { msg: "Capped by RPS * duration", @@ -54,11 +60,15 @@ func TestBenchmark(t *testing.T) { rps: 120, want: 60, wantDuration: 500 * time.Millisecond, + wantMarker: 60, + wantUnit: pb.U_NO, }, { msg: "Capped by duration", d: 500 * time.Millisecond, wantDuration: 500 * time.Millisecond, + wantMarker: int(500 * time.Millisecond), + wantUnit: pb.U_DURATION, }, } @@ -71,18 +81,20 @@ func TestBenchmark(t *testing.T) { })) m := benchmarkMethodForTest(t, fooMethod, transport.TChannel) + buf, _, out := getOutput(t) for _, tt := range tests { + buf.Reset() requests.Store(0) start := time.Now() - buf, _, out := getOutput(t) + runBenchmark(out, Options{ BOpts: BenchmarkOptions{ MaxRequests: tt.n, MaxDuration: tt.d, RPS: tt.rps, - Connections: 50, + Connections: 25, Concurrency: 2, }, TOpts: s.transportOpts(), @@ -96,6 +108,8 @@ func TestBenchmark(t *testing.T) { assert.EqualValues(t, tt.want, requests.Load(), "%v: Invalid number of requests", tt.msg) } + assert.Equal(t, tt.wantMarker, progressMarker, "progress bar total should be: %v", tt.wantMarker) + assert.Equal(t, tt.wantUnit, progressUnit, "progress bar unit should be %v", tt.wantUnit) if tt.wantDuration != 0 { // Make sure the total duration is within a delta. @@ -149,3 +163,77 @@ func TestRunBenchmarkErrors(t *testing.T) { assert.Contains(t, fatalMessage, tt.wantErr, "Missing error for %+v", tt.opts) } } + +func TestBenchmarkProgressBar(t *testing.T) { + tests := []struct { + msg string + maxRequests int + d time.Duration + rps int + + wantOut string + wantMarker int + wantUnit pb.Units + }{ + { + msg: "RPS simple progrss bar", + maxRequests: 100, + wantMarker: 100, + wantUnit: pb.U_NO, + wantOut: "100 / 100 100.00%", + }, + { + msg: "Duration progress bar", + d: 1 * time.Second, + wantMarker: 1, + wantUnit: pb.U_DURATION, + wantOut: "1s", + }, + { + msg: "RPS times duration in seconds 1 second", + d: 1 * time.Second, + rps: 177, + wantMarker: 177, + wantUnit: pb.U_NO, + wantOut: "177 / 177 100.00%", + }, + { + msg: "RPS times duration in seconds", + d: 500 * time.Millisecond, + rps: 100, + wantMarker: 50, + wantUnit: pb.U_NO, + wantOut: "50 / 50 100.00%", + }, + } + var requests atomic.Int32 + s := newServer(t) + defer s.shutdown() + s.register(fooMethod, methods.errorIf(func() bool { + requests.Inc() + return false + })) + + m := benchmarkMethodForTest(t, fooMethod, transport.TChannel) + + buf, _, out := getOutput(t) + for _, tt := range tests { + buf.Reset() + requests.Store(0) + + t.Run(fmt.Sprintf(tt.msg), func(t *testing.T) { + runBenchmark(out, Options{ + BOpts: BenchmarkOptions{ + MaxRequests: tt.maxRequests, + MaxDuration: tt.d, + RPS: tt.rps, + Connections: 50, + Concurrency: 2, + ProgressBar: true, + }, + TOpts: s.transportOpts(), + }, m) + assert.Contains(t, buf.String(), tt.wantOut) + }) + } +} diff --git a/glide.lock b/glide.lock index 3bcb2c01..1d1c7d90 100644 --- a/glide.lock +++ b/glide.lock @@ -1,8 +1,8 @@ -hash: da28cdd6cc388e155bde58369bf3f1d979954e11446888c23f3bb51187d1806a -updated: 2017-04-25T15:16:31.435502477-07:00 +hash: 449b3c6f4996b7e8fc006fe20f57caac03d3a38d85ed4a4707541bdbe3fd9a5f +updated: 2017-05-31T11:04:29.93385975-07:00 imports: - name: github.com/apache/thrift - version: 3311a9b2375276441234218f4351c6a8f66a6bc2 + version: e41e47c2b4b2407bac525d203b281c63fb253978 subpackages: - lib/go/thrift - name: github.com/cactus/go-statsd-client @@ -17,8 +17,10 @@ imports: - spew - name: github.com/jessevdk/go-flags version: 460c7bb0abd6e927f2767cadc91aa6ef776a98b4 +- name: github.com/mattn/go-runewidth + version: 14207d285c6c197daabb5c9793d63e7af9ab2d50 - name: github.com/opentracing/opentracing-go - version: 6edb48674bd9467b8e91fda004f2bd7202d60ce4 + version: 1949ddbfd147afd4d964a9f00b24eb291e0e7c38 subpackages: - ext - log @@ -32,7 +34,7 @@ imports: - assert - require - name: github.com/uber-go/atomic - version: 4e336646b2ef9fc6e47be8e21594178f98e5ebcf + version: 0c9e689d64f004564b79d9a663634756df322902 - name: github.com/uber/jaeger-client-go version: e39d0f1b622558cae3d9db0062a739cc6ffa700f subpackages: @@ -75,14 +77,53 @@ imports: - version - wire - name: golang.org/x/net - version: da118f7b8e5954f39d0d2130ab35d4bf0e3cb344 + version: 513929065c19401a1c7b76ecd942f9f86a0c061b subpackages: - context +- name: gopkg.in/cheggaaa/pb.v1 + version: f6ccf2184de4dd34495277e38dc19b6e7fbe0ea2 - name: gopkg.in/yaml.v2 version: a83829b6f1293c91addabc89d0571c246397bbf4 testImports: +- name: github.com/beorn7/perks + version: 4c0e84591b9aa9e6dcfdf3e020114cd81f89d5f9 + subpackages: + - quantile +- name: github.com/facebookgo/clock + version: 600d898af40aa09a7a93ecb9265d87b0504b6f03 +- name: github.com/golang/protobuf + version: 7cc19b78d562895b13596ddce7aafb59dd789318 + subpackages: + - proto +- name: github.com/matttproud/golang_protobuf_extensions + version: c12348ce28de40eed0136aa2b644d0ee0650e56c + subpackages: + - pbutil +- name: github.com/prometheus/client_golang + version: c5b7fccd204277076155f10851dad72b76a49317 + subpackages: + - prometheus + - prometheus/promhttp +- name: github.com/prometheus/client_model + version: 6f3806018612930941127f2a7c6c453ba2c527d2 + subpackages: + - go +- name: github.com/prometheus/common + version: dd2f054febf4a6c00f2343686efb775948a8bff4 + subpackages: + - expfmt + - internal/bitbucket.org/ww/goautoneg + - model +- name: github.com/prometheus/procfs + version: a1dba9ce8baed984a2495b658c82687f8157b98f + subpackages: + - xfs +- name: github.com/uber-go/tally + version: e9b601813b0be4771b1b3995390567b67ab2b3fc +- name: go.uber.org/multierr + version: a3d1fc1f1316d4132fc61f4ea1159ae0613fb474 - name: go.uber.org/yarpc - version: 6ae533f0810337028ef055b690360e050aa13219 + version: 6b9245d126758870690c44f27f695cc7e811f913 subpackages: - api/encoding - api/middleware @@ -98,7 +139,10 @@ testImports: - internal/introspection - internal/iopool - internal/net + - internal/observability - internal/outboundmiddleware + - internal/pally + - internal/procedure - internal/request - internal/sync - peer @@ -106,3 +150,12 @@ testImports: - transport/http - transport/tchannel - transport/tchannel/internal +- name: go.uber.org/zap + version: fab453050a7a08c35f31fc5fff6f2dbd962285ab + subpackages: + - buffer + - internal/bufferpool + - internal/color + - internal/exit + - internal/multierror + - zapcore diff --git a/glide.yaml b/glide.yaml index cb42a04d..91d9efb4 100644 --- a/glide.yaml +++ b/glide.yaml @@ -16,6 +16,8 @@ import: version: ^1 - package: github.com/uber/jaeger-client-go version: ^1 +- package: gopkg.in/cheggaaa/pb.v1 + version: ^1 testImport: - package: github.com/apache/thrift version: master diff --git a/options.go b/options.go index bac3986e..44335943 100644 --- a/options.go +++ b/options.go @@ -105,6 +105,9 @@ type BenchmarkOptions struct { // Benchmark metrics can optionally be reported via statsd. StatsdHostPort string `long:"statsd" description:"Optional host:port of a StatsD server to report metrics"` + + // Output options + ProgressBar bool `long:"progress-bar" description:"show a progress bar for the benchmark."` } func newOptions() *Options {