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