1/*
2 *
3 * Copyright 2017 gRPC authors.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 *
17 */
18
19package stats
20
21import (
22	"bufio"
23	"bytes"
24	"fmt"
25	"os"
26	"runtime"
27	"sort"
28	"strings"
29	"sync"
30	"testing"
31)
32
33var (
34	curB         *testing.B
35	curBenchName string
36	curStats     map[string]*Stats
37
38	orgStdout  *os.File
39	nextOutPos int
40
41	injectCond *sync.Cond
42	injectDone chan struct{}
43)
44
45// AddStats adds a new unnamed Stats instance to the current benchmark. You need
46// to run benchmarks by calling RunTestMain() to inject the stats to the
47// benchmark results. If numBuckets is not positive, the default value (16) will
48// be used. Please note that this calls b.ResetTimer() since it may be blocked
49// until the previous benchmark stats is printed out. So AddStats() should
50// typically be called at the very beginning of each benchmark function.
51func AddStats(b *testing.B, numBuckets int) *Stats {
52	return AddStatsWithName(b, "", numBuckets)
53}
54
55// AddStatsWithName adds a new named Stats instance to the current benchmark.
56// With this, you can add multiple stats in a single benchmark. You need
57// to run benchmarks by calling RunTestMain() to inject the stats to the
58// benchmark results. If numBuckets is not positive, the default value (16) will
59// be used. Please note that this calls b.ResetTimer() since it may be blocked
60// until the previous benchmark stats is printed out. So AddStatsWithName()
61// should typically be called at the very beginning of each benchmark function.
62func AddStatsWithName(b *testing.B, name string, numBuckets int) *Stats {
63	var benchName string
64	for i := 1; ; i++ {
65		pc, _, _, ok := runtime.Caller(i)
66		if !ok {
67			panic("benchmark function not found")
68		}
69		p := strings.Split(runtime.FuncForPC(pc).Name(), ".")
70		benchName = p[len(p)-1]
71		if strings.HasPrefix(benchName, "run") {
72			break
73		}
74	}
75	procs := runtime.GOMAXPROCS(-1)
76	if procs != 1 {
77		benchName = fmt.Sprintf("%s-%d", benchName, procs)
78	}
79
80	stats := NewStats(numBuckets)
81
82	if injectCond != nil {
83		// We need to wait until the previous benchmark stats is printed out.
84		injectCond.L.Lock()
85		for curB != nil && curBenchName != benchName {
86			injectCond.Wait()
87		}
88
89		curB = b
90		curBenchName = benchName
91		curStats[name] = stats
92
93		injectCond.L.Unlock()
94	}
95
96	b.ResetTimer()
97	return stats
98}
99
100// RunTestMain runs the tests with enabling injection of benchmark stats. It
101// returns an exit code to pass to os.Exit.
102func RunTestMain(m *testing.M) int {
103	startStatsInjector()
104	defer stopStatsInjector()
105	return m.Run()
106}
107
108// startStatsInjector starts stats injection to benchmark results.
109func startStatsInjector() {
110	orgStdout = os.Stdout
111	r, w, _ := os.Pipe()
112	os.Stdout = w
113	nextOutPos = 0
114
115	resetCurBenchStats()
116
117	injectCond = sync.NewCond(&sync.Mutex{})
118	injectDone = make(chan struct{})
119	go func() {
120		defer close(injectDone)
121
122		scanner := bufio.NewScanner(r)
123		scanner.Split(splitLines)
124		for scanner.Scan() {
125			injectStatsIfFinished(scanner.Text())
126		}
127		if err := scanner.Err(); err != nil {
128			panic(err)
129		}
130	}()
131}
132
133// stopStatsInjector stops stats injection and restores os.Stdout.
134func stopStatsInjector() {
135	os.Stdout.Close()
136	<-injectDone
137	injectCond = nil
138	os.Stdout = orgStdout
139}
140
141// splitLines is a split function for a bufio.Scanner that returns each line
142// of text, teeing texts to the original stdout even before each line ends.
143func splitLines(data []byte, eof bool) (advance int, token []byte, err error) {
144	if eof && len(data) == 0 {
145		return 0, nil, nil
146	}
147
148	if i := bytes.IndexByte(data, '\n'); i >= 0 {
149		orgStdout.Write(data[nextOutPos : i+1])
150		nextOutPos = 0
151		return i + 1, data[0:i], nil
152	}
153
154	orgStdout.Write(data[nextOutPos:])
155	nextOutPos = len(data)
156
157	if eof {
158		// This is a final, non-terminated line. Return it.
159		return len(data), data, nil
160	}
161
162	return 0, nil, nil
163}
164
165// injectStatsIfFinished prints out the stats if the current benchmark finishes.
166func injectStatsIfFinished(line string) {
167	injectCond.L.Lock()
168	defer injectCond.L.Unlock()
169	// We assume that the benchmark results start with "Benchmark".
170	if curB == nil || !strings.HasPrefix(line, "Benchmark") {
171		return
172	}
173
174	if !curB.Failed() {
175		// Output all stats in alphabetical order.
176		names := make([]string, 0, len(curStats))
177		for name := range curStats {
178			names = append(names, name)
179		}
180		sort.Strings(names)
181		for _, name := range names {
182			stats := curStats[name]
183			// The output of stats starts with a header like "Histogram (unit: ms)"
184			// followed by statistical properties and the buckets. Add the stats name
185			// if it is a named stats and indent them as Go testing outputs.
186			lines := strings.Split(stats.String(), "\n")
187			if n := len(lines); n > 0 {
188				if name != "" {
189					name = ": " + name
190				}
191				fmt.Fprintf(orgStdout, "--- %s%s\n", lines[0], name)
192				for _, line := range lines[1 : n-1] {
193					fmt.Fprintf(orgStdout, "\t%s\n", line)
194				}
195			}
196		}
197	}
198
199	resetCurBenchStats()
200	injectCond.Signal()
201}
202
203// resetCurBenchStats resets the current benchmark stats.
204func resetCurBenchStats() {
205	curB = nil
206	curBenchName = ""
207	curStats = make(map[string]*Stats)
208}
209