diff --git a/autoinstrument/init.go b/autoinstrument/init.go index d76e4b7f..68f41d24 100644 --- a/autoinstrument/init.go +++ b/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 } diff --git a/instrumentation/testing/init.go b/instrumentation/testing/init.go index 062c1e0f..038f0566 100644 --- a/instrumentation/testing/init.go +++ b/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 } } diff --git a/instrumentation/testing/logger.go b/instrumentation/testing/logger.go index 26ef389c..53d248e1 100644 --- a/instrumentation/testing/logger.go +++ b/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 {