profiler: add hooks to override functions for testing (#1259)

The profiler uses global startCPUProfile/stopCPUProfile/lookupProfile
functions that are overridden by tests.

Any test that wants to override these functions has to reset them. This
results in duplicated code, and can cause difficult-to-understand issues
if a new test is added that overrides the values and fails to reset
them. For example, this has contributed to race conditions in our unit
tests that are unrelated to actual bugs in our code. See PR #1255 for
example. Having bugs in our tests makes it more difficult to diagnose
real bugs in our code.

Instead of having global values, add override hooks to the profiler
struct that can be set during tests and fall back to the default
behavior if no hooks are set. So long as each test uses its own profiler
struct, each test can set its own hooks without affecting any other
test.

profiler/profile.go

@@ -11,7 +11,6 @@ import (
 	"fmt"
 	"io"
 	"runtime"
-	"runtime/pprof"
 	"time"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/profiler/internal/pprofutils"
@@ -88,11 +87,11 @@ var profileTypes = map[ProfileType]profileType{
 				// rate itself.
 				runtime.SetCPUProfileRate(p.cfg.cpuProfileRate)
 			}
-			if err := startCPUProfile(&buf); err != nil {
+			if err := p.startCPUProfile(&buf); err != nil {
 				return nil, err
 			}
 			p.interruptibleSleep(p.cfg.cpuDuration)
-			stopCPUProfile()
+			p.stopCPUProfile()
 			return buf.Bytes(), nil
 		},
 	},
@@ -140,7 +139,7 @@ var profileTypes = map[ProfileType]profileType{
 				text  = &bytes.Buffer{}
 				pprof = &bytes.Buffer{}
 			)
-			if err := lookupProfile("goroutine", text, 2); err != nil {
+			if err := p.lookupProfile("goroutine", text, 2); err != nil {
 				return nil, err
 			}
 			err := goroutineDebug2ToPprof(text, pprof, now)
@@ -158,9 +157,9 @@ var profileTypes = map[ProfileType]profileType{
 	},
 }
 
-func collectGenericProfile(t profileType, _ *profiler) ([]byte, error) {
+func collectGenericProfile(t profileType, p *profiler) ([]byte, error) {
 	var buf bytes.Buffer
-	err := lookupProfile(t.Name, &buf, 0)
+	err := p.lookupProfile(t.Name, &buf, 0)
 	return buf.Bytes(), err
 }
 
@@ -293,23 +292,6 @@ func (p *profiler) deltaProfile(t profileType, curData []byte) (*profile, error)
 	}, nil
 }
 
-var (
-	// startCPUProfile starts the CPU profile; replaced in tests
-	startCPUProfile = pprof.StartCPUProfile
-	// stopCPUProfile stops the CPU profile; replaced in tests
-	stopCPUProfile = pprof.StopCPUProfile
-)
-
-// lookpupProfile looks up the profile with the given name and writes it to w. It returns
-// any errors encountered in the process. It is replaced in tests.
-var lookupProfile = func(name string, w io.Writer, debug int) error {
-	prof := pprof.Lookup(name)
-	if prof == nil {
-		return errors.New("profile not found")
-	}
-	return prof.WriteTo(w, debug)
-}
-
 func goroutineDebug2ToPprof(r io.Reader, w io.Writer, t time.Time) (err error) {
 	// gostackparse.Parse() has been extensively tested and should not crash
 	// under any circumstances, but we really want to avoid crashing a customers

profiler/profile_test.go

@@ -117,15 +117,14 @@ main;bar 0 0 8 16
 				// followed by prof2 when calling runProfile().
 				deltaProfiler := func(prof1, prof2 []byte, opts ...Option) (*profiler, func()) {
 					returnProfs := [][]byte{prof1, prof2}
-					old := lookupProfile
-					lookupProfile = func(_ string, w io.Writer, _ int) error {
+					p, err := unstartedProfiler(opts...)
+					p.testHooks.lookupProfile = func(_ string, w io.Writer, _ int) error {
 						_, err := w.Write(returnProfs[0])
 						returnProfs = returnProfs[1:]
 						return err
 					}
-					p, err := unstartedProfiler(opts...)
 					require.NoError(t, err)
-					return p, func() { lookupProfile = old }
+					return p, func() {}
 				}
 
 				t.Run(profType.String(), func(t *testing.T) {
@@ -173,15 +172,12 @@ main;bar 0 0 8 16
 	})
 
 	t.Run("cpu", func(t *testing.T) {
-		defer func(old func(_ io.Writer) error) { startCPUProfile = old }(startCPUProfile)
-		startCPUProfile = func(w io.Writer) error {
+		p, err := unstartedProfiler(CPUDuration(10 * time.Millisecond))
+		p.testHooks.startCPUProfile = func(w io.Writer) error {
 			_, err := w.Write([]byte("my-cpu-profile"))
 			return err
 		}
-		defer func(old func()) { stopCPUProfile = old }(stopCPUProfile)
-		stopCPUProfile = func() {}
-
-		p, err := unstartedProfiler(CPUDuration(10 * time.Millisecond))
+		p.testHooks.stopCPUProfile = func() {}
 		require.NoError(t, err)
 		start := time.Now()
 		profs, err := p.runProfile(CPUProfile)
@@ -193,13 +189,11 @@ main;bar 0 0 8 16
 	})
 
 	t.Run("goroutine", func(t *testing.T) {
-		defer func(old func(_ string, _ io.Writer, _ int) error) { lookupProfile = old }(lookupProfile)
-		lookupProfile = func(name string, w io.Writer, _ int) error {
+		p, err := unstartedProfiler()
+		p.testHooks.lookupProfile = func(name string, w io.Writer, _ int) error {
 			_, err := w.Write([]byte(name))
 			return err
 		}
-
-		p, err := unstartedProfiler()
 		require.NoError(t, err)
 		profs, err := p.runProfile(GoroutineProfile)
 		require.NoError(t, err)
@@ -230,13 +224,12 @@ main.main()
 ...additional frames elided...
 `
 
-		defer func(old func(_ string, _ io.Writer, _ int) error) { lookupProfile = old }(lookupProfile)
-		lookupProfile = func(_ string, w io.Writer, _ int) error {
+		p, err := unstartedProfiler()
+		p.testHooks.lookupProfile = func(_ string, w io.Writer, _ int) error {
 			_, err := w.Write([]byte(sample))
 			return err
 		}
 
-		p, err := unstartedProfiler()
 		require.NoError(t, err)
 		profs, err := p.runProfile(expGoroutineWaitProfile)
 		require.NoError(t, err)
@@ -315,13 +308,11 @@ main.main()
 		os.Setenv(envVar, strconv.Itoa(limit))
 		defer os.Setenv(envVar, oldVal)
 
-		defer func(old func(_ string, _ io.Writer, _ int) error) { lookupProfile = old }(lookupProfile)
-		lookupProfile = func(_ string, w io.Writer, _ int) error {
+		p, err := unstartedProfiler()
+		p.testHooks.lookupProfile = func(_ string, w io.Writer, _ int) error {
 			_, err := w.Write([]byte(""))
 			return err
 		}
-
-		p, err := unstartedProfiler()
 		require.NoError(t, err)
 		_, err = p.runProfile(expGoroutineWaitProfile)
 		var errRoutines, errLimit int

profiler/profiler.go

@@ -8,10 +8,12 @@ package profiler
 import (
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
 	"runtime"
+	"runtime/pprof"
 	"sync"
 	"time"
 
@@ -69,6 +71,42 @@ type profiler struct {
 	wg         sync.WaitGroup                    // wg waits for all goroutines to exit when stopping.
 	met        *metrics                          // metric collector state
 	prev       map[ProfileType]*pprofile.Profile // previous collection results for delta profiling
+
+	testHooks testHooks
+}
+
+// testHooks are functions that are replaced during testing which would normally
+// depend on accessing runtime state that is not needed/available for the test
+type testHooks struct {
+	startCPUProfile func(w io.Writer) error
+	stopCPUProfile  func()
+	lookupProfile   func(name string, w io.Writer, debug int) error
+}
+
+func (p *profiler) startCPUProfile(w io.Writer) error {
+	if p.testHooks.startCPUProfile != nil {
+		return p.testHooks.startCPUProfile(w)
+	}
+	return pprof.StartCPUProfile(w)
+}
+
+func (p *profiler) stopCPUProfile() {
+	if p.testHooks.startCPUProfile != nil {
+		p.testHooks.stopCPUProfile()
+		return
+	}
+	pprof.StopCPUProfile()
+}
+
+func (p *profiler) lookupProfile(name string, w io.Writer, debug int) error {
+	if p.testHooks.lookupProfile != nil {
+		return p.testHooks.lookupProfile(name, w, debug)
+	}
+	prof := pprof.Lookup(name)
+	if prof == nil {
+		return errors.New("profile not found")
+	}
+	return prof.WriteTo(w, debug)
 }
 
 // newProfiler creates a new, unstarted profiler.

profiler/profiler_test.go

@@ -232,15 +232,12 @@ func TestProfilerInternal(t *testing.T) {
 		)
 		require.NoError(t, err)
 		var startCPU, stopCPU, writeHeap uint64
-		defer func(old func(_ io.Writer) error) { startCPUProfile = old }(startCPUProfile)
-		startCPUProfile = func(_ io.Writer) error {
+		p.testHooks.startCPUProfile = func(_ io.Writer) error {
 			atomic.AddUint64(&startCPU, 1)
 			return nil
 		}
-		defer func(old func()) { stopCPUProfile = old }(stopCPUProfile)
-		stopCPUProfile = func() { atomic.AddUint64(&stopCPU, 1) }
-		defer func(old func(_ string, w io.Writer, _ int) error) { lookupProfile = old }(lookupProfile)
-		lookupProfile = func(name string, w io.Writer, _ int) error {
+		p.testHooks.stopCPUProfile = func() { atomic.AddUint64(&stopCPU, 1) }
+		p.testHooks.lookupProfile = func(name string, w io.Writer, _ int) error {
 			if name == "heap" {
 				atomic.AddUint64(&writeHeap, 1)
 			}