From c63da93d52a32d1f2228b53ce222058849bf0536 Mon Sep 17 00:00:00 2001 From: iamqizhao Date: Thu, 23 Apr 2015 13:10:23 -0700 Subject: [PATCH] stats lib for benchmark --- benchmark/stats/counter.go | 135 ++++++++++++++++++ benchmark/stats/histogram.go | 255 ++++++++++++++++++++++++++++++++++ benchmark/stats/stats.go | 116 ++++++++++++++++ benchmark/stats/timeseries.go | 154 ++++++++++++++++++++ benchmark/stats/tracker.go | 159 +++++++++++++++++++++ benchmark/stats/util.go | 191 +++++++++++++++++++++++++ 6 files changed, 1010 insertions(+) create mode 100644 benchmark/stats/counter.go create mode 100644 benchmark/stats/histogram.go create mode 100644 benchmark/stats/stats.go create mode 100644 benchmark/stats/timeseries.go create mode 100644 benchmark/stats/tracker.go create mode 100644 benchmark/stats/util.go diff --git a/benchmark/stats/counter.go b/benchmark/stats/counter.go new file mode 100644 index 00000000..326b88d9 --- /dev/null +++ b/benchmark/stats/counter.go @@ -0,0 +1,135 @@ +package stats + +import ( + "sync" + "time" +) + +var ( + // Used for testing. + TimeNow func() time.Time = time.Now +) + +const ( + hour = 0 + tenminutes = 1 + minute = 2 +) + +// Counter is a counter that keeps track of its recent values over a given +// period of time, and with a given resolution. Use newCounter() to instantiate. +type Counter struct { + mu sync.RWMutex + ts [3]*timeseries + lastUpdate time.Time +} + +// newCounter returns a new Counter. +func newCounter() *Counter { + now := TimeNow() + c := &Counter{} + c.ts[hour] = newTimeSeries(now, time.Hour, time.Minute) + c.ts[tenminutes] = newTimeSeries(now, 10*time.Minute, 10*time.Second) + c.ts[minute] = newTimeSeries(now, time.Minute, time.Second) + return c +} + +func (c *Counter) advance() time.Time { + now := TimeNow() + for _, ts := range c.ts { + ts.advanceTime(now) + } + return now +} + +// Value returns the current value of the counter. +func (c *Counter) Value() int64 { + c.mu.RLock() + defer c.mu.RUnlock() + return c.ts[minute].headValue() +} + +// LastUpdate returns the last update time of the counter. +func (c *Counter) LastUpdate() time.Time { + c.mu.RLock() + defer c.mu.RUnlock() + return c.lastUpdate +} + +// Set updates the current value of the counter. +func (c *Counter) Set(value int64) { + c.mu.Lock() + defer c.mu.Unlock() + c.lastUpdate = c.advance() + for _, ts := range c.ts { + ts.set(value) + } +} + +// Incr increments the current value of the counter by 'delta'. +func (c *Counter) Incr(delta int64) { + c.mu.Lock() + defer c.mu.Unlock() + c.lastUpdate = c.advance() + for _, ts := range c.ts { + ts.incr(delta) + } +} + +// Delta1h returns the delta for the last hour. +func (c *Counter) Delta1h() int64 { + c.mu.RLock() + defer c.mu.RUnlock() + c.advance() + return c.ts[hour].delta() +} + +// Delta10m returns the delta for the last 10 minutes. +func (c *Counter) Delta10m() int64 { + c.mu.RLock() + defer c.mu.RUnlock() + c.advance() + return c.ts[tenminutes].delta() +} + +// Delta1m returns the delta for the last minute. +func (c *Counter) Delta1m() int64 { + c.mu.RLock() + defer c.mu.RUnlock() + c.advance() + return c.ts[minute].delta() +} + +// Rate1h returns the rate of change of the counter in the last hour. +func (c *Counter) Rate1h() float64 { + c.mu.RLock() + defer c.mu.RUnlock() + c.advance() + return c.ts[hour].rate() +} + +// Rate10m returns the rate of change of the counter in the last 10 minutes. +func (c *Counter) Rate10m() float64 { + c.mu.RLock() + defer c.mu.RUnlock() + c.advance() + return c.ts[tenminutes].rate() +} + +// Rate1m returns the rate of change of the counter in the last minute. +func (c *Counter) Rate1m() float64 { + c.mu.RLock() + defer c.mu.RUnlock() + c.advance() + return c.ts[minute].rate() +} + +// Reset resets the counter to an empty state. +func (c *Counter) Reset() { + c.mu.Lock() + defer c.mu.Unlock() + now := TimeNow() + for _, ts := range c.ts { + ts.reset(now) + } +} diff --git a/benchmark/stats/histogram.go b/benchmark/stats/histogram.go new file mode 100644 index 00000000..30381182 --- /dev/null +++ b/benchmark/stats/histogram.go @@ -0,0 +1,255 @@ +package stats + +import ( + "bytes" + "fmt" + "io" + "strconv" + "strings" + "time" +) + +// HistogramValue is the value of Histogram objects. +type HistogramValue struct { + // Count is the total number of values added to the histogram. + Count int64 + // Sum is the sum of all the values added to the histogram. + Sum int64 + // Min is the minimum of all the values added to the histogram. + Min int64 + // Max is the maximum of all the values added to the histogram. + Max int64 + // Buckets contains all the buckets of the histogram. + Buckets []HistogramBucket +} + +// HistogramBucket is one histogram bucket. +type HistogramBucket struct { + // LowBound is the lower bound of the bucket. + LowBound int64 + // Count is the number of values in the bucket. + Count int64 +} + +// Print writes textual output of the histogram values. +func (v HistogramValue) Print(w io.Writer) { + avg := float64(v.Sum) / float64(v.Count) + fmt.Fprintf(w, "Count: %d Min: %d Max: %d Avg: %.2f\n", v.Count, v.Min, v.Max, avg) + fmt.Fprintf(w, "%s\n", strings.Repeat("-", 60)) + if v.Count <= 0 { + return + } + + maxBucketDigitLen := len(strconv.FormatInt(v.Buckets[len(v.Buckets)-1].LowBound, 10)) + if maxBucketDigitLen < 3 { + // For "inf". + maxBucketDigitLen = 3 + } + maxCountDigitLen := len(strconv.FormatInt(v.Count, 10)) + percentMulti := 100 / float64(v.Count) + + accCount := int64(0) + for i, b := range v.Buckets { + fmt.Fprintf(w, "[%*d, ", maxBucketDigitLen, b.LowBound) + if i+1 < len(v.Buckets) { + fmt.Fprintf(w, "%*d)", maxBucketDigitLen, v.Buckets[i+1].LowBound) + } else { + fmt.Fprintf(w, "%*s)", maxBucketDigitLen, "inf") + } + + accCount += b.Count + fmt.Fprintf(w, " %*d %5.1f%% %5.1f%%", maxCountDigitLen, b.Count, float64(b.Count)*percentMulti, float64(accCount)*percentMulti) + + const barScale = 0.1 + barLength := int(float64(b.Count)*percentMulti*barScale + 0.5) + fmt.Fprintf(w, " %s\n", strings.Repeat("#", barLength)) + } +} + +// String returns the textual output of the histogram values as string. +func (v HistogramValue) String() string { + var b bytes.Buffer + v.Print(&b) + return b.String() +} + +// A Histogram accumulates values in the form of a histogram. The type of the +// values is int64, which is suitable for keeping track of things like RPC +// latency in milliseconds. New histogram objects should be obtained via the +// New() function. +type Histogram struct { + opts HistogramOptions + buckets []bucketInternal + count *Counter + sum *Counter + tracker *Tracker +} + +// HistogramOptions contains the parameters that define the histogram's buckets. +type HistogramOptions struct { + // NumBuckets is the number of buckets. + NumBuckets int + // GrowthFactor is the growth factor of the buckets. A value of 0.1 + // indicates that bucket N+1 will be 10% larger than bucket N. + GrowthFactor float64 + // SmallestBucketSize is the size of the first bucket. Bucket sizes are + // rounded down to the nearest integer. + SmallestBucketSize float64 + // MinValue is the lower bound of the first bucket. + MinValue int64 +} + +// bucketInternal is the internal representation of a bucket, which includes a +// rate counter. +type bucketInternal struct { + lowBound int64 + count *Counter +} + +// NewHistogram returns a pointer to a new Histogram object that was created +// with the provided options. +func NewHistogram(opts HistogramOptions) *Histogram { + if opts.NumBuckets == 0 { + opts.NumBuckets = 32 + } + if opts.SmallestBucketSize == 0.0 { + opts.SmallestBucketSize = 1.0 + } + h := Histogram{ + opts: opts, + buckets: make([]bucketInternal, opts.NumBuckets), + count: newCounter(), + sum: newCounter(), + tracker: newTracker(), + } + low := opts.MinValue + delta := opts.SmallestBucketSize + for i := 0; i < opts.NumBuckets; i++ { + h.buckets[i].lowBound = low + h.buckets[i].count = newCounter() + low = low + int64(delta) + delta = delta * (1.0 + opts.GrowthFactor) + } + return &h +} + +// Opts returns a copy of the options used to create the Histogram. +func (h *Histogram) Opts() HistogramOptions { + return h.opts +} + +// Add adds a value to the histogram. +func (h *Histogram) Add(value int64) error { + bucket, err := h.findBucket(value) + if err != nil { + return err + } + h.buckets[bucket].count.Incr(1) + h.count.Incr(1) + h.sum.Incr(value) + h.tracker.Push(value) + return nil +} + +// LastUpdate returns the time at which the object was last updated. +func (h *Histogram) LastUpdate() time.Time { + return h.count.LastUpdate() +} + +// Value returns the accumulated state of the histogram since it was created. +func (h *Histogram) Value() HistogramValue { + b := make([]HistogramBucket, len(h.buckets)) + for i, v := range h.buckets { + b[i] = HistogramBucket{ + LowBound: v.lowBound, + Count: v.count.Value(), + } + } + + v := HistogramValue{ + Count: h.count.Value(), + Sum: h.sum.Value(), + Min: h.tracker.Min(), + Max: h.tracker.Max(), + Buckets: b, + } + return v +} + +// Delta1h returns the change in the last hour. +func (h *Histogram) Delta1h() HistogramValue { + b := make([]HistogramBucket, len(h.buckets)) + for i, v := range h.buckets { + b[i] = HistogramBucket{ + LowBound: v.lowBound, + Count: v.count.Delta1h(), + } + } + + v := HistogramValue{ + Count: h.count.Delta1h(), + Sum: h.sum.Delta1h(), + Min: h.tracker.Min1h(), + Max: h.tracker.Max1h(), + Buckets: b, + } + return v +} + +// Delta10m returns the change in the last 10 minutes. +func (h *Histogram) Delta10m() HistogramValue { + b := make([]HistogramBucket, len(h.buckets)) + for i, v := range h.buckets { + b[i] = HistogramBucket{ + LowBound: v.lowBound, + Count: v.count.Delta10m(), + } + } + + v := HistogramValue{ + Count: h.count.Delta10m(), + Sum: h.sum.Delta10m(), + Min: h.tracker.Min10m(), + Max: h.tracker.Max10m(), + Buckets: b, + } + return v +} + +// Delta1m returns the change in the last 10 minutes. +func (h *Histogram) Delta1m() HistogramValue { + b := make([]HistogramBucket, len(h.buckets)) + for i, v := range h.buckets { + b[i] = HistogramBucket{ + LowBound: v.lowBound, + Count: v.count.Delta1m(), + } + } + + v := HistogramValue{ + Count: h.count.Delta1m(), + Sum: h.sum.Delta1m(), + Min: h.tracker.Min1m(), + Max: h.tracker.Max1m(), + Buckets: b, + } + return v +} + +// findBucket does a binary search to find in which bucket the value goes. +func (h *Histogram) findBucket(value int64) (int, error) { + lastBucket := len(h.buckets) - 1 + min, max := 0, lastBucket + for max >= min { + b := (min + max) / 2 + if value >= h.buckets[b].lowBound && (b == lastBucket || value < h.buckets[b+1].lowBound) { + return b, nil + } + if value < h.buckets[b].lowBound { + max = b - 1 + continue + } + min = b + 1 + } + return 0, fmt.Errorf("no bucket for value: %f", value) +} diff --git a/benchmark/stats/stats.go b/benchmark/stats/stats.go new file mode 100644 index 00000000..4290ad77 --- /dev/null +++ b/benchmark/stats/stats.go @@ -0,0 +1,116 @@ +package stats + +import ( + "bytes" + "fmt" + "io" + "math" + "time" +) + +// Stats is a simple helper for gathering additional statistics like histogram +// during benchmarks. This is not thread safe. +type Stats struct { + numBuckets int + unit time.Duration + min, max int64 + histogram *Histogram + + durations durationSlice + dirty bool +} + +type durationSlice []time.Duration + +// NewStats creates a new Stats instance. If numBuckets is not positive, +// the default value (16) will be used. +func NewStats(numBuckets int) *Stats { + if numBuckets <= 0 { + numBuckets = 16 + } + return &Stats{ + // Use one more bucket for the last unbounded bucket. + numBuckets: numBuckets + 1, + durations: make(durationSlice, 0, 100000), + } +} + +// Add adds an elapsed time per operation to the stats. +func (stats *Stats) Add(d time.Duration) { + stats.durations = append(stats.durations, d) + stats.dirty = true +} + +// Clear resets the stats, removing all values. +func (stats *Stats) Clear() { + stats.durations = stats.durations[:0] + stats.histogram = nil + stats.dirty = false +} + +// maybeUpdate updates internal stat data if there was any newly added +// stats since this was updated. +func (stats *Stats) maybeUpdate() { + if !stats.dirty { + return + } + + stats.min = math.MaxInt64 + stats.max = 0 + for _, d := range stats.durations { + if stats.min > int64(d) { + stats.min = int64(d) + } + if stats.max < int64(d) { + stats.max = int64(d) + } + } + + // Use the largest unit that can represent the minimum time duration. + stats.unit = time.Nanosecond + for _, u := range []time.Duration{time.Microsecond, time.Millisecond, time.Second} { + if stats.min <= int64(u) { + break + } + stats.unit = u + } + + // Adjust the min/max according to the new unit. + stats.min /= int64(stats.unit) + stats.max /= int64(stats.unit) + numBuckets := stats.numBuckets + if n := int(stats.max - stats.min + 1); n < numBuckets { + numBuckets = n + } + stats.histogram = NewHistogram(HistogramOptions{ + NumBuckets: numBuckets, + // max(i.e., Nth lower bound) = min + (1 + growthFactor)^(numBuckets-2). + GrowthFactor: math.Pow(float64(stats.max-stats.min), 1/float64(stats.numBuckets-2)) - 1, + SmallestBucketSize: 1.0, + MinValue: stats.min}) + + for _, d := range stats.durations { + stats.histogram.Add(int64(d / stats.unit)) + } + + stats.dirty = false +} + +// Print writes textual output of the Stats. +func (stats *Stats) Print(w io.Writer) { + stats.maybeUpdate() + + if stats.histogram == nil { + fmt.Fprint(w, "Histogram (empty)\n") + } else { + fmt.Fprintf(w, "Histogram (unit: %s)\n", fmt.Sprintf("%v", stats.unit)[1:]) + stats.histogram.Value().Print(w) + } +} + +// String returns the textual output of the Stats as string. +func (stats *Stats) String() string { + var b bytes.Buffer + stats.Print(&b) + return b.String() +} diff --git a/benchmark/stats/timeseries.go b/benchmark/stats/timeseries.go new file mode 100644 index 00000000..2ba18a4d --- /dev/null +++ b/benchmark/stats/timeseries.go @@ -0,0 +1,154 @@ +package stats + +import ( + "math" + "time" +) + +// timeseries holds the history of a changing value over a predefined period of +// time. +type timeseries struct { + size int // The number of time slots. Equivalent to len(slots). + resolution time.Duration // The time resolution of each slot. + stepCount int64 // The number of intervals seen since creation. + head int // The position of the current time in slots. + time time.Time // The time at the beginning of the current time slot. + slots []int64 // A circular buffer of time slots. +} + +// newTimeSeries returns a newly allocated timeseries that covers the requested +// period with the given resolution. +func newTimeSeries(initialTime time.Time, period, resolution time.Duration) *timeseries { + size := int(period.Nanoseconds()/resolution.Nanoseconds()) + 1 + return ×eries{ + size: size, + resolution: resolution, + stepCount: 1, + time: initialTime, + slots: make([]int64, size), + } +} + +// advanceTimeWithFill moves the timeseries forward to time t and fills in any +// slots that get skipped in the process with the given value. Values older than +// the timeseries period are lost. +func (ts *timeseries) advanceTimeWithFill(t time.Time, value int64) { + advanceTo := t.Truncate(ts.resolution) + if !advanceTo.After(ts.time) { + // This is shortcut for the most common case of a busy counter + // where updates come in many times per ts.resolution. + ts.time = advanceTo + return + } + steps := int(advanceTo.Sub(ts.time).Nanoseconds() / ts.resolution.Nanoseconds()) + ts.stepCount += int64(steps) + if steps > ts.size { + steps = ts.size + } + for steps > 0 { + ts.head = (ts.head + 1) % ts.size + ts.slots[ts.head] = value + steps-- + } + ts.time = advanceTo +} + +// advanceTime moves the timeseries forward to time t and fills in any slots +// that get skipped in the process with the head value. Values older than the +// timeseries period are lost. +func (ts *timeseries) advanceTime(t time.Time) { + ts.advanceTimeWithFill(t, ts.slots[ts.head]) +} + +// set sets the current value of the timeseries. +func (ts *timeseries) set(value int64) { + ts.slots[ts.head] = value +} + +// incr sets the current value of the timeseries. +func (ts *timeseries) incr(delta int64) { + ts.slots[ts.head] += delta +} + +// headValue returns the latest value from the timeseries. +func (ts *timeseries) headValue() int64 { + return ts.slots[ts.head] +} + +// headTime returns the time of the latest value from the timeseries. +func (ts *timeseries) headTime() time.Time { + return ts.time +} + +// tailValue returns the oldest value from the timeseries. +func (ts *timeseries) tailValue() int64 { + if ts.stepCount < int64(ts.size) { + return 0 + } + return ts.slots[(ts.head+1)%ts.size] +} + +// tailTime returns the time of the oldest value from the timeseries. +func (ts *timeseries) tailTime() time.Time { + size := int64(ts.size) + if ts.stepCount < size { + size = ts.stepCount + } + return ts.time.Add(-time.Duration(size-1) * ts.resolution) +} + +// delta returns the difference between the newest and oldest values from the +// timeseries. +func (ts *timeseries) delta() int64 { + return ts.headValue() - ts.tailValue() +} + +// rate returns the rate of change between the oldest and newest values from +// the timeseries in units per second. +func (ts *timeseries) rate() float64 { + deltaTime := ts.headTime().Sub(ts.tailTime()).Seconds() + if deltaTime == 0 { + return 0 + } + return float64(ts.delta()) / deltaTime +} + +// min returns the smallest value from the timeseries. +func (ts *timeseries) min() int64 { + to := ts.size + if ts.stepCount < int64(ts.size) { + to = ts.head + 1 + } + tail := (ts.head + 1) % ts.size + min := int64(math.MaxInt64) + for b := 0; b < to; b++ { + if b != tail && ts.slots[b] < min { + min = ts.slots[b] + } + } + return min +} + +// max returns the largest value from the timeseries. +func (ts *timeseries) max() int64 { + to := ts.size + if ts.stepCount < int64(ts.size) { + to = ts.head + 1 + } + tail := (ts.head + 1) % ts.size + max := int64(math.MinInt64) + for b := 0; b < to; b++ { + if b != tail && ts.slots[b] > max { + max = ts.slots[b] + } + } + return max +} + +// reset resets the timeseries to an empty state. +func (ts *timeseries) reset(t time.Time) { + ts.head = 0 + ts.time = t + ts.stepCount = 1 + ts.slots = make([]int64, ts.size) +} diff --git a/benchmark/stats/tracker.go b/benchmark/stats/tracker.go new file mode 100644 index 00000000..802f7295 --- /dev/null +++ b/benchmark/stats/tracker.go @@ -0,0 +1,159 @@ +package stats + +import ( + "math" + "sync" + "time" +) + +// Tracker is a min/max value tracker that keeps track of its min/max values +// over a given period of time, and with a given resolution. The initial min +// and max values are math.MaxInt64 and math.MinInt64 respectively. +type Tracker struct { + mu sync.RWMutex + min, max int64 // All time min/max. + minTS, maxTS [3]*timeseries + lastUpdate time.Time +} + +// newTracker returns a new Tracker. +func newTracker() *Tracker { + now := TimeNow() + t := &Tracker{} + t.minTS[hour] = newTimeSeries(now, time.Hour, time.Minute) + t.minTS[tenminutes] = newTimeSeries(now, 10*time.Minute, 10*time.Second) + t.minTS[minute] = newTimeSeries(now, time.Minute, time.Second) + t.maxTS[hour] = newTimeSeries(now, time.Hour, time.Minute) + t.maxTS[tenminutes] = newTimeSeries(now, 10*time.Minute, 10*time.Second) + t.maxTS[minute] = newTimeSeries(now, time.Minute, time.Second) + t.init() + return t +} + +func (t *Tracker) init() { + t.min = math.MaxInt64 + t.max = math.MinInt64 + for _, ts := range t.minTS { + ts.set(math.MaxInt64) + } + for _, ts := range t.maxTS { + ts.set(math.MinInt64) + } +} + +func (t *Tracker) advance() time.Time { + now := TimeNow() + for _, ts := range t.minTS { + ts.advanceTimeWithFill(now, math.MaxInt64) + } + for _, ts := range t.maxTS { + ts.advanceTimeWithFill(now, math.MinInt64) + } + return now +} + +// LastUpdate returns the last update time of the range. +func (t *Tracker) LastUpdate() time.Time { + t.mu.RLock() + defer t.mu.RUnlock() + return t.lastUpdate +} + +// Push adds a new value if it is a new minimum or maximum. +func (t *Tracker) Push(value int64) { + t.mu.Lock() + defer t.mu.Unlock() + t.lastUpdate = t.advance() + if t.min > value { + t.min = value + } + if t.max < value { + t.max = value + } + for _, ts := range t.minTS { + if ts.headValue() > value { + ts.set(value) + } + } + for _, ts := range t.maxTS { + if ts.headValue() < value { + ts.set(value) + } + } +} + +// Min returns the minimum value of the tracker +func (t *Tracker) Min() int64 { + t.mu.RLock() + defer t.mu.RUnlock() + return t.min +} + +// Max returns the maximum value of the tracker. +func (t *Tracker) Max() int64 { + t.mu.RLock() + defer t.mu.RUnlock() + return t.max +} + +// Min1h returns the minimum value for the last hour. +func (t *Tracker) Min1h() int64 { + t.mu.Lock() + defer t.mu.Unlock() + t.advance() + return t.minTS[hour].min() +} + +// Max1h returns the maximum value for the last hour. +func (t *Tracker) Max1h() int64 { + t.mu.Lock() + defer t.mu.Unlock() + t.advance() + return t.maxTS[hour].max() +} + +// Min10m returns the minimum value for the last 10 minutes. +func (t *Tracker) Min10m() int64 { + t.mu.Lock() + defer t.mu.Unlock() + t.advance() + return t.minTS[tenminutes].min() +} + +// Max10m returns the maximum value for the last 10 minutes. +func (t *Tracker) Max10m() int64 { + t.mu.Lock() + defer t.mu.Unlock() + t.advance() + return t.maxTS[tenminutes].max() +} + +// Min1m returns the minimum value for the last 1 minute. +func (t *Tracker) Min1m() int64 { + t.mu.Lock() + defer t.mu.Unlock() + t.advance() + return t.minTS[minute].min() +} + +// Max1m returns the maximum value for the last 1 minute. +func (t *Tracker) Max1m() int64 { + t.mu.Lock() + defer t.mu.Unlock() + t.advance() + return t.maxTS[minute].max() +} + +// Reset resets the range to an empty state. +func (t *Tracker) Reset() { + t.mu.Lock() + defer t.mu.Unlock() + now := TimeNow() + for _, ts := range t.minTS { + ts.reset(now) + } + for _, ts := range t.maxTS { + ts.reset(now) + } + t.init() +} diff --git a/benchmark/stats/util.go b/benchmark/stats/util.go new file mode 100644 index 00000000..a9922f98 --- /dev/null +++ b/benchmark/stats/util.go @@ -0,0 +1,191 @@ +package stats + +import ( + "bufio" + "bytes" + "fmt" + "os" + "runtime" + "sort" + "strings" + "sync" + "testing" +) + +var ( + curB *testing.B + curBenchName string + curStats map[string]*Stats + + orgStdout *os.File + nextOutPos int + + injectCond *sync.Cond + injectDone chan struct{} +) + +// AddStats adds a new unnamed Stats instance to the current benchmark. You need +// to run benchmarks by calling RunTestMain() to inject the stats to the +// benchmark results. If numBuckets is not positive, the default value (16) will +// be used. Please note that this calls b.ResetTimer() since it may be blocked +// until the previous benchmark stats is printed out. So AddStats() should +// typically be called at the very beginning of each benchmark function. +func AddStats(b *testing.B, numBuckets int) *Stats { + return AddStatsWithName(b, "", numBuckets) +} + +// AddStatsWithName adds a new named Stats instance to the current benchmark. +// With this, you can add multiple stats in a single benchmark. You need +// to run benchmarks by calling RunTestMain() to inject the stats to the +// benchmark results. If numBuckets is not positive, the default value (16) will +// be used. Please note that this calls b.ResetTimer() since it may be blocked +// until the previous benchmark stats is printed out. So AddStatsWithName() +// should typically be called at the very beginning of each benchmark function. +func AddStatsWithName(b *testing.B, name string, numBuckets int) *Stats { + var benchName string + for i := 1; ; i++ { + pc, _, _, ok := runtime.Caller(i) + if !ok { + panic("benchmark function not found") + } + p := strings.Split(runtime.FuncForPC(pc).Name(), ".") + benchName = p[len(p)-1] + if strings.HasPrefix(benchName, "Benchmark") { + break + } + } + procs := runtime.GOMAXPROCS(-1) + if procs != 1 { + benchName = fmt.Sprintf("%s-%d", benchName, procs) + } + + stats := NewStats(numBuckets) + + if injectCond != nil { + // We need to wait until the previous benchmark stats is printed out. + injectCond.L.Lock() + for curB != nil && curBenchName != benchName { + injectCond.Wait() + } + + curB = b + curBenchName = benchName + curStats[name] = stats + + injectCond.L.Unlock() + } + + b.ResetTimer() + return stats +} + +// RunTestMain runs the tests with enabling injection of benchmark stats. It +// returns an exit code to pass to os.Exit. +func RunTestMain(m *testing.M) int { + startStatsInjector() + defer stopStatsInjector() + return m.Run() +} + +// startStatsInjector starts stats injection to benchmark results. +func startStatsInjector() { + orgStdout = os.Stdout + r, w, _ := os.Pipe() + os.Stdout = w + nextOutPos = 0 + + resetCurBenchStats() + + injectCond = sync.NewCond(&sync.Mutex{}) + injectDone = make(chan struct{}) + go func() { + defer close(injectDone) + + scanner := bufio.NewScanner(r) + scanner.Split(splitLines) + for scanner.Scan() { + injectStatsIfFinished(scanner.Text()) + } + if err := scanner.Err(); err != nil { + panic(err) + } + }() +} + +// stopStatsInjector stops stats injection and restores os.Stdout. +func stopStatsInjector() { + os.Stdout.Close() + <-injectDone + injectCond = nil + os.Stdout = orgStdout +} + +// splitLines is a split function for a bufio.Scanner that returns each line +// of text, teeing texts to the original stdout even before each line ends. +func splitLines(data []byte, eof bool) (advance int, token []byte, err error) { + if eof && len(data) == 0 { + return 0, nil, nil + } + + if i := bytes.IndexByte(data, '\n'); i >= 0 { + orgStdout.Write(data[nextOutPos : i+1]) + nextOutPos = 0 + return i + 1, data[0:i], nil + } + + orgStdout.Write(data[nextOutPos:]) + nextOutPos = len(data) + + if eof { + // This is a final, non-terminated line. Return it. + return len(data), data, nil + } + + return 0, nil, nil +} + +// injectStatsIfFinished prints out the stats if the current benchmark finishes. +func injectStatsIfFinished(line string) { + injectCond.L.Lock() + defer injectCond.L.Unlock() + + // We assume that the benchmark results start with the benchmark name. + if curB == nil || !strings.HasPrefix(line, curBenchName) { + return + } + + if !curB.Failed() { + // Output all stats in alphabetical order. + names := make([]string, 0, len(curStats)) + for name := range curStats { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + stats := curStats[name] + // The output of stats starts with a header like "Histogram (unit: ms)" + // followed by statistical properties and the buckets. Add the stats name + // if it is a named stats and indent them as Go testing outputs. + lines := strings.Split(stats.String(), "\n") + if n := len(lines); n > 0 { + if name != "" { + name = ": " + name + } + fmt.Fprintf(orgStdout, "--- %s%s\n", lines[0], name) + for _, line := range lines[1 : n-1] { + fmt.Fprintf(orgStdout, "\t%s\n", line) + } + } + } + } + + resetCurBenchStats() + injectCond.Signal() +} + +// resetCurBenchStats resets the current benchmark stats. +func resetCurBenchStats() { + curB = nil + curBenchName = "" + curStats = make(map[string]*Stats) +}