pomerium/internal/testenv/envutil/pprof.go

101 lines
2.8 KiB
Go

package envutil
import (
"flag"
"io"
"os"
"os/exec"
"path/filepath"
"runtime"
"runtime/pprof"
"sync"
"testing"
"github.com/stretchr/testify/require"
)
// PauseProfiling will suspend CPU and memory profiling, if started using the
// -cpuprofile and/or -memprofile test flags. The returned function will restart
// profiling when called. Existing CPU profile data is overwritten, but
// existing memory profile data is kept.
func PauseProfiling(t testing.TB) (resume func()) {
resumeFuncs := []func(){}
if filename := ProfileOutputPath("cpuprofile"); filename != "" {
if _, err := os.Stat(filename); err == nil {
pprof.StopCPUProfile()
t.Logf("pausing cpu profiling (%s)", filename)
resumeFuncs = append(resumeFuncs, func() {
t.Logf("resuming cpu profiling (%s)", filename)
f, err := os.Create(filename)
require.NoError(t, err)
require.NoError(t, pprof.StartCPUProfile(f))
})
}
}
if filename := ProfileOutputPath("memprofile"); filename != "" {
rate := runtime.MemProfileRate
runtime.MemProfileRate = 0
t.Log("pausing memory profiling")
resumeFuncs = append(resumeFuncs, func() {
t.Log("resuming memory profiling")
runtime.MemProfileRate = rate
})
}
return sync.OnceFunc(func() {
for _, f := range resumeFuncs {
f()
}
})
}
// Returns the file path set by the '-test.<name>profile' flag, or empty string
// if the flag was not set.
func ProfileOutputPath(name string) string {
outputdir := flag.Lookup("test.outputdir")
if f := flag.Lookup("test." + name); f != nil {
if filename := f.Value.String(); filename != "" {
if outputdir != nil {
filename = filepath.Join(outputdir.Value.String(), filename)
}
return filename
}
}
return ""
}
// Returns true if the envoy binary at the given path was compiled with
// gperftools profiler support.
func EnvoyProfilerAvailable(binary string) bool {
// There are a few symbols that will only show up if envoy is compiled with
// tcmalloc=gperftools. Specifically, symbols defined in these headers:
// https://github.com/gperftools/gperftools/tree/master/src/gperftools
// The symbols are not mangled, so pick one that is unlikely to be ambiguous
// or part of another function name.
err := exec.Command("/usr/bin/grep", "-q", "ProfilingIsEnabledForAllThreads", binary).Run()
return err == nil
}
func CollectEnvoyHeapProfiles(base string) error {
combined, err := os.OpenFile(base, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
if err != nil {
return err
}
defer combined.Close()
parts, _ := filepath.Glob(base + ".*.heap")
for _, part := range parts {
pf, err := os.Open(part)
if err != nil {
return err
}
if _, err := io.Copy(combined, pf); err != nil {
return err
}
_ = pf.Close()
if err := os.Remove(part); err != nil {
return err
}
}
return nil
}