From c3c4a6adabf26936f1d9c446ee1237589e073490 Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Fri, 6 Sep 2024 15:51:31 +0200
Subject: [PATCH 01/10] internal/civisibility/integrations/gotesting: fixes for
 orchestrion autoinstrumentation

For orchestrion we are going to instrument testing.M.Run(), testing.T.Run() and testing.B.Run() methods, for that we need to make some fixes and add some checks to avoid multiple instrumentation over the same Func
---
 .../integrations/gotesting/instrumentation.go | 73 +++++++++++++------
 .../integrations/gotesting/testing.go         | 15 +++-
 .../integrations/gotesting/testingB.go        |  7 --
 .../integrations/gotesting/testingT.go        | 25 +++++++
 4 files changed, 87 insertions(+), 33 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index 619de787ba..b2dd85a506 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -56,18 +56,33 @@ func instrumentTestingM(m *testing.M) func(exitCode int) {
 
 // instrumentTestingTFunc helper function to instrument a testing function func(*testing.T)
 func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
-	// Reflect the function to obtain its pointer.
-	fReflect := reflect.Indirect(reflect.ValueOf(f))
-	moduleName, suiteName := utils.GetModuleAndSuiteName(fReflect.Pointer())
-	originalFunc := runtime.FuncForPC(fReflect.Pointer())
+	// Avoid instrumenting twice
+	if hasCiVisibilityTestFunc(&f) {
+		return f
+	}
 
-	// Increment the test count in the module.
-	atomic.AddInt32(modulesCounters[moduleName], 1)
+	instrumentedFunc := func(t *testing.T) {
+		// Reflect the function to obtain its pointer.
+		fReflect := reflect.Indirect(reflect.ValueOf(f))
+		moduleName, suiteName := utils.GetModuleAndSuiteName(fReflect.Pointer())
+		originalFunc := runtime.FuncForPC(fReflect.Pointer())
 
-	// Increment the test count in the suite.
-	atomic.AddInt32(suitesCounters[suiteName], 1)
+		// Initialize module counters if not already present.
+		if _, ok := modulesCounters[moduleName]; !ok {
+			var v int32
+			modulesCounters[moduleName] = &v
+		}
+		// Increment the test count in the module.
+		atomic.AddInt32(modulesCounters[moduleName], 1)
+
+		// Initialize suite counters if not already present.
+		if _, ok := suitesCounters[suiteName]; !ok {
+			var v int32
+			suitesCounters[suiteName] = &v
+		}
+		// Increment the test count in the suite.
+		atomic.AddInt32(suitesCounters[suiteName], 1)
 
-	return func(t *testing.T) {
 		// Create or retrieve the module, suite, and test for CI visibility.
 		module := session.GetOrCreateModuleWithFramework(moduleName, testFramework, runtime.Version())
 		suite := module.GetOrCreateSuite(suiteName)
@@ -101,6 +116,9 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 		// Execute the original test function.
 		f(t)
 	}
+
+	setCiVisibilityTestFunc(&instrumentedFunc)
+	return instrumentedFunc
 }
 
 // instrumentTestingTSetErrorInfo helper function to set an error in the `testing.T` CI Visibility span
@@ -134,18 +152,7 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 		return name, f
 	}
 
-	// Reflect the function to obtain its pointer.
-	fReflect := reflect.Indirect(reflect.ValueOf(f))
-	moduleName, suiteName := utils.GetModuleAndSuiteName(fReflect.Pointer())
-	originalFunc := runtime.FuncForPC(fReflect.Pointer())
-
-	// Increment the test count in the module.
-	atomic.AddInt32(modulesCounters[moduleName], 1)
-
-	// Increment the test count in the suite.
-	atomic.AddInt32(suitesCounters[suiteName], 1)
-
-	return subBenchmarkAutoName, func(b *testing.B) {
+	instrumentedFunc := func(b *testing.B) {
 		// The sub-benchmark implementation relies on creating a dummy sub benchmark (called [DD:TestVisibility]) with
 		// a Run over the original sub benchmark function to get the child results without interfering measurements
 		// By doing this the name of the sub-benchmark are changed
@@ -155,6 +162,27 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 		//		benchmark/[DD:TestVisibility]/child
 		// We use regex and decrement the depth level of the benchmark to restore the original name
 
+		// Reflect the function to obtain its pointer.
+		fReflect := reflect.Indirect(reflect.ValueOf(f))
+		moduleName, suiteName := utils.GetModuleAndSuiteName(fReflect.Pointer())
+		originalFunc := runtime.FuncForPC(fReflect.Pointer())
+
+		// Initialize module counters if not already present.
+		if _, ok := modulesCounters[moduleName]; !ok {
+			var v int32
+			modulesCounters[moduleName] = &v
+		}
+		// Increment the test count in the module.
+		atomic.AddInt32(modulesCounters[moduleName], 1)
+
+		// Initialize suite counters if not already present.
+		if _, ok := suitesCounters[suiteName]; !ok {
+			var v int32
+			suitesCounters[suiteName] = &v
+		}
+		// Increment the test count in the suite.
+		atomic.AddInt32(suitesCounters[suiteName], 1)
+
 		// Decrement level.
 		bpf := getBenchmarkPrivateFields(b)
 		bpf.AddLevel(-1)
@@ -201,7 +229,6 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 		}
 
 		setCiVisibilityBenchmarkFunc(&instrumentedFunc)
-		defer deleteCiVisibilityBenchmarkFunc(&instrumentedFunc)
 		b.Run(name, instrumentedFunc)
 
 		endTime := time.Now()
@@ -258,6 +285,8 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 
 		checkModuleAndSuite(module, suite)
 	}
+	setCiVisibilityBenchmarkFunc(&instrumentedFunc)
+	return subBenchmarkAutoName, instrumentedFunc
 }
 
 // instrumentTestingBSetErrorInfo helper function to set an error in the `testing.B` CI Visibility span
diff --git a/internal/civisibility/integrations/gotesting/testing.go b/internal/civisibility/integrations/gotesting/testing.go
index a90a2bcd94..3d1c3e3c6f 100644
--- a/internal/civisibility/integrations/gotesting/testing.go
+++ b/internal/civisibility/integrations/gotesting/testing.go
@@ -128,7 +128,7 @@ func (ddm *M) instrumentInternalTests(internalTests *[]testing.InternalTest) {
 // executeInternalTest wraps the original test function to include CI visibility instrumentation.
 func (ddm *M) executeInternalTest(testInfo *testingTInfo) func(*testing.T) {
 	originalFunc := runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(testInfo.originalFunc)).Pointer())
-	return func(t *testing.T) {
+	instrumentedFunc := func(t *testing.T) {
 		// Create or retrieve the module, suite, and test for CI visibility.
 		module := session.GetOrCreateModuleWithFramework(testInfo.moduleName, testFramework, runtime.Version())
 		suite := module.GetOrCreateSuite(testInfo.suiteName)
@@ -165,6 +165,8 @@ func (ddm *M) executeInternalTest(testInfo *testingTInfo) func(*testing.T) {
 		// Execute the original test function.
 		testInfo.originalFunc(t)
 	}
+	setCiVisibilityTestFunc(&instrumentedFunc)
+	return instrumentedFunc
 }
 
 // instrumentInternalBenchmarks instruments the internal benchmarks for CI visibility.
@@ -216,7 +218,7 @@ func (ddm *M) instrumentInternalBenchmarks(internalBenchmarks *[]testing.Interna
 
 // executeInternalBenchmark wraps the original benchmark function to include CI visibility instrumentation.
 func (ddm *M) executeInternalBenchmark(benchmarkInfo *testingBInfo) func(*testing.B) {
-	return func(b *testing.B) {
+	instrumentedInternalFunc := func(b *testing.B) {
 
 		// decrement level
 		getBenchmarkPrivateFields(b).AddLevel(-1)
@@ -231,7 +233,7 @@ func (ddm *M) executeInternalBenchmark(benchmarkInfo *testingBInfo) func(*testin
 		// Run the original benchmark function.
 		var iPfOfB *benchmarkPrivateFields
 		var recoverFunc *func(r any)
-		b.Run(b.Name(), func(b *testing.B) {
+		instrumentedFunc := func(b *testing.B) {
 			// Stop the timer to perform initialization and replacements.
 			b.StopTimer()
 
@@ -259,7 +261,10 @@ func (ddm *M) executeInternalBenchmark(benchmarkInfo *testingBInfo) func(*testin
 			b.ResetTimer()
 			b.StartTimer()
 			benchmarkInfo.originalFunc(b)
-		})
+		}
+
+		setCiVisibilityBenchmarkFunc(&instrumentedFunc)
+		b.Run(b.Name(), instrumentedFunc)
 
 		endTime := time.Now()
 		results := iPfOfB.result
@@ -315,6 +320,8 @@ func (ddm *M) executeInternalBenchmark(benchmarkInfo *testingBInfo) func(*testin
 
 		checkModuleAndSuite(module, suite)
 	}
+	setCiVisibilityBenchmarkFunc(&instrumentedInternalFunc)
+	return instrumentedInternalFunc
 }
 
 // RunM runs the tests and benchmarks using CI visibility.
diff --git a/internal/civisibility/integrations/gotesting/testingB.go b/internal/civisibility/integrations/gotesting/testingB.go
index 2fc80eafbd..2538daf6ed 100644
--- a/internal/civisibility/integrations/gotesting/testingB.go
+++ b/internal/civisibility/integrations/gotesting/testingB.go
@@ -226,10 +226,3 @@ func setCiVisibilityBenchmarkFunc(fn *func(*testing.B)) {
 	defer civisibilityBenchmarksFuncsMutex.RUnlock()
 	civisibilityBenchmarksFuncs[fn] = struct{}{}
 }
-
-// deleteCiVisibilityBenchmarkFunc untracks a func(*testing.B) as instrumented benchmark.
-func deleteCiVisibilityBenchmarkFunc(fn *func(*testing.B)) {
-	civisibilityBenchmarksFuncsMutex.RLock()
-	defer civisibilityBenchmarksFuncsMutex.RUnlock()
-	delete(civisibilityBenchmarksFuncs, fn)
-}
diff --git a/internal/civisibility/integrations/gotesting/testingT.go b/internal/civisibility/integrations/gotesting/testingT.go
index 40e72a2122..34a15f9c06 100644
--- a/internal/civisibility/integrations/gotesting/testingT.go
+++ b/internal/civisibility/integrations/gotesting/testingT.go
@@ -21,6 +21,12 @@ var (
 
 	// ciVisibilityTestsMutex is a read-write mutex for synchronizing access to ciVisibilityTests.
 	ciVisibilityTestsMutex sync.RWMutex
+
+	// civisibilityTestsFuncs holds a map of *func(*testing.T) for tracking instrumented functions
+	civisibilityTestsFuncs = map[*func(*testing.T)]struct{}{}
+
+	// civisibilityTestsFuncsMutex is a read-write mutex for synchronizing access to civisibilityTestsFuncs.
+	civisibilityTestsFuncsMutex sync.RWMutex
 )
 
 // T is a type alias for testing.T to provide additional methods for CI visibility.
@@ -156,3 +162,22 @@ func setCiVisibilityTest(t *testing.T, ciTest integrations.DdTest) {
 	defer ciVisibilityTestsMutex.Unlock()
 	ciVisibilityTests[t] = ciTest
 }
+
+// hasCiVisibilityTestFunc gets if a func(*testing.T) is being instrumented.
+func hasCiVisibilityTestFunc(fn *func(*testing.T)) bool {
+	civisibilityTestsFuncsMutex.RLock()
+	defer civisibilityTestsFuncsMutex.RUnlock()
+
+	if _, ok := civisibilityTestsFuncs[fn]; ok {
+		return true
+	}
+
+	return false
+}
+
+// setCiVisibilityTestFunc tracks a func(*testing.T) as instrumented benchmark.
+func setCiVisibilityTestFunc(fn *func(*testing.T)) {
+	civisibilityTestsFuncsMutex.RLock()
+	defer civisibilityTestsFuncsMutex.RUnlock()
+	civisibilityTestsFuncs[fn] = struct{}{}
+}

From a97c7b767706f99f2e4f9a44a929dcef172f5159 Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Fri, 6 Sep 2024 16:54:17 +0200
Subject: [PATCH 02/10] internal/civisibility/integrations/gotesting: fixes for
 orchestrion autoinstrumentation

fixes
---
 .../integrations/gotesting/instrumentation.go | 32 ++++++++++---------
 .../integrations/gotesting/testing.go         | 10 +++---
 .../integrations/gotesting/testingB.go        |  7 ++--
 .../integrations/gotesting/testingT.go        |  7 ++--
 4 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index b2dd85a506..6457a62a38 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -56,17 +56,17 @@ func instrumentTestingM(m *testing.M) func(exitCode int) {
 
 // instrumentTestingTFunc helper function to instrument a testing function func(*testing.T)
 func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
+	// Reflect the function to obtain its pointer.
+	fReflect := reflect.Indirect(reflect.ValueOf(f))
+	moduleName, suiteName := utils.GetModuleAndSuiteName(fReflect.Pointer())
+	originalFunc := runtime.FuncForPC(fReflect.Pointer())
+
 	// Avoid instrumenting twice
-	if hasCiVisibilityTestFunc(&f) {
+	if hasCiVisibilityTestFunc(originalFunc) {
 		return f
 	}
 
 	instrumentedFunc := func(t *testing.T) {
-		// Reflect the function to obtain its pointer.
-		fReflect := reflect.Indirect(reflect.ValueOf(f))
-		moduleName, suiteName := utils.GetModuleAndSuiteName(fReflect.Pointer())
-		originalFunc := runtime.FuncForPC(fReflect.Pointer())
-
 		// Initialize module counters if not already present.
 		if _, ok := modulesCounters[moduleName]; !ok {
 			var v int32
@@ -117,7 +117,8 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 		f(t)
 	}
 
-	setCiVisibilityTestFunc(&instrumentedFunc)
+	setCiVisibilityTestFunc(originalFunc)
+	setCiVisibilityTestFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()))
 	return instrumentedFunc
 }
 
@@ -147,8 +148,13 @@ func instrumentTestingTSkipNow(t *testing.T) {
 
 // instrumentTestingBFunc helper function to instrument a benchmark function func(*testing.B)
 func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (string, func(*testing.B)) {
+	// Reflect the function to obtain its pointer.
+	fReflect := reflect.Indirect(reflect.ValueOf(f))
+	moduleName, suiteName := utils.GetModuleAndSuiteName(fReflect.Pointer())
+	originalFunc := runtime.FuncForPC(fReflect.Pointer())
+
 	// Avoid instrumenting twice
-	if hasCiVisibilityBenchmarkFunc(&f) {
+	if hasCiVisibilityBenchmarkFunc(originalFunc) {
 		return name, f
 	}
 
@@ -162,11 +168,6 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 		//		benchmark/[DD:TestVisibility]/child
 		// We use regex and decrement the depth level of the benchmark to restore the original name
 
-		// Reflect the function to obtain its pointer.
-		fReflect := reflect.Indirect(reflect.ValueOf(f))
-		moduleName, suiteName := utils.GetModuleAndSuiteName(fReflect.Pointer())
-		originalFunc := runtime.FuncForPC(fReflect.Pointer())
-
 		// Initialize module counters if not already present.
 		if _, ok := modulesCounters[moduleName]; !ok {
 			var v int32
@@ -228,7 +229,7 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 			f(b)
 		}
 
-		setCiVisibilityBenchmarkFunc(&instrumentedFunc)
+		setCiVisibilityBenchmarkFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()))
 		b.Run(name, instrumentedFunc)
 
 		endTime := time.Now()
@@ -285,7 +286,8 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 
 		checkModuleAndSuite(module, suite)
 	}
-	setCiVisibilityBenchmarkFunc(&instrumentedFunc)
+	setCiVisibilityBenchmarkFunc(originalFunc)
+	setCiVisibilityBenchmarkFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()))
 	return subBenchmarkAutoName, instrumentedFunc
 }
 
diff --git a/internal/civisibility/integrations/gotesting/testing.go b/internal/civisibility/integrations/gotesting/testing.go
index 3d1c3e3c6f..f7f6558b50 100644
--- a/internal/civisibility/integrations/gotesting/testing.go
+++ b/internal/civisibility/integrations/gotesting/testing.go
@@ -165,7 +165,8 @@ func (ddm *M) executeInternalTest(testInfo *testingTInfo) func(*testing.T) {
 		// Execute the original test function.
 		testInfo.originalFunc(t)
 	}
-	setCiVisibilityTestFunc(&instrumentedFunc)
+	setCiVisibilityTestFunc(originalFunc)
+	setCiVisibilityTestFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()))
 	return instrumentedFunc
 }
 
@@ -218,13 +219,13 @@ func (ddm *M) instrumentInternalBenchmarks(internalBenchmarks *[]testing.Interna
 
 // executeInternalBenchmark wraps the original benchmark function to include CI visibility instrumentation.
 func (ddm *M) executeInternalBenchmark(benchmarkInfo *testingBInfo) func(*testing.B) {
+	originalFunc := runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(benchmarkInfo.originalFunc)).Pointer())
 	instrumentedInternalFunc := func(b *testing.B) {
 
 		// decrement level
 		getBenchmarkPrivateFields(b).AddLevel(-1)
 
 		startTime := time.Now()
-		originalFunc := runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(benchmarkInfo.originalFunc)).Pointer())
 		module := session.GetOrCreateModuleWithFrameworkAndStartTime(benchmarkInfo.moduleName, testFramework, runtime.Version(), startTime)
 		suite := module.GetOrCreateSuiteWithStartTime(benchmarkInfo.suiteName, startTime)
 		test := suite.CreateTestWithStartTime(benchmarkInfo.testName, startTime)
@@ -263,7 +264,7 @@ func (ddm *M) executeInternalBenchmark(benchmarkInfo *testingBInfo) func(*testin
 			benchmarkInfo.originalFunc(b)
 		}
 
-		setCiVisibilityBenchmarkFunc(&instrumentedFunc)
+		setCiVisibilityBenchmarkFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()))
 		b.Run(b.Name(), instrumentedFunc)
 
 		endTime := time.Now()
@@ -320,7 +321,8 @@ func (ddm *M) executeInternalBenchmark(benchmarkInfo *testingBInfo) func(*testin
 
 		checkModuleAndSuite(module, suite)
 	}
-	setCiVisibilityBenchmarkFunc(&instrumentedInternalFunc)
+	setCiVisibilityBenchmarkFunc(originalFunc)
+	setCiVisibilityBenchmarkFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedInternalFunc)).Pointer()))
 	return instrumentedInternalFunc
 }
 
diff --git a/internal/civisibility/integrations/gotesting/testingB.go b/internal/civisibility/integrations/gotesting/testingB.go
index 2538daf6ed..9a602f6f24 100644
--- a/internal/civisibility/integrations/gotesting/testingB.go
+++ b/internal/civisibility/integrations/gotesting/testingB.go
@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"regexp"
+	"runtime"
 	"sync"
 	"testing"
 	"time"
@@ -30,7 +31,7 @@ var (
 	subBenchmarkAutoNameRegex = regexp.MustCompile(`(?si)\/\[DD:TestVisibility\].*`)
 
 	// civisibilityBenchmarksFuncs holds a map of *func(*testing.B) for tracking instrumented functions
-	civisibilityBenchmarksFuncs = map[*func(*testing.B)]struct{}{}
+	civisibilityBenchmarksFuncs = map[*runtime.Func]struct{}{}
 
 	// civisibilityBenchmarksFuncsMutex is a read-write mutex for synchronizing access to civisibilityBenchmarksFuncs.
 	civisibilityBenchmarksFuncsMutex sync.RWMutex
@@ -209,7 +210,7 @@ func setCiVisibilityBenchmark(b *testing.B, ciTest integrations.DdTest) {
 }
 
 // hasCiVisibilityBenchmarkFunc gets if a func(*testing.B) is being instrumented.
-func hasCiVisibilityBenchmarkFunc(fn *func(*testing.B)) bool {
+func hasCiVisibilityBenchmarkFunc(fn *runtime.Func) bool {
 	civisibilityBenchmarksFuncsMutex.RLock()
 	defer civisibilityBenchmarksFuncsMutex.RUnlock()
 
@@ -221,7 +222,7 @@ func hasCiVisibilityBenchmarkFunc(fn *func(*testing.B)) bool {
 }
 
 // setCiVisibilityBenchmarkFunc tracks a func(*testing.B) as instrumented benchmark.
-func setCiVisibilityBenchmarkFunc(fn *func(*testing.B)) {
+func setCiVisibilityBenchmarkFunc(fn *runtime.Func) {
 	civisibilityBenchmarksFuncsMutex.RLock()
 	defer civisibilityBenchmarksFuncsMutex.RUnlock()
 	civisibilityBenchmarksFuncs[fn] = struct{}{}
diff --git a/internal/civisibility/integrations/gotesting/testingT.go b/internal/civisibility/integrations/gotesting/testingT.go
index 34a15f9c06..88ef23cd58 100644
--- a/internal/civisibility/integrations/gotesting/testingT.go
+++ b/internal/civisibility/integrations/gotesting/testingT.go
@@ -8,6 +8,7 @@ package gotesting
 import (
 	"context"
 	"fmt"
+	"runtime"
 	"sync"
 	"testing"
 	"time"
@@ -23,7 +24,7 @@ var (
 	ciVisibilityTestsMutex sync.RWMutex
 
 	// civisibilityTestsFuncs holds a map of *func(*testing.T) for tracking instrumented functions
-	civisibilityTestsFuncs = map[*func(*testing.T)]struct{}{}
+	civisibilityTestsFuncs = map[*runtime.Func]struct{}{}
 
 	// civisibilityTestsFuncsMutex is a read-write mutex for synchronizing access to civisibilityTestsFuncs.
 	civisibilityTestsFuncsMutex sync.RWMutex
@@ -164,7 +165,7 @@ func setCiVisibilityTest(t *testing.T, ciTest integrations.DdTest) {
 }
 
 // hasCiVisibilityTestFunc gets if a func(*testing.T) is being instrumented.
-func hasCiVisibilityTestFunc(fn *func(*testing.T)) bool {
+func hasCiVisibilityTestFunc(fn *runtime.Func) bool {
 	civisibilityTestsFuncsMutex.RLock()
 	defer civisibilityTestsFuncsMutex.RUnlock()
 
@@ -176,7 +177,7 @@ func hasCiVisibilityTestFunc(fn *func(*testing.T)) bool {
 }
 
 // setCiVisibilityTestFunc tracks a func(*testing.T) as instrumented benchmark.
-func setCiVisibilityTestFunc(fn *func(*testing.T)) {
+func setCiVisibilityTestFunc(fn *runtime.Func) {
 	civisibilityTestsFuncsMutex.RLock()
 	defer civisibilityTestsFuncsMutex.RUnlock()
 	civisibilityTestsFuncs[fn] = struct{}{}

From 232c421b02ac5845a7605d0882168a05df19924d Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Mon, 9 Sep 2024 12:06:41 +0200
Subject: [PATCH 03/10] internal/civisibility/integrations/gotesting: refactor
 and more fixes for orchestrion autoinstrumentation

---
 .../integrations/gotesting/instrumentation.go | 123 ++++++++++++------
 .../integrations/gotesting/testing.go         |  13 +-
 .../integrations/gotesting/testingB.go        |  33 +----
 .../integrations/gotesting/testingT.go        |  60 +--------
 4 files changed, 102 insertions(+), 127 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index 6457a62a38..fc47407880 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -11,6 +11,7 @@ import (
 	"reflect"
 	"runtime"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -23,6 +24,64 @@ import (
 // The following functions are being used by the gotesting package for manual instrumentation and the orchestrion
 // automatic instrumentation
 
+type instrumentationMetadata struct {
+	IsInternal       bool
+	OriginalTest     *func(*testing.T)
+	InstrumentedTest *func(*testing.T)
+}
+
+var (
+	// instrumentationMap holds a map of *runtime.Func for tracking instrumented functions
+	instrumentationMap = map[*runtime.Func]*instrumentationMetadata{}
+
+	// instrumentationMapMutex is a read-write mutex for synchronizing access to instrumentationMap.
+	instrumentationMapMutex sync.RWMutex
+
+	// ciVisibilityTests holds a map of *testing.T or *testing.B to civisibility.DdTest for tracking tests.
+	ciVisibilityTests = map[testing.TB]integrations.DdTest{}
+
+	// ciVisibilityTestsMutex is a read-write mutex for synchronizing access to ciVisibilityTests.
+	ciVisibilityTestsMutex sync.RWMutex
+)
+
+// getInstrumentationMetadata gets the stored instrumentation metadata for a given *runtime.Func.
+func getInstrumentationMetadata(fn *runtime.Func) *instrumentationMetadata {
+	instrumentationMapMutex.RLock()
+	defer instrumentationMapMutex.RUnlock()
+
+	if v, ok := instrumentationMap[fn]; ok {
+		return v
+	}
+
+	return nil
+}
+
+// setInstrumentationMetadata stores an instrumentation metadata for a given *runtime.Func.
+func setInstrumentationMetadata(fn *runtime.Func, metadata *instrumentationMetadata) {
+	instrumentationMapMutex.RLock()
+	defer instrumentationMapMutex.RUnlock()
+	instrumentationMap[fn] = metadata
+}
+
+// getCiVisibilityTest retrieves the CI visibility test associated with a given *testing.T or *testing.B
+func getCiVisibilityTest(tb testing.TB) integrations.DdTest {
+	ciVisibilityTestsMutex.RLock()
+	defer ciVisibilityTestsMutex.RUnlock()
+
+	if v, ok := ciVisibilityTests[tb]; ok {
+		return v
+	}
+
+	return nil
+}
+
+// setCiVisibilityTest associates a CI visibility test with a given *testing.T or *testing.B
+func setCiVisibilityTest(tb testing.TB, ciTest integrations.DdTest) {
+	ciVisibilityTestsMutex.Lock()
+	defer ciVisibilityTestsMutex.Unlock()
+	ciVisibilityTests[tb] = ciTest
+}
+
 // instrumentTestingM helper function to instrument internalTests and internalBenchmarks in a `*testing.M` instance.
 func instrumentTestingM(m *testing.M) func(exitCode int) {
 	// Initialize CI Visibility
@@ -62,8 +121,14 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 	originalFunc := runtime.FuncForPC(fReflect.Pointer())
 
 	// Avoid instrumenting twice
-	if hasCiVisibilityTestFunc(originalFunc) {
-		return f
+	metadata := getInstrumentationMetadata(originalFunc)
+	if metadata != nil {
+		// If is an internal test, we don't instrument because f is already the instrumented func by executeInternalTest
+		if metadata.IsInternal {
+			return f
+		}
+
+		return *metadata.InstrumentedTest
 	}
 
 	instrumentedFunc := func(t *testing.T) {
@@ -117,30 +182,36 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 		f(t)
 	}
 
-	setCiVisibilityTestFunc(originalFunc)
-	setCiVisibilityTestFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()))
+	metadata = &instrumentationMetadata{
+		IsInternal:       false,
+		OriginalTest:     &f,
+		InstrumentedTest: &instrumentedFunc,
+	}
+
+	setInstrumentationMetadata(originalFunc, metadata)
+	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()), metadata)
 	return instrumentedFunc
 }
 
-// instrumentTestingTSetErrorInfo helper function to set an error in the `testing.T` CI Visibility span
-func instrumentTestingTSetErrorInfo(t *testing.T, errType string, errMessage string, skip int) {
-	ciTest := getCiVisibilityTest(t)
+// instrumentSetErrorInfo helper function to set an error in the `testing.T or testing.B` CI Visibility span
+func instrumentSetErrorInfo(tb testing.TB, errType string, errMessage string, skip int) {
+	ciTest := getCiVisibilityTest(tb)
 	if ciTest != nil {
 		ciTest.SetErrorInfo(errType, errMessage, utils.GetStacktrace(2+skip))
 	}
 }
 
-// instrumentTestingTCloseAndSkip helper function to close and skip with a reason a `testing.T` CI Visibility span
-func instrumentTestingTCloseAndSkip(t *testing.T, skipReason string) {
-	ciTest := getCiVisibilityTest(t)
+// instrumentCloseAndSkip helper function to close and skip with a reason a `testing.T or testing.B` CI Visibility span
+func instrumentCloseAndSkip(tb testing.TB, skipReason string) {
+	ciTest := getCiVisibilityTest(tb)
 	if ciTest != nil {
 		ciTest.CloseWithFinishTimeAndSkipReason(integrations.ResultStatusSkip, time.Now(), skipReason)
 	}
 }
 
-// instrumentTestingTSkipNow helper function to close and skip a `testing.T` CI Visibility span
-func instrumentTestingTSkipNow(t *testing.T) {
-	ciTest := getCiVisibilityTest(t)
+// instrumentSkipNow helper function to close and skip a `testing.T or testing.B` CI Visibility span
+func instrumentSkipNow(tb testing.TB) {
+	ciTest := getCiVisibilityTest(tb)
 	if ciTest != nil {
 		ciTest.Close(integrations.ResultStatusSkip)
 	}
@@ -219,7 +290,7 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 			// Replace this function with the original one (executed only once - the first iteration[b.run1]).
 			*iPfOfB.benchFunc = f
 			// Set b to the CI visibility test.
-			setCiVisibilityBenchmark(b, test)
+			setCiVisibilityTest(b, test)
 
 			// Enable the timer again.
 			b.ResetTimer()
@@ -290,27 +361,3 @@ func instrumentTestingBFunc(pb *testing.B, name string, f func(*testing.B)) (str
 	setCiVisibilityBenchmarkFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()))
 	return subBenchmarkAutoName, instrumentedFunc
 }
-
-// instrumentTestingBSetErrorInfo helper function to set an error in the `testing.B` CI Visibility span
-func instrumentTestingBSetErrorInfo(b *testing.B, errType string, errMessage string, skip int) {
-	ciTest := getCiVisibilityBenchmark(b)
-	if ciTest != nil {
-		ciTest.SetErrorInfo(errType, errMessage, utils.GetStacktrace(2+skip))
-	}
-}
-
-// instrumentTestingBCloseAndSkip helper function to close and skip with a reason a `testing.B` CI Visibility span
-func instrumentTestingBCloseAndSkip(b *testing.B, skipReason string) {
-	ciTest := getCiVisibilityBenchmark(b)
-	if ciTest != nil {
-		ciTest.CloseWithFinishTimeAndSkipReason(integrations.ResultStatusSkip, time.Now(), skipReason)
-	}
-}
-
-// instrumentTestingBSkipNow helper function to close and skip a `testing.B` CI Visibility span
-func instrumentTestingBSkipNow(b *testing.B) {
-	ciTest := getCiVisibilityBenchmark(b)
-	if ciTest != nil {
-		ciTest.Close(integrations.ResultStatusSkip)
-	}
-}
diff --git a/internal/civisibility/integrations/gotesting/testing.go b/internal/civisibility/integrations/gotesting/testing.go
index f7f6558b50..1722248dd9 100644
--- a/internal/civisibility/integrations/gotesting/testing.go
+++ b/internal/civisibility/integrations/gotesting/testing.go
@@ -165,8 +165,15 @@ func (ddm *M) executeInternalTest(testInfo *testingTInfo) func(*testing.T) {
 		// Execute the original test function.
 		testInfo.originalFunc(t)
 	}
-	setCiVisibilityTestFunc(originalFunc)
-	setCiVisibilityTestFunc(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()))
+
+	metadata := &instrumentationMetadata{
+		IsInternal:       true,
+		OriginalTest:     &testInfo.originalFunc,
+		InstrumentedTest: &instrumentedFunc,
+	}
+
+	setInstrumentationMetadata(originalFunc, metadata)
+	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()), metadata)
 	return instrumentedFunc
 }
 
@@ -256,7 +263,7 @@ func (ddm *M) executeInternalBenchmark(benchmarkInfo *testingBInfo) func(*testin
 			// Replace the benchmark function with the original one (this must be executed only once - the first iteration[b.run1]).
 			*iPfOfB.benchFunc = benchmarkInfo.originalFunc
 			// Set the CI visibility benchmark.
-			setCiVisibilityBenchmark(b, test)
+			setCiVisibilityTest(b, test)
 
 			// Restart the timer and execute the original benchmark function.
 			b.ResetTimer()
diff --git a/internal/civisibility/integrations/gotesting/testingB.go b/internal/civisibility/integrations/gotesting/testingB.go
index 9a602f6f24..c7fbcc5247 100644
--- a/internal/civisibility/integrations/gotesting/testingB.go
+++ b/internal/civisibility/integrations/gotesting/testingB.go
@@ -18,12 +18,6 @@ import (
 )
 
 var (
-	// ciVisibilityBenchmarks holds a map of *testing.B to civisibility.DdTest for tracking benchmarks.
-	ciVisibilityBenchmarks = map[*testing.B]integrations.DdTest{}
-
-	// ciVisibilityBenchmarksMutex is a read-write mutex for synchronizing access to ciVisibilityBenchmarks.
-	ciVisibilityBenchmarksMutex sync.RWMutex
-
 	// subBenchmarkAutoName is a placeholder name for CI Visibility sub-benchmarks.
 	subBenchmarkAutoName = "[DD:TestVisibility]"
 
@@ -60,7 +54,7 @@ func (ddb *B) Run(name string, f func(*testing.B)) bool {
 // integration tests.
 func (ddb *B) Context() context.Context {
 	b := (*testing.B)(ddb)
-	ciTest := getCiVisibilityBenchmark(b)
+	ciTest := getCiVisibilityTest(b)
 	if ciTest != nil {
 		return ciTest.Context()
 	}
@@ -114,7 +108,7 @@ func (ddb *B) Skipf(format string, args ...any) {
 // during the test. Calling SkipNow does not stop those other goroutines.
 func (ddb *B) SkipNow() {
 	b := (*testing.B)(ddb)
-	instrumentTestingBSkipNow(b)
+	instrumentSkipNow(b)
 	b.SkipNow()
 }
 
@@ -180,35 +174,16 @@ func (ddb *B) SetParallelism(p int) { (*testing.B)(ddb).SetParallelism(p) }
 
 func (ddb *B) getBWithError(errType string, errMessage string) *testing.B {
 	b := (*testing.B)(ddb)
-	instrumentTestingBSetErrorInfo(b, errType, errMessage, 1)
+	instrumentSetErrorInfo(b, errType, errMessage, 1)
 	return b
 }
 
 func (ddb *B) getBWithSkip(skipReason string) *testing.B {
 	b := (*testing.B)(ddb)
-	instrumentTestingBCloseAndSkip(b, skipReason)
+	instrumentCloseAndSkip(b, skipReason)
 	return b
 }
 
-// getCiVisibilityBenchmark retrieves the CI visibility benchmark associated with a given *testing.B.
-func getCiVisibilityBenchmark(b *testing.B) integrations.DdTest {
-	ciVisibilityBenchmarksMutex.RLock()
-	defer ciVisibilityBenchmarksMutex.RUnlock()
-
-	if v, ok := ciVisibilityBenchmarks[b]; ok {
-		return v
-	}
-
-	return nil
-}
-
-// setCiVisibilityBenchmark associates a CI visibility benchmark with a given *testing.B.
-func setCiVisibilityBenchmark(b *testing.B, ciTest integrations.DdTest) {
-	ciVisibilityBenchmarksMutex.Lock()
-	defer ciVisibilityBenchmarksMutex.Unlock()
-	ciVisibilityBenchmarks[b] = ciTest
-}
-
 // hasCiVisibilityBenchmarkFunc gets if a func(*testing.B) is being instrumented.
 func hasCiVisibilityBenchmarkFunc(fn *runtime.Func) bool {
 	civisibilityBenchmarksFuncsMutex.RLock()
diff --git a/internal/civisibility/integrations/gotesting/testingT.go b/internal/civisibility/integrations/gotesting/testingT.go
index 88ef23cd58..c53fd273eb 100644
--- a/internal/civisibility/integrations/gotesting/testingT.go
+++ b/internal/civisibility/integrations/gotesting/testingT.go
@@ -8,28 +8,12 @@ package gotesting
 import (
 	"context"
 	"fmt"
-	"runtime"
-	"sync"
 	"testing"
 	"time"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/civisibility/integrations"
 )
 
-var (
-	// ciVisibilityTests holds a map of *testing.T to civisibility.DdTest for tracking tests.
-	ciVisibilityTests = map[*testing.T]integrations.DdTest{}
-
-	// ciVisibilityTestsMutex is a read-write mutex for synchronizing access to ciVisibilityTests.
-	ciVisibilityTestsMutex sync.RWMutex
-
-	// civisibilityTestsFuncs holds a map of *func(*testing.T) for tracking instrumented functions
-	civisibilityTestsFuncs = map[*runtime.Func]struct{}{}
-
-	// civisibilityTestsFuncsMutex is a read-write mutex for synchronizing access to civisibilityTestsFuncs.
-	civisibilityTestsFuncsMutex sync.RWMutex
-)
-
 // T is a type alias for testing.T to provide additional methods for CI visibility.
 type T testing.T
 
@@ -110,7 +94,7 @@ func (ddt *T) Skipf(format string, args ...any) {
 // during the test. Calling SkipNow does not stop those other goroutines.
 func (ddt *T) SkipNow() {
 	t := (*testing.T)(ddt)
-	instrumentTestingTSkipNow(t)
+	instrumentSkipNow(t)
 	t.SkipNow()
 }
 
@@ -135,50 +119,12 @@ func (ddt *T) Setenv(key, value string) { (*testing.T)(ddt).Setenv(key, value) }
 
 func (ddt *T) getTWithError(errType string, errMessage string) *testing.T {
 	t := (*testing.T)(ddt)
-	instrumentTestingTSetErrorInfo(t, errType, errMessage, 1)
+	instrumentSetErrorInfo(t, errType, errMessage, 1)
 	return t
 }
 
 func (ddt *T) getTWithSkip(skipReason string) *testing.T {
 	t := (*testing.T)(ddt)
-	instrumentTestingTCloseAndSkip(t, skipReason)
+	instrumentCloseAndSkip(t, skipReason)
 	return t
 }
-
-// getCiVisibilityTest retrieves the CI visibility test associated with a given *testing.T.
-func getCiVisibilityTest(t *testing.T) integrations.DdTest {
-	ciVisibilityTestsMutex.RLock()
-	defer ciVisibilityTestsMutex.RUnlock()
-
-	if v, ok := ciVisibilityTests[t]; ok {
-		return v
-	}
-
-	return nil
-}
-
-// setCiVisibilityTest associates a CI visibility test with a given *testing.T.
-func setCiVisibilityTest(t *testing.T, ciTest integrations.DdTest) {
-	ciVisibilityTestsMutex.Lock()
-	defer ciVisibilityTestsMutex.Unlock()
-	ciVisibilityTests[t] = ciTest
-}
-
-// hasCiVisibilityTestFunc gets if a func(*testing.T) is being instrumented.
-func hasCiVisibilityTestFunc(fn *runtime.Func) bool {
-	civisibilityTestsFuncsMutex.RLock()
-	defer civisibilityTestsFuncsMutex.RUnlock()
-
-	if _, ok := civisibilityTestsFuncs[fn]; ok {
-		return true
-	}
-
-	return false
-}
-
-// setCiVisibilityTestFunc tracks a func(*testing.T) as instrumented benchmark.
-func setCiVisibilityTestFunc(fn *runtime.Func) {
-	civisibilityTestsFuncsMutex.RLock()
-	defer civisibilityTestsFuncsMutex.RUnlock()
-	civisibilityTestsFuncs[fn] = struct{}{}
-}

From 919400b79ca527ee19f0f763553c6331e2a3ccca Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Mon, 9 Sep 2024 15:53:58 +0200
Subject: [PATCH 04/10] ddtrace/tracer: adding log.Debug to ci visibility
 types.

---
 ddtrace/tracer/civisibility_payload.go   |  3 +++
 ddtrace/tracer/civisibility_transport.go |  3 +++
 ddtrace/tracer/civisibility_writer.go    | 11 ++++++-----
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/ddtrace/tracer/civisibility_payload.go b/ddtrace/tracer/civisibility_payload.go
index df8ffc04cc..34685e8bbe 100644
--- a/ddtrace/tracer/civisibility_payload.go
+++ b/ddtrace/tracer/civisibility_payload.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/tinylib/msgp/msgp"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/version"
 )
 
@@ -46,6 +47,7 @@ func (p *ciVisibilityPayload) push(event *ciVisibilityEvent) error {
 //
 //	A pointer to a newly initialized civisibilitypayload instance.
 func newCiVisibilityPayload() *ciVisibilityPayload {
+	log.Debug("ciVisibilityPayload: creating payload instance")
 	return &ciVisibilityPayload{newPayload()}
 }
 
@@ -61,6 +63,7 @@ func newCiVisibilityPayload() *ciVisibilityPayload {
 //	A pointer to a bytes.Buffer containing the encoded CI Visibility payload.
 //	An error if reading from the buffer or encoding the payload fails.
 func (p *ciVisibilityPayload) getBuffer(config *config) (*bytes.Buffer, error) {
+	log.Debug("ciVisibilityPayload: .getBuffer (count: %v)", p.itemCount())
 
 	/*
 			The Payload format in the CI Visibility protocol is like this:
diff --git a/ddtrace/tracer/civisibility_transport.go b/ddtrace/tracer/civisibility_transport.go
index db64b5d73d..0731332022 100644
--- a/ddtrace/tracer/civisibility_transport.go
+++ b/ddtrace/tracer/civisibility_transport.go
@@ -105,6 +105,8 @@ func newCiVisibilityTransport(config *config) *ciVisibilityTransport {
 		testCycleURL = fmt.Sprintf("%s/%s/%s", config.agentURL.String(), EvpProxyPath, TestCyclePath)
 	}
 
+	log.Debug("ciVisibilityTransport: creating transport instance [agentless: %v, testcycleurl: %v]", agentlessEnabled, testCycleURL)
+
 	return &ciVisibilityTransport{
 		config:           config,
 		testCycleURLPath: testCycleURL,
@@ -157,6 +159,7 @@ func (t *ciVisibilityTransport) send(p *payload) (body io.ReadCloser, err error)
 		req.Header.Set("Content-Encoding", "gzip")
 	}
 
+	log.Debug("ciVisibilityTransport: sending transport request: %v bytes", buffer.Len())
 	response, err := t.config.httpClient.Do(req)
 	if err != nil {
 		return nil, err
diff --git a/ddtrace/tracer/civisibility_writer.go b/ddtrace/tracer/civisibility_writer.go
index 1582b200a8..969b5edea6 100644
--- a/ddtrace/tracer/civisibility_writer.go
+++ b/ddtrace/tracer/civisibility_writer.go
@@ -45,6 +45,7 @@ type ciVisibilityTraceWriter struct {
 //
 //	A pointer to an initialized ciVisibilityTraceWriter.
 func newCiVisibilityTraceWriter(c *config) *ciVisibilityTraceWriter {
+	log.Debug("ciVisibilityTraceWriter: creating trace writer instance")
 	return &ciVisibilityTraceWriter{
 		config:  c,
 		payload: newCiVisibilityPayload(),
@@ -62,7 +63,7 @@ func (w *ciVisibilityTraceWriter) add(trace []*span) {
 	for _, s := range trace {
 		cvEvent := getCiVisibilityEvent(s)
 		if err := w.payload.push(cvEvent); err != nil {
-			log.Error("Error encoding msgpack: %v", err)
+			log.Error("ciVisibilityTraceWriter: Error encoding msgpack: %v", err)
 		}
 		if w.payload.size() > agentlessPayloadSizeLimit {
 			w.flush()
@@ -104,16 +105,16 @@ func (w *ciVisibilityTraceWriter) flush() {
 		var err error
 		for attempt := 0; attempt <= w.config.sendRetries; attempt++ {
 			size, count = p.size(), p.itemCount()
-			log.Debug("Sending payload: size: %d events: %d\n", size, count)
+			log.Debug("ciVisibilityTraceWriter: sending payload: size: %d events: %d\n", size, count)
 			_, err = w.config.transport.send(p.payload)
 			if err == nil {
-				log.Debug("sent events after %d attempts", attempt+1)
+				log.Debug("ciVisibilityTraceWriter: sent events after %d attempts", attempt+1)
 				return
 			}
-			log.Error("failure sending events (attempt %d), will retry: %v", attempt+1, err)
+			log.Error("ciVisibilityTraceWriter: failure sending events (attempt %d), will retry: %v", attempt+1, err)
 			p.reset()
 			time.Sleep(time.Millisecond)
 		}
-		log.Error("lost %d events: %v", count, err)
+		log.Error("ciVisibilityTraceWriter: lost %d events: %v", count, err)
 	}(oldp)
 }

From a8d9bd8d118f15e9f4eb86261fbc8264e3dbe933 Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Tue, 10 Sep 2024 12:01:24 +0200
Subject: [PATCH 05/10] ddtrace/tracer: fix json schema error when ParentId is
 not 0 for suites, modules and sessions.

---
 ddtrace/tracer/civisibility_tslv.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/ddtrace/tracer/civisibility_tslv.go b/ddtrace/tracer/civisibility_tslv.go
index 377f6d5656..ef6614d48f 100644
--- a/ddtrace/tracer/civisibility_tslv.go
+++ b/ddtrace/tracer/civisibility_tslv.go
@@ -273,6 +273,7 @@ func getCiVisibilityEvent(span *span) *ciVisibilityEvent {
 //	A pointer to the created ciVisibilityEvent.
 func createTestEventFromSpan(span *span) *ciVisibilityEvent {
 	tSpan := createTslvSpan(span)
+	tSpan.ParentID = 0
 	tSpan.SessionID = getAndRemoveMetaToUInt64(span, constants.TestSessionIDTag)
 	tSpan.ModuleID = getAndRemoveMetaToUInt64(span, constants.TestModuleIDTag)
 	tSpan.SuiteID = getAndRemoveMetaToUInt64(span, constants.TestSuiteIDTag)
@@ -298,6 +299,7 @@ func createTestEventFromSpan(span *span) *ciVisibilityEvent {
 //	A pointer to the created ciVisibilityEvent.
 func createTestSuiteEventFromSpan(span *span) *ciVisibilityEvent {
 	tSpan := createTslvSpan(span)
+	tSpan.ParentID = 0
 	tSpan.SessionID = getAndRemoveMetaToUInt64(span, constants.TestSessionIDTag)
 	tSpan.ModuleID = getAndRemoveMetaToUInt64(span, constants.TestModuleIDTag)
 	tSpan.SuiteID = getAndRemoveMetaToUInt64(span, constants.TestSuiteIDTag)
@@ -320,6 +322,7 @@ func createTestSuiteEventFromSpan(span *span) *ciVisibilityEvent {
 //	A pointer to the created ciVisibilityEvent.
 func createTestModuleEventFromSpan(span *span) *ciVisibilityEvent {
 	tSpan := createTslvSpan(span)
+	tSpan.ParentID = 0
 	tSpan.SessionID = getAndRemoveMetaToUInt64(span, constants.TestSessionIDTag)
 	tSpan.ModuleID = getAndRemoveMetaToUInt64(span, constants.TestModuleIDTag)
 	return &ciVisibilityEvent{
@@ -341,6 +344,7 @@ func createTestModuleEventFromSpan(span *span) *ciVisibilityEvent {
 //	A pointer to the created ciVisibilityEvent.
 func createTestSessionEventFromSpan(span *span) *ciVisibilityEvent {
 	tSpan := createTslvSpan(span)
+	tSpan.ParentID = 0
 	tSpan.SessionID = getAndRemoveMetaToUInt64(span, constants.TestSessionIDTag)
 	return &ciVisibilityEvent{
 		span:    span,

From 90698088f35b331428de9038c5321aacbc202806 Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Tue, 10 Sep 2024 13:31:28 +0200
Subject: [PATCH 06/10] internal/civisibility/integrations/gotesting: fix
 functions to get and store the civisibility test from `testing.T` and/or
 `testing.common` both structs implements `testing.TB` we want to get always
 the same test no matter the struct used.

---
 .../civisibility/integrations/gotesting/instrumentation.go | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index fc47407880..5089c94a51 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -15,6 +15,7 @@ import (
 	"sync/atomic"
 	"testing"
 	"time"
+	"unsafe"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/civisibility/integrations"
@@ -38,7 +39,7 @@ var (
 	instrumentationMapMutex sync.RWMutex
 
 	// ciVisibilityTests holds a map of *testing.T or *testing.B to civisibility.DdTest for tracking tests.
-	ciVisibilityTests = map[testing.TB]integrations.DdTest{}
+	ciVisibilityTests = map[unsafe.Pointer]integrations.DdTest{}
 
 	// ciVisibilityTestsMutex is a read-write mutex for synchronizing access to ciVisibilityTests.
 	ciVisibilityTestsMutex sync.RWMutex
@@ -68,7 +69,7 @@ func getCiVisibilityTest(tb testing.TB) integrations.DdTest {
 	ciVisibilityTestsMutex.RLock()
 	defer ciVisibilityTestsMutex.RUnlock()
 
-	if v, ok := ciVisibilityTests[tb]; ok {
+	if v, ok := ciVisibilityTests[reflect.ValueOf(tb).UnsafePointer()]; ok {
 		return v
 	}
 
@@ -79,7 +80,7 @@ func getCiVisibilityTest(tb testing.TB) integrations.DdTest {
 func setCiVisibilityTest(tb testing.TB, ciTest integrations.DdTest) {
 	ciVisibilityTestsMutex.Lock()
 	defer ciVisibilityTestsMutex.Unlock()
-	ciVisibilityTests[tb] = ciTest
+	ciVisibilityTests[reflect.ValueOf(tb).UnsafePointer()] = ciTest
 }
 
 // instrumentTestingM helper function to instrument internalTests and internalBenchmarks in a `*testing.M` instance.

From 34fb70feddd32790ec06dbe4f2deb0b2a8906a2a Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Tue, 10 Sep 2024 16:04:12 +0200
Subject: [PATCH 07/10] internal/civisibility/integrations/gotesting: add a new
 type to store if a ddtest has been failed or skipped so we can avoid
 rewriting existing useful messages by a generic one. The case of
 Error()/Errorf() calling Fail()/Failf() internally

---
 .../integrations/gotesting/instrumentation.go | 45 ++++++++++++-------
 .../integrations/gotesting/testingB.go        |  6 +--
 .../integrations/gotesting/testingT.go        |  6 +--
 3 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index 5089c94a51..8836df0d7e 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -25,11 +25,19 @@ import (
 // The following functions are being used by the gotesting package for manual instrumentation and the orchestrion
 // automatic instrumentation
 
-type instrumentationMetadata struct {
-	IsInternal       bool
-	OriginalTest     *func(*testing.T)
-	InstrumentedTest *func(*testing.T)
-}
+type (
+	instrumentationMetadata struct {
+		IsInternal       bool
+		OriginalTest     *func(*testing.T)
+		InstrumentedTest *func(*testing.T)
+	}
+
+	ddTestItem struct {
+		test    integrations.DdTest
+		error   bool
+		skipped bool
+	}
+)
 
 var (
 	// instrumentationMap holds a map of *runtime.Func for tracking instrumented functions
@@ -39,7 +47,7 @@ var (
 	instrumentationMapMutex sync.RWMutex
 
 	// ciVisibilityTests holds a map of *testing.T or *testing.B to civisibility.DdTest for tracking tests.
-	ciVisibilityTests = map[unsafe.Pointer]integrations.DdTest{}
+	ciVisibilityTests = map[unsafe.Pointer]*ddTestItem{}
 
 	// ciVisibilityTestsMutex is a read-write mutex for synchronizing access to ciVisibilityTests.
 	ciVisibilityTestsMutex sync.RWMutex
@@ -65,7 +73,7 @@ func setInstrumentationMetadata(fn *runtime.Func, metadata *instrumentationMetad
 }
 
 // getCiVisibilityTest retrieves the CI visibility test associated with a given *testing.T or *testing.B
-func getCiVisibilityTest(tb testing.TB) integrations.DdTest {
+func getCiVisibilityTest(tb testing.TB) *ddTestItem {
 	ciVisibilityTestsMutex.RLock()
 	defer ciVisibilityTestsMutex.RUnlock()
 
@@ -80,7 +88,7 @@ func getCiVisibilityTest(tb testing.TB) integrations.DdTest {
 func setCiVisibilityTest(tb testing.TB, ciTest integrations.DdTest) {
 	ciVisibilityTestsMutex.Lock()
 	defer ciVisibilityTestsMutex.Unlock()
-	ciVisibilityTests[reflect.ValueOf(tb).UnsafePointer()] = ciTest
+	ciVisibilityTests[reflect.ValueOf(tb).UnsafePointer()] = &ddTestItem{test: ciTest}
 }
 
 // instrumentTestingM helper function to instrument internalTests and internalBenchmarks in a `*testing.M` instance.
@@ -196,25 +204,28 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 
 // instrumentSetErrorInfo helper function to set an error in the `testing.T or testing.B` CI Visibility span
 func instrumentSetErrorInfo(tb testing.TB, errType string, errMessage string, skip int) {
-	ciTest := getCiVisibilityTest(tb)
-	if ciTest != nil {
-		ciTest.SetErrorInfo(errType, errMessage, utils.GetStacktrace(2+skip))
+	ciTestItem := getCiVisibilityTest(tb)
+	if ciTestItem != nil && !ciTestItem.error && ciTestItem.test != nil {
+		ciTestItem.error = true
+		ciTestItem.test.SetErrorInfo(errType, errMessage, utils.GetStacktrace(2+skip))
 	}
 }
 
 // instrumentCloseAndSkip helper function to close and skip with a reason a `testing.T or testing.B` CI Visibility span
 func instrumentCloseAndSkip(tb testing.TB, skipReason string) {
-	ciTest := getCiVisibilityTest(tb)
-	if ciTest != nil {
-		ciTest.CloseWithFinishTimeAndSkipReason(integrations.ResultStatusSkip, time.Now(), skipReason)
+	ciTestItem := getCiVisibilityTest(tb)
+	if ciTestItem != nil && !ciTestItem.skipped && ciTestItem.test != nil {
+		ciTestItem.skipped = true
+		ciTestItem.test.CloseWithFinishTimeAndSkipReason(integrations.ResultStatusSkip, time.Now(), skipReason)
 	}
 }
 
 // instrumentSkipNow helper function to close and skip a `testing.T or testing.B` CI Visibility span
 func instrumentSkipNow(tb testing.TB) {
-	ciTest := getCiVisibilityTest(tb)
-	if ciTest != nil {
-		ciTest.Close(integrations.ResultStatusSkip)
+	ciTestItem := getCiVisibilityTest(tb)
+	if ciTestItem != nil && !ciTestItem.skipped && ciTestItem.test != nil {
+		ciTestItem.skipped = true
+		ciTestItem.test.Close(integrations.ResultStatusSkip)
 	}
 }
 
diff --git a/internal/civisibility/integrations/gotesting/testingB.go b/internal/civisibility/integrations/gotesting/testingB.go
index c7fbcc5247..d7ecd7c01c 100644
--- a/internal/civisibility/integrations/gotesting/testingB.go
+++ b/internal/civisibility/integrations/gotesting/testingB.go
@@ -54,9 +54,9 @@ func (ddb *B) Run(name string, f func(*testing.B)) bool {
 // integration tests.
 func (ddb *B) Context() context.Context {
 	b := (*testing.B)(ddb)
-	ciTest := getCiVisibilityTest(b)
-	if ciTest != nil {
-		return ciTest.Context()
+	ciTestItem := getCiVisibilityTest(b)
+	if ciTestItem != nil && ciTestItem.test != nil {
+		return ciTestItem.test.Context()
 	}
 
 	return context.Background()
diff --git a/internal/civisibility/integrations/gotesting/testingT.go b/internal/civisibility/integrations/gotesting/testingT.go
index c53fd273eb..e05f60abc8 100644
--- a/internal/civisibility/integrations/gotesting/testingT.go
+++ b/internal/civisibility/integrations/gotesting/testingT.go
@@ -40,9 +40,9 @@ func (ddt *T) Run(name string, f func(*testing.T)) bool {
 // integration tests.
 func (ddt *T) Context() context.Context {
 	t := (*testing.T)(ddt)
-	ciTest := getCiVisibilityTest(t)
-	if ciTest != nil {
-		return ciTest.Context()
+	ciTestItem := getCiVisibilityTest(t)
+	if ciTestItem != nil && ciTestItem.test != nil {
+		return ciTestItem.test.Context()
 	}
 
 	return context.Background()

From 05323f4feba022ee43afdc48edc29357fb2a01e2 Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Wed, 11 Sep 2024 10:23:08 +0200
Subject: [PATCH 08/10] internal/civisibility/integrations/gotesting: change
 ddTestItem.error and ddTestItem.skipped bool type to an atomic.Int32 to avoid
 any race condition.

---
 .../integrations/gotesting/instrumentation.go       | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index 8836df0d7e..6077c34e10 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -34,8 +34,8 @@ type (
 
 	ddTestItem struct {
 		test    integrations.DdTest
-		error   bool
-		skipped bool
+		error   atomic.Int32
+		skipped atomic.Int32
 	}
 )
 
@@ -205,8 +205,7 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 // instrumentSetErrorInfo helper function to set an error in the `testing.T or testing.B` CI Visibility span
 func instrumentSetErrorInfo(tb testing.TB, errType string, errMessage string, skip int) {
 	ciTestItem := getCiVisibilityTest(tb)
-	if ciTestItem != nil && !ciTestItem.error && ciTestItem.test != nil {
-		ciTestItem.error = true
+	if ciTestItem != nil && ciTestItem.error.CompareAndSwap(0, 1) && ciTestItem.test != nil {
 		ciTestItem.test.SetErrorInfo(errType, errMessage, utils.GetStacktrace(2+skip))
 	}
 }
@@ -214,8 +213,7 @@ func instrumentSetErrorInfo(tb testing.TB, errType string, errMessage string, sk
 // instrumentCloseAndSkip helper function to close and skip with a reason a `testing.T or testing.B` CI Visibility span
 func instrumentCloseAndSkip(tb testing.TB, skipReason string) {
 	ciTestItem := getCiVisibilityTest(tb)
-	if ciTestItem != nil && !ciTestItem.skipped && ciTestItem.test != nil {
-		ciTestItem.skipped = true
+	if ciTestItem != nil && ciTestItem.skipped.CompareAndSwap(0, 1) && ciTestItem.test != nil {
 		ciTestItem.test.CloseWithFinishTimeAndSkipReason(integrations.ResultStatusSkip, time.Now(), skipReason)
 	}
 }
@@ -223,8 +221,7 @@ func instrumentCloseAndSkip(tb testing.TB, skipReason string) {
 // instrumentSkipNow helper function to close and skip a `testing.T or testing.B` CI Visibility span
 func instrumentSkipNow(tb testing.TB) {
 	ciTestItem := getCiVisibilityTest(tb)
-	if ciTestItem != nil && !ciTestItem.skipped && ciTestItem.test != nil {
-		ciTestItem.skipped = true
+	if ciTestItem != nil && ciTestItem.skipped.CompareAndSwap(0, 1) && ciTestItem.test != nil {
 		ciTestItem.test.Close(integrations.ResultStatusSkip)
 	}
 }

From 2246373ac55db45e8580016d189e9ac2bd49b8f4 Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Wed, 11 Sep 2024 15:59:47 +0200
Subject: [PATCH 09/10] internal/civisibility/integrations/gotesting: fixes an
 issue where we were caching the instrumented function from an original one,
 caching also the closure data of the function, now we don't catch the
 instrumented function. Also, added some tests to handle this scenario.

---
 .../integrations/gotesting/instrumentation.go | 30 ++++---------------
 .../integrations/gotesting/testing.go         |  9 +-----
 .../integrations/gotesting/testing_test.go    | 21 ++++++++++---
 3 files changed, 24 insertions(+), 36 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index 6077c34e10..0098505219 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -27,9 +27,7 @@ import (
 
 type (
 	instrumentationMetadata struct {
-		IsInternal       bool
-		OriginalTest     *func(*testing.T)
-		InstrumentedTest *func(*testing.T)
+		IsInternal bool
 	}
 
 	ddTestItem struct {
@@ -57,11 +55,9 @@ var (
 func getInstrumentationMetadata(fn *runtime.Func) *instrumentationMetadata {
 	instrumentationMapMutex.RLock()
 	defer instrumentationMapMutex.RUnlock()
-
 	if v, ok := instrumentationMap[fn]; ok {
 		return v
 	}
-
 	return nil
 }
 
@@ -76,11 +72,9 @@ func setInstrumentationMetadata(fn *runtime.Func, metadata *instrumentationMetad
 func getCiVisibilityTest(tb testing.TB) *ddTestItem {
 	ciVisibilityTestsMutex.RLock()
 	defer ciVisibilityTestsMutex.RUnlock()
-
 	if v, ok := ciVisibilityTests[reflect.ValueOf(tb).UnsafePointer()]; ok {
 		return v
 	}
-
 	return nil
 }
 
@@ -131,16 +125,12 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 
 	// Avoid instrumenting twice
 	metadata := getInstrumentationMetadata(originalFunc)
-	if metadata != nil {
+	if metadata != nil && metadata.IsInternal {
 		// If is an internal test, we don't instrument because f is already the instrumented func by executeInternalTest
-		if metadata.IsInternal {
-			return f
-		}
-
-		return *metadata.InstrumentedTest
+		return f
 	}
 
-	instrumentedFunc := func(t *testing.T) {
+	instrumentedFn := func(t *testing.T) {
 		// Initialize module counters if not already present.
 		if _, ok := modulesCounters[moduleName]; !ok {
 			var v int32
@@ -190,16 +180,8 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 		// Execute the original test function.
 		f(t)
 	}
-
-	metadata = &instrumentationMetadata{
-		IsInternal:       false,
-		OriginalTest:     &f,
-		InstrumentedTest: &instrumentedFunc,
-	}
-
-	setInstrumentationMetadata(originalFunc, metadata)
-	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()), metadata)
-	return instrumentedFunc
+	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFn)).Pointer()), &instrumentationMetadata{IsInternal: true})
+	return instrumentedFn
 }
 
 // instrumentSetErrorInfo helper function to set an error in the `testing.T or testing.B` CI Visibility span
diff --git a/internal/civisibility/integrations/gotesting/testing.go b/internal/civisibility/integrations/gotesting/testing.go
index 1722248dd9..91cc246fc9 100644
--- a/internal/civisibility/integrations/gotesting/testing.go
+++ b/internal/civisibility/integrations/gotesting/testing.go
@@ -166,14 +166,7 @@ func (ddm *M) executeInternalTest(testInfo *testingTInfo) func(*testing.T) {
 		testInfo.originalFunc(t)
 	}
 
-	metadata := &instrumentationMetadata{
-		IsInternal:       true,
-		OriginalTest:     &testInfo.originalFunc,
-		InstrumentedTest: &instrumentedFunc,
-	}
-
-	setInstrumentationMetadata(originalFunc, metadata)
-	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()), metadata)
+	setInstrumentationMetadata(runtime.FuncForPC(reflect.Indirect(reflect.ValueOf(instrumentedFunc)).Pointer()), &instrumentationMetadata{IsInternal: true})
 	return instrumentedFunc
 }
 
diff --git a/internal/civisibility/integrations/gotesting/testing_test.go b/internal/civisibility/integrations/gotesting/testing_test.go
index e45e62d15d..4345e9e822 100644
--- a/internal/civisibility/integrations/gotesting/testing_test.go
+++ b/internal/civisibility/integrations/gotesting/testing_test.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"slices"
 	"strconv"
 	"testing"
 
@@ -35,7 +36,10 @@ func TestMain(m *testing.M) {
 	// or use a helper method gotesting.RunM(m)
 
 	// os.Exit((*M)(m).Run())
-	_ = RunM(m)
+	exit := RunM(m)
+	if exit != 0 {
+		os.Exit(exit)
+	}
 
 	finishedSpans := mTracer.FinishedSpans()
 	// 1 session span
@@ -107,19 +111,28 @@ func Test_Foo(gt *testing.T) {
 	assertTest(gt)
 	t := (*T)(gt)
 	var tests = []struct {
+		index byte
 		name  string
 		input string
 		want  string
 	}{
-		{"yellow should return color", "yellow", "color"},
-		{"banana should return fruit", "banana", "fruit"},
-		{"duck should return animal", "duck", "animal"},
+		{1, "yellow should return color", "yellow", "color"},
+		{2, "banana should return fruit", "banana", "fruit"},
+		{3, "duck should return animal", "duck", "animal"},
 	}
+	buf := []byte{}
 	for _, test := range tests {
+		test := test
 		t.Run(test.name, func(t *testing.T) {
 			t.Log(test.name)
+			buf = append(buf, test.index)
 		})
 	}
+
+	expected := []byte{1, 2, 3}
+	if !slices.Equal(buf, expected) {
+		t.Error("error in subtests closure")
+	}
 }
 
 // TestWithExternalCalls demonstrates testing with external HTTP calls.

From 50434cc5b96c9401c25551a46251702cd92d8266 Mon Sep 17 00:00:00 2001
From: Tony Redondo <tony.redondo@datadoghq.com>
Date: Thu, 12 Sep 2024 10:45:43 +0200
Subject: [PATCH 10/10] internal/civisibility/integrations/gotesting: fix
 comments.

---
 .../integrations/gotesting/instrumentation.go          | 10 +++++-----
 .../civisibility/integrations/gotesting/testingB.go    |  6 +++---
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/internal/civisibility/integrations/gotesting/instrumentation.go b/internal/civisibility/integrations/gotesting/instrumentation.go
index 0098505219..01abacc910 100644
--- a/internal/civisibility/integrations/gotesting/instrumentation.go
+++ b/internal/civisibility/integrations/gotesting/instrumentation.go
@@ -68,7 +68,7 @@ func setInstrumentationMetadata(fn *runtime.Func, metadata *instrumentationMetad
 	instrumentationMap[fn] = metadata
 }
 
-// getCiVisibilityTest retrieves the CI visibility test associated with a given *testing.T or *testing.B
+// getCiVisibilityTest retrieves the CI visibility test associated with a given *testing.T, *testing.B, *testing.common
 func getCiVisibilityTest(tb testing.TB) *ddTestItem {
 	ciVisibilityTestsMutex.RLock()
 	defer ciVisibilityTestsMutex.RUnlock()
@@ -78,7 +78,7 @@ func getCiVisibilityTest(tb testing.TB) *ddTestItem {
 	return nil
 }
 
-// setCiVisibilityTest associates a CI visibility test with a given *testing.T or *testing.B
+// setCiVisibilityTest associates a CI visibility test with a given *testing.T, *testing.B, *testing.common
 func setCiVisibilityTest(tb testing.TB, ciTest integrations.DdTest) {
 	ciVisibilityTestsMutex.Lock()
 	defer ciVisibilityTestsMutex.Unlock()
@@ -184,7 +184,7 @@ func instrumentTestingTFunc(f func(*testing.T)) func(*testing.T) {
 	return instrumentedFn
 }
 
-// instrumentSetErrorInfo helper function to set an error in the `testing.T or testing.B` CI Visibility span
+// instrumentSetErrorInfo helper function to set an error in the `*testing.T, *testing.B, *testing.common` CI Visibility span
 func instrumentSetErrorInfo(tb testing.TB, errType string, errMessage string, skip int) {
 	ciTestItem := getCiVisibilityTest(tb)
 	if ciTestItem != nil && ciTestItem.error.CompareAndSwap(0, 1) && ciTestItem.test != nil {
@@ -192,7 +192,7 @@ func instrumentSetErrorInfo(tb testing.TB, errType string, errMessage string, sk
 	}
 }
 
-// instrumentCloseAndSkip helper function to close and skip with a reason a `testing.T or testing.B` CI Visibility span
+// instrumentCloseAndSkip helper function to close and skip with a reason a `*testing.T, *testing.B, *testing.common` CI Visibility span
 func instrumentCloseAndSkip(tb testing.TB, skipReason string) {
 	ciTestItem := getCiVisibilityTest(tb)
 	if ciTestItem != nil && ciTestItem.skipped.CompareAndSwap(0, 1) && ciTestItem.test != nil {
@@ -200,7 +200,7 @@ func instrumentCloseAndSkip(tb testing.TB, skipReason string) {
 	}
 }
 
-// instrumentSkipNow helper function to close and skip a `testing.T or testing.B` CI Visibility span
+// instrumentSkipNow helper function to close and skip a `*testing.T, *testing.B, *testing.common` CI Visibility span
 func instrumentSkipNow(tb testing.TB) {
 	ciTestItem := getCiVisibilityTest(tb)
 	if ciTestItem != nil && ciTestItem.skipped.CompareAndSwap(0, 1) && ciTestItem.test != nil {
diff --git a/internal/civisibility/integrations/gotesting/testingB.go b/internal/civisibility/integrations/gotesting/testingB.go
index d7ecd7c01c..b37bef009d 100644
--- a/internal/civisibility/integrations/gotesting/testingB.go
+++ b/internal/civisibility/integrations/gotesting/testingB.go
@@ -24,7 +24,7 @@ var (
 	// subBenchmarkAutoNameRegex is a regex pattern to match the sub-benchmark auto name.
 	subBenchmarkAutoNameRegex = regexp.MustCompile(`(?si)\/\[DD:TestVisibility\].*`)
 
-	// civisibilityBenchmarksFuncs holds a map of *func(*testing.B) for tracking instrumented functions
+	// civisibilityBenchmarksFuncs holds a map of *runtime.Func for tracking instrumented functions
 	civisibilityBenchmarksFuncs = map[*runtime.Func]struct{}{}
 
 	// civisibilityBenchmarksFuncsMutex is a read-write mutex for synchronizing access to civisibilityBenchmarksFuncs.
@@ -184,7 +184,7 @@ func (ddb *B) getBWithSkip(skipReason string) *testing.B {
 	return b
 }
 
-// hasCiVisibilityBenchmarkFunc gets if a func(*testing.B) is being instrumented.
+// hasCiVisibilityBenchmarkFunc gets if a *runtime.Func is being instrumented.
 func hasCiVisibilityBenchmarkFunc(fn *runtime.Func) bool {
 	civisibilityBenchmarksFuncsMutex.RLock()
 	defer civisibilityBenchmarksFuncsMutex.RUnlock()
@@ -196,7 +196,7 @@ func hasCiVisibilityBenchmarkFunc(fn *runtime.Func) bool {
 	return false
 }
 
-// setCiVisibilityBenchmarkFunc tracks a func(*testing.B) as instrumented benchmark.
+// setCiVisibilityBenchmarkFunc tracks a *runtime.Func as instrumented benchmark.
 func setCiVisibilityBenchmarkFunc(fn *runtime.Func) {
 	civisibilityBenchmarksFuncsMutex.RLock()
 	defer civisibilityBenchmarksFuncsMutex.RUnlock()