--- /dev/null
+/*
+ *
+ * Copyright 2017 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+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, "run") {
+ 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 "Benchmark".
+ if curB == nil || !strings.HasPrefix(line, "Benchmark") {
+ 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)
+}
OSDN Git Service
