diff --git a/profiler/profile.go b/profiler/profile.go index a8d04cfd7e..038d077dc8 100644 --- a/profiler/profile.go +++ b/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 diff --git a/profiler/profile_test.go b/profiler/profile_test.go index 7c14422c92..cd9d23f86f 100644 --- a/profiler/profile_test.go +++ b/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 diff --git a/profiler/profiler.go b/profiler/profiler.go index ee9088ff35..4e33a6d998 100644 --- a/profiler/profiler.go +++ b/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. diff --git a/profiler/profiler_test.go b/profiler/profiler_test.go index db02f1781b..48d919ab2c 100644 --- a/profiler/profiler_test.go +++ b/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) }