go.undefinedlabs.com/scopeagent@v0.4.2/instrumentation/testing/benchmark.go (about)

     1  package testing
     2  
     3  import (
     4  	"context"
     5  	"go.undefinedlabs.com/scopeagent/tags"
     6  	"math"
     7  	"regexp"
     8  	"runtime"
     9  	"strings"
    10  	"sync"
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/opentracing/opentracing-go"
    15  
    16  	"go.undefinedlabs.com/scopeagent/instrumentation"
    17  	"go.undefinedlabs.com/scopeagent/reflection"
    18  	"go.undefinedlabs.com/scopeagent/runner"
    19  )
    20  
    21  type (
    22  	Benchmark struct {
    23  		b *testing.B
    24  	}
    25  )
    26  
    27  var (
    28  	benchmarkMapMutex  sync.RWMutex
    29  	benchmarkMap       = map[*testing.B]*Benchmark{}
    30  	benchmarkNameRegex = regexp.MustCompile(`([\w-_:!@#\$%&()=]*)(\/\*\&\/)?`)
    31  )
    32  
    33  // Starts a new benchmark using a pc as caller
    34  func StartBenchmark(b *testing.B, pc uintptr, benchFunc func(b *testing.B)) {
    35  	if !hasBenchmark(b) {
    36  		// If the current benchmark is not instrumented, we instrument it.
    37  		startBenchmark(b, pc, benchFunc)
    38  	} else {
    39  		// If the benchmark is already instrumented, we passthrough to the benchFunc
    40  		benchFunc(b)
    41  	}
    42  }
    43  
    44  // Runs an auto instrumented sub benchmark
    45  func (bench *Benchmark) Run(name string, f func(b *testing.B)) bool {
    46  	pc, _, _, _ := runtime.Caller(1)
    47  	return bench.b.Run(name, func(innerB *testing.B) {
    48  		startBenchmark(innerB, pc, f)
    49  	})
    50  }
    51  
    52  // Adds a benchmark struct to the map
    53  func addBenchmark(b *testing.B, value *Benchmark) {
    54  	benchmarkMapMutex.Lock()
    55  	defer benchmarkMapMutex.Unlock()
    56  	benchmarkMap[b] = value
    57  }
    58  
    59  // Gets if the benchmark struct exist
    60  func hasBenchmark(b *testing.B) bool {
    61  	benchmarkMapMutex.RLock()
    62  	defer benchmarkMapMutex.RUnlock()
    63  	_, ok := benchmarkMap[b]
    64  	return ok
    65  }
    66  
    67  // Gets the Benchmark struct from *testing.Benchmark
    68  func GetBenchmark(b *testing.B) *Benchmark {
    69  	benchmarkMapMutex.RLock()
    70  	defer benchmarkMapMutex.RUnlock()
    71  	if bench, ok := benchmarkMap[b]; ok {
    72  		return bench
    73  	}
    74  	return &Benchmark{b: b}
    75  }
    76  
    77  func startBenchmark(b *testing.B, pc uintptr, benchFunc func(b *testing.B)) {
    78  	var bChild *testing.B
    79  	b.ReportAllocs()
    80  	b.ResetTimer()
    81  	startTime := time.Now()
    82  	result := b.Run("*&", func(b1 *testing.B) {
    83  		addBenchmark(b1, &Benchmark{b: b1})
    84  		benchFunc(b1)
    85  		bChild = b1
    86  	})
    87  	if bChild == nil {
    88  		return
    89  	}
    90  	if reflection.GetBenchmarkHasSub(bChild) > 0 {
    91  		return
    92  	}
    93  	results, err := reflection.GetBenchmarkResult(bChild)
    94  	if err != nil {
    95  		instrumentation.Logger().Printf("Error while extracting the benchmark result object: %v\n", err)
    96  		return
    97  	}
    98  
    99  	// Extracting the benchmark func name (by removing any possible sub-benchmark suffix `{bench_func}/{sub_benchmark}`)
   100  	// to search the func source code bounds and to calculate the package name.
   101  	fullTestName := runner.GetOriginalTestName(b.Name())
   102  
   103  	// We detect if the parent benchmark is instrumented, and if so we remove the "*" SubBenchmark from the previous instrumentation
   104  	parentBenchmark := reflection.GetParentBenchmark(b)
   105  	if parentBenchmark != nil && hasBenchmark(parentBenchmark) {
   106  		var nameSegments []string
   107  		for _, match := range benchmarkNameRegex.FindAllStringSubmatch(fullTestName, -1) {
   108  			if match[1] != "" {
   109  				nameSegments = append(nameSegments, match[1])
   110  			}
   111  		}
   112  		fullTestName = strings.Join(nameSegments, "/")
   113  	}
   114  	packageName := reflection.GetBenchmarkSuiteName(b)
   115  	pName, _, tCode := instrumentation.GetPackageAndNameAndBoundaries(pc)
   116  
   117  	if packageName == "" {
   118  		packageName = pName
   119  	}
   120  
   121  	oTags := opentracing.Tags{
   122  		"span.kind":      "test",
   123  		"test.name":      fullTestName,
   124  		"test.suite":     packageName,
   125  		"test.framework": "testing",
   126  		"test.language":  "go",
   127  		"test.type":      "benchmark",
   128  	}
   129  
   130  	if tCode != "" {
   131  		oTags["test.code"] = tCode
   132  	}
   133  
   134  	var startOptions []opentracing.StartSpanOption
   135  	startOptions = append(startOptions, oTags, opentracing.StartTime(startTime))
   136  
   137  	span, _ := opentracing.StartSpanFromContextWithTracer(context.Background(), instrumentation.Tracer(), fullTestName, startOptions...)
   138  	span.SetBaggageItem("trace.kind", "test")
   139  	avg := math.Round((float64(results.T.Nanoseconds())/float64(results.N))*100) / 100
   140  	span.SetTag("benchmark.runs", results.N)
   141  	span.SetTag("benchmark.duration.mean", avg)
   142  	span.SetTag("benchmark.memory.mean_allocations", results.AllocsPerOp())
   143  	span.SetTag("benchmark.memory.mean_bytes_allocations", results.AllocedBytesPerOp())
   144  	if result {
   145  		span.SetTag("test.status", tags.TestStatus_PASS)
   146  	} else {
   147  		span.SetTag("test.status", tags.TestStatus_FAIL)
   148  	}
   149  	span.FinishWithOptions(opentracing.FinishOptions{
   150  		FinishTime: startTime.Add(results.T),
   151  	})
   152  }