vendor/google.golang.org/grpc/benchmark/stats/util.go

   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
  19 package stats
  20
  21 import (
  22         "bufio"
  23         "bytes"
  24         "fmt"
  25         "os"
  26         "runtime"
  27         "sort"
  28         "strings"
  29         "sync"
  30         "testing"
  31 )
  32
  33 var (
  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.
  51 func 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.
  62 func 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.
 102 func RunTestMain(m *testing.M) int {
 103         startStatsInjector()
 104         defer stopStatsInjector()
 105         return m.Run()
 106 }
 107
 108 // startStatsInjector starts stats injection to benchmark results.
 109 func 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.
 134 func 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.
 143 func 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.
 166 func 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.
 204 func resetCurBenchStats() {
 205         curB = nil
 206         curBenchName = ""
 207         curStats = make(map[string]*Stats)
 208 }