sub test and sub benchmark auto instrumentation

Author: tonyredondo

instrumentation/testing/benchmark.go

@@ -43,7 +43,7 @@ func StartBenchmark(b *testing.B, pc uintptr, benchFunc func(b *testing.B)) {
 // Runs an auto instrumented sub benchmark
 func (bench *Benchmark) Run(name string, f func(b *testing.B)) bool {
 	pc, _, _, _ := runtime.Caller(1)
-	return bench.b.Run(name, func(innerB *testing.B) {
+	return FromTestingB(bench.b).Run(name, func(innerB *testing.B) {
 		startBenchmark(innerB, pc, f)
 	})
 }
@@ -78,7 +78,7 @@ func startBenchmark(b *testing.B, pc uintptr, benchFunc func(b *testing.B)) {
 	b.ReportAllocs()
 	b.ResetTimer()
 	startTime := time.Now()
-	result := b.Run("*&", func(b1 *testing.B) {
+	result := FromTestingB(b).Run("*&", func(b1 *testing.B) {
 		addBenchmark(b1, &Benchmark{b: b1})
 		benchFunc(b1)
 		bChild = b1

instrumentation/testing/go_benchmark.go

@@ -0,0 +1,115 @@
+/*
+	The purpose with this file is to clone the struct alignment of the testing.B struct so we can assign a *testing.B
+	pointer to the *goB to have access to the internal private fields.
+
+	We use this to create a Run clone method to be called from the subtest auto instrumentation
+*/
+
+package testing
+
+import (
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+	"unsafe"
+)
+
+// clone of testing.B struct
+type goB struct {
+	goCommon
+	importPath       string
+	context          *goBenchContext
+	N                int
+	previousN        int
+	previousDuration time.Duration
+	benchFunc        func(b *testing.B)
+	benchTime        goBenchTimeFlag
+	bytes            int64
+	missingBytes     bool
+	timerOn          bool
+	showAllocResult  bool
+	result           testing.BenchmarkResult
+	parallelism      int
+	startAllocs      uint64
+	startBytes       uint64
+	netAllocs        uint64
+	netBytes         uint64
+	extra            map[string]float64
+}
+
+// clone of testing.benchContext struct
+type goBenchContext struct {
+	match  *goMatcher
+	maxLen int
+	extLen int
+}
+
+// clone of testing.benchTimeFlag struct
+type goBenchTimeFlag struct {
+	d time.Duration
+	n int
+}
+
+// Convert *goB to *testing.B
+func (b *goB) ToTestingB() *testing.B {
+	return *(**testing.B)(unsafe.Pointer(&b))
+}
+
+// Convert *testing.B to *goB
+func FromTestingB(b *testing.B) *goB {
+	return *(**goB)(unsafe.Pointer(&b))
+}
+
+//go:linkname benchmarkLock testing.benchmarkLock
+var benchmarkLock sync.Mutex
+
+//go:linkname (*goB).run1 testing.(*B).run1
+func (b *goB) run1() bool
+
+//go:linkname (*goB).run testing.(*B).run
+func (b *goB) run() bool
+
+//go:linkname (*goB).add testing.(*B).add
+func (b *goB) add(other testing.BenchmarkResult)
+
+// we clone the same (*testing.B).Run implementation because the Patch overwrite the original implementation with the jump
+func (b *goB) Run(name string, f func(b *testing.B)) bool {
+	atomic.StoreInt32(&b.hasSub, 1)
+	benchmarkLock.Unlock()
+	defer benchmarkLock.Lock()
+
+	benchName, ok, partial := b.name, true, false
+	if b.context != nil {
+		benchName, ok, partial = b.context.match.fullName(&b.goCommon, name)
+	}
+	if !ok {
+		return true
+	}
+	var pc [maxStackLen]uintptr
+	n := runtime.Callers(2, pc[:])
+	sub := &goB{
+		goCommon: goCommon{
+			signal:  make(chan bool),
+			name:    benchName,
+			parent:  &b.goCommon,
+			level:   b.level + 1,
+			creator: pc[:n],
+			w:       b.w,
+			chatty:  b.chatty,
+		},
+		importPath: b.importPath,
+		benchFunc:  f,
+		benchTime:  b.benchTime,
+		context:    b.context,
+	}
+	if partial {
+		atomic.StoreInt32(&sub.hasSub, 1)
+	}
+	if sub.run1() {
+		sub.run()
+	}
+	b.add(sub.result)
+	return !sub.failed
+}

instrumentation/testing/go_testing.go

@@ -0,0 +1,135 @@
+/*
+	The purpose with this file is to clone the struct alignment of the testing.T struct so we can assign a *testing.T
+	pointer to the *goT to have access to the internal private fields.
+
+	We use this to create a Run clone method to be called from the subtest auto instrumentation
+*/
+package testing
+
+import (
+	"bytes"
+	"fmt"
+	"runtime"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"unsafe"
+)
+
+// clone of testing.T struct
+type goT struct {
+	goCommon
+	isParallel bool
+	context    *goTestContext
+}
+
+// clone of testing.testContext struct
+type goTestContext struct {
+	match         *goMatcher
+	mu            sync.Mutex
+	startParallel chan bool
+	running       int
+	numWaiting    int
+	maxParallel   int
+}
+
+// clone of testing.matcher struct
+type goMatcher struct {
+	filter    []string
+	matchFunc func(pat, str string) (bool, error)
+	mu        sync.Mutex
+	subNames  map[string]int64
+}
+
+// clone of testing.indenter struct
+type goIndenter struct {
+	c *goCommon
+}
+
+// Convert *goT to *testing.T
+func (t *goT) ToTestingT() *testing.T {
+	return *(**testing.T)(unsafe.Pointer(&t))
+}
+
+// Convert *testing.T to *goT
+func FromTestingT(t *testing.T) *goT {
+	return *(**goT)(unsafe.Pointer(&t))
+}
+
+const maxStackLen = 50
+
+//go:linkname matchMutex testing.matchMutex
+var matchMutex sync.Mutex
+
+//go:linkname goTRunner testing.tRunner
+func goTRunner(t *testing.T, fn func(t *testing.T))
+
+//go:linkname rewrite testing.rewrite
+func rewrite(s string) string
+
+//go:linkname shouldFailFast testing.shouldFailFast
+func shouldFailFast() bool
+
+//go:linkname (*goMatcher).fullName testing.(*matcher).fullName
+func (m *goMatcher) fullName(c *goCommon, subname string) (name string, ok, partial bool)
+
+// this method calls the original testing.tRunner by converting *goT to *testing.T
+func tRunner(t *goT, fn func(t *goT)) {
+	goTRunner(t.ToTestingT(), func(t *testing.T) { fn(FromTestingT(t)) })
+}
+
+// we clone the same (*testing.T).Run implementation because the Patch overwrite the original implementation with the jump
+func (t *goT) Run(name string, f func(t *goT)) bool {
+	atomic.StoreInt32(&t.hasSub, 1)
+	testName, ok, _ := t.context.match.fullName(&t.goCommon, name)
+	if !ok || shouldFailFast() {
+		return true
+	}
+	var pc [maxStackLen]uintptr
+	n := runtime.Callers(2, pc[:])
+	t = &goT{
+		goCommon: goCommon{
+			barrier: make(chan bool),
+			signal:  make(chan bool),
+			name:    testName,
+			parent:  &t.goCommon,
+			level:   t.level + 1,
+			creator: pc[:n],
+			chatty:  t.chatty,
+		},
+		context: t.context,
+	}
+	t.w = goIndenter{&t.goCommon}
+
+	if t.chatty {
+		root := t.parent
+		for ; root.parent != nil; root = root.parent {
+		}
+		root.mu.Lock()
+		fmt.Fprintf(root.w, "=== RUN   %s\n", t.name)
+		root.mu.Unlock()
+	}
+	go tRunner(t, f)
+	if !<-t.signal {
+		runtime.Goexit()
+	}
+	return !t.failed
+}
+
+// we can't link an instance method without a struct pointer
+func (w goIndenter) Write(b []byte) (n int, err error) {
+	n = len(b)
+	for len(b) > 0 {
+		end := bytes.IndexByte(b, '\n')
+		if end == -1 {
+			end = len(b)
+		} else {
+			end++
+		}
+		const indent = "    "
+		w.c.output = append(w.c.output, indent...)
+		w.c.output = append(w.c.output, b[:end]...)
+		b = b[end:]
+	}
+	return
+}

instrumentation/testing/go_testing.s

@@ -0,0 +1,4 @@
+// The testing package uses //go:linkname to push a few functions into this
+// package but we still need a .s file so the Go tool does not pass -complete
+// to the go tool compile so the latter does not complain about Go functions
+// with no bodies.

instrumentation/testing/go_testing_go13.go

@@ -0,0 +1,35 @@
+// +build !go1.14
+
+package testing
+
+import (
+	"io"
+	"sync"
+	"time"
+)
+
+// clone of testing.common struct
+type goCommon struct {
+	mu         sync.RWMutex
+	output     []byte
+	w          io.Writer
+	ran        bool
+	failed     bool
+	skipped    bool
+	done       bool
+	helpers    map[string]struct{}
+	chatty     bool
+	finished   bool
+	hasSub     int32
+	raceErrors int
+	runner     string
+	parent     *goCommon
+	level      int
+	creator    []uintptr
+	name       string
+	start      time.Time
+	duration   time.Duration
+	barrier    chan bool
+	signal     chan bool
+	sub        []*goT
+}

instrumentation/testing/go_testing_go14.go

@@ -0,0 +1,36 @@
+// +build go1.14
+
+package testing
+
+import (
+	"io"
+	"sync"
+	"time"
+)
+
+// clone of testing.common struct
+type goCommon struct {
+	mu         sync.RWMutex
+	output     []byte
+	w          io.Writer
+	ran        bool
+	failed     bool
+	skipped    bool
+	done       bool
+	helpers    map[string]struct{}
+	cleanup    func() // New in golang 1.14
+	chatty     bool
+	finished   bool
+	hasSub     int32
+	raceErrors int
+	runner     string
+	parent     *goCommon
+	level      int
+	creator    []uintptr
+	name       string
+	start      time.Time
+	duration   time.Duration
+	barrier    chan bool
+	signal     chan bool
+	sub        []*goT
+}

instrumentation/testing/init.go

@@ -54,7 +54,11 @@ func Init(m *testing.M) {
 			benchmarks = append(benchmarks, testing.InternalBenchmark{
 				Name: benchmark.Name,
 				F: func(b *testing.B) { // Indirection of the original benchmark
-					startBenchmark(b, funcPointer, funcValue)
+					if envDMPatch, set := os.LookupEnv("SCOPE_DISABLE_MONKEY_PATCHING"); !set || envDMPatch == "" {
+						funcValue(b)
+					} else {
+						startBenchmark(b, funcPointer, funcValue)
+					}
 				},
 			})
 		}
@@ -66,14 +70,11 @@ func Init(m *testing.M) {
 		var t *testing.T
 		tType := reflect.TypeOf(t)
 		if tRunMethod, ok := tType.MethodByName("Run"); ok {
-			var runPatch *mpatch.Patch
-			var err error
-			runPatch, err = mpatch.PatchMethodByReflect(tRunMethod, func(t *testing.T, name string, f func(t *testing.T)) bool {
+			_, err := mpatch.PatchMethodByReflect(tRunMethod, func(t *testing.T, name string, f func(t *testing.T)) bool {
 				pc, _, _, _ := runtime.Caller(1)
-				logOnError(runPatch.Unpatch())
-				defer runPatch.Patch()
-				return t.Run(name, func(childT *testing.T) {
-					_ = runPatch.Patch()
+				gT := FromTestingT(t)
+				return gT.Run(name, func(childGoT *goT) {
+					childT := childGoT.ToTestingT()
 					addAutoInstrumentedTest(childT)
 					childTest := StartTestFromCaller(childT, pc)
 					defer childTest.end()
@@ -82,5 +83,18 @@ func Init(m *testing.M) {
 			})
 			logOnError(err)
 		}
+
+		// We monkey patch the `testing.B.Run()` func to auto instrument sub benchmark
+		var b *testing.B
+		bType := reflect.TypeOf(b)
+		if bRunMethod, ok := bType.MethodByName("Run"); ok {
+			_, err := mpatch.PatchMethodByReflect(bRunMethod, func(b *testing.B, name string, f func(b *testing.B)) bool {
+				pc, _, _, _ := runtime.Caller(1)
+				return FromTestingB(b).Run(name, func(b *testing.B) {
+					StartBenchmark(b, pc, f)
+				})
+			})
+			logOnError(err)
+		}
 	}
 }