mirror of
https://github.com/pomerium/pomerium.git
synced 2025-04-30 10:56:28 +02:00
101 lines
2.8 KiB
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
|
|
}
|