changes

Author: tonyredondo

autoinstrument/init.go

@@ -18,31 +18,39 @@ var (
 
 func init() {
 	once.Do(func() {
-		var m *testing.M
-		var mRunMethod reflect.Method
-		var ok bool
-		mType := reflect.TypeOf(m)
-		if mRunMethod, ok = mType.MethodByName("Run"); !ok {
-			return
-		}
-
-		var runPatch *mpatch.Patch
-		var err error
-		runPatch, err = mpatch.PatchMethodByReflect(mRunMethod, func(m *testing.M) int {
-			logOnError(runPatch.Unpatch())
-			defer func() {
-				logOnError(runPatch.Patch())
-			}()
-			scopetesting.PatchTestingLogger()
-			defer scopetesting.UnpatchTestingLogger()
-			return scopeagent.Run(m)
-		})
-		logOnError(err)
+		patchMRun()
+		scopetesting.PatchTRun()
+		scopetesting.PatchBRun()
 	})
 }
 
-func logOnError(err error) {
+func patchMRun() {
+	var m *testing.M
+	var mRunMethod reflect.Method
+	var ok bool
+	mType := reflect.TypeOf(m)
+	if mRunMethod, ok = mType.MethodByName("Run"); !ok {
+		return
+	}
+
+	var runPatch *mpatch.Patch
+	var err error
+	runPatch, err = mpatch.PatchMethodByReflect(mRunMethod, func(m *testing.M) int {
+		logOnError(runPatch.Unpatch())
+		defer func() {
+			logOnError(runPatch.Patch())
+		}()
+		scopetesting.PatchTestingLogger()
+		defer scopetesting.UnpatchTestingLogger()
+		return scopeagent.Run(m)
+	})
+	logOnError(err)
+}
+
+func logOnError(err error) bool {
 	if err != nil {
 		instrumentation.Logger().Println(err)
+		return true
 	}
+	return false
 }

instrumentation/testing/init.go

@@ -1,7 +1,6 @@
 package testing
 
 import (
-	"os"
 	"reflect"
 	"runtime"
 	"testing"
@@ -11,6 +10,11 @@ import (
 	"go.undefinedlabs.com/scopeagent/reflection"
 )
 
+var (
+	tRunPatched bool
+	bRunPatched bool
+)
+
 // Initialize the testing instrumentation
 func Init(m *testing.M) {
 	if tPointer, err := reflection.GetFieldPointerOf(m, "tests"); err == nil {
@@ -41,7 +45,7 @@ func Init(m *testing.M) {
 			benchmarks = append(benchmarks, testing.InternalBenchmark{
 				Name: benchmark.Name,
 				F: func(b *testing.B) { // Indirection of the original benchmark
-					if envDMPatch, set := os.LookupEnv("SCOPE_DISABLE_MONKEY_PATCHING"); !set || envDMPatch == "" {
+					if bRunPatched {
 						funcValue(b)
 					} else {
 						startBenchmark(b, funcPointer, funcValue)
@@ -51,36 +55,50 @@ func Init(m *testing.M) {
 		}
 		*intBenchmarks = benchmarks
 	}
+}
 
-	if envDMPatch, set := os.LookupEnv("SCOPE_DISABLE_MONKEY_PATCHING"); !set || envDMPatch == "" {
-		// We monkey patch the `testing.T.Run()` func to auto instrument sub tests
-		var t *testing.T
-		tType := reflect.TypeOf(t)
-		if tRunMethod, ok := tType.MethodByName("Run"); ok {
-			_, err := mpatch.PatchMethodByReflect(tRunMethod, func(t *testing.T, name string, f func(t *testing.T)) bool {
-				pc, _, _, _ := runtime.Caller(1)
-				gT := FromTestingT(t)
-				return gT.Run(name, func(childT *testing.T) {
-					addAutoInstrumentedTest(childT)
-					childTest := StartTestFromCaller(childT, pc)
-					defer childTest.end()
-					f(childT)
-				})
-			})
-			logOnError(err)
-		}
+func PatchTRun() {
+	// We monkey patch the `testing.T.Run()` func to auto instrument sub tests
+	var t *testing.T
+	var tRunMethod reflect.Method
+	var ok bool
+	tType := reflect.TypeOf(t)
+	if tRunMethod, ok = tType.MethodByName("Run"); !ok {
+		return
+	}
 
-		// 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)
-		}
+	_, err := mpatch.PatchMethodByReflect(tRunMethod, func(t *testing.T, name string, f func(t *testing.T)) bool {
+		pc, _, _, _ := runtime.Caller(1)
+		gT := FromTestingT(t)
+		return gT.Run(name, func(childT *testing.T) {
+			addAutoInstrumentedTest(childT)
+			childTest := StartTestFromCaller(childT, pc)
+			defer childTest.end()
+			f(childT)
+		})
+	})
+	if !logOnError(err) {
+		tRunPatched = true
+	}
+}
+
+func PatchBRun() {
+	// We monkey patch the `testing.B.Run()` func to auto instrument sub benchmark
+	var b *testing.B
+	var bRunMethod reflect.Method
+	var ok bool
+	bType := reflect.TypeOf(b)
+	if bRunMethod, ok = bType.MethodByName("Run"); !ok {
+		return
+	}
+
+	_, 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)
+		})
+	})
+	if !logOnError(err) {
+		bRunPatched = true
 	}
 }

instrumentation/testing/logger.go

@@ -149,17 +149,18 @@ func patch(methodName string, methodBody func(test *Test, argsValues []reflect.V
 		methodBody(test, in[1:])
 		return nil
 	})
-	logOnError(err)
-	if err == nil {
+	if !logOnError(err) {
 		patches[methodName] = methodPatch
 		patchPointers[reflect.ValueOf(methodBody).Pointer()] = true
 	}
 }
 
-func logOnError(err error) {
+func logOnError(err error) bool {
 	if err != nil {
 		instrumentation.Logger().Println(err)
+		return true
 	}
+	return false
 }
 
 func isAPatchPointer(ptr uintptr) bool {