diff --git a/internal/pprof/pprof.go b/internal/pprof/pprof.go index bcb68df180e..c3009ae35de 100644 --- a/internal/pprof/pprof.go +++ b/internal/pprof/pprof.go @@ -7,7 +7,7 @@ import ( "os" "path/filepath" "runtime" - "runtime/pprof" + runtimepprof "runtime/pprof" "sync" "time" ) @@ -19,6 +19,9 @@ type ProfileSession struct { logWriter io.Writer } +// startCPUProfile is a test seam for simulating runtime/pprof startup failures. +var startCPUProfile = runtimepprof.StartCPUProfile + // BeginProfiling starts CPU and memory profiling, writing the profiles to the specified directory. func BeginProfiling(profileDir string, logWriter io.Writer) *ProfileSession { if err := os.MkdirAll(profileDir, 0o755); err != nil { @@ -34,7 +37,9 @@ func BeginProfiling(profileDir string, logWriter io.Writer) *ProfileSession { panic(err) } - if err := pprof.StartCPUProfile(cpuFile); err != nil { + if err := startCPUProfile(cpuFile); err != nil { + cpuFile.Close() + os.Remove(cpuProfilePath) panic(err) } @@ -47,7 +52,7 @@ func BeginProfiling(profileDir string, logWriter io.Writer) *ProfileSession { } func (p *ProfileSession) Stop() { - pprof.StopCPUProfile() + runtimepprof.StopCPUProfile() p.cpuFile.Close() if p.memFilePath != "" { @@ -55,7 +60,7 @@ func (p *ProfileSession) Stop() { if err != nil { panic(err) } - if err := pprof.Lookup("allocs").WriteTo(memFile, 0); err != nil { + if err := runtimepprof.Lookup("allocs").WriteTo(memFile, 0); err != nil { panic(err) } memFile.Close() @@ -90,7 +95,7 @@ func (c *CPUProfiler) StartCPUProfile(profileDir string) error { return fmt.Errorf("failed to create CPU profile file: %w", err) } - if err := pprof.StartCPUProfile(cpuFile); err != nil { + if err := startCPUProfile(cpuFile); err != nil { cpuFile.Close() os.Remove(cpuProfilePath) return fmt.Errorf("failed to start CPU profile: %w", err) @@ -134,7 +139,7 @@ func SaveHeapProfile(profileDir string) (string, error) { defer heapFile.Close() runtime.GC() - if err := pprof.Lookup("heap").WriteTo(heapFile, 0); err != nil { + if err := runtimepprof.Lookup("heap").WriteTo(heapFile, 0); err != nil { os.Remove(heapProfilePath) return "", fmt.Errorf("failed to write heap profile: %w", err) } @@ -155,7 +160,7 @@ func SaveAllocProfile(profileDir string) (string, error) { } defer allocFile.Close() - if err := pprof.Lookup("allocs").WriteTo(allocFile, 0); err != nil { + if err := runtimepprof.Lookup("allocs").WriteTo(allocFile, 0); err != nil { os.Remove(allocProfilePath) return "", fmt.Errorf("failed to write alloc profile: %w", err) } diff --git a/internal/pprof/pprof_test.go b/internal/pprof/pprof_test.go new file mode 100644 index 00000000000..bba4cbf9518 --- /dev/null +++ b/internal/pprof/pprof_test.go @@ -0,0 +1,34 @@ +package pprof + +import ( + "errors" + "fmt" + "io" + "os" + "path/filepath" + "testing" + + "github.com/microsoft/typescript-go/internal/testutil" + "gotest.tools/v3/assert" +) + +func TestBeginProfilingCleansUpCPUFileWhenStartFails(t *testing.T) { + wantErr := errors.New("failed to start profile") + originalStartCPUProfile := startCPUProfile + startCPUProfile = func(io.Writer) error { + return wantErr + } + t.Cleanup(func() { + startCPUProfile = originalStartCPUProfile + }) + + profileDir := t.TempDir() + cpuProfilePath := filepath.Join(profileDir, fmt.Sprintf("%d-cpuprofile.pb.gz", os.Getpid())) + + testutil.AssertPanics(t, func() { + BeginProfiling(profileDir, io.Discard) + }, wantErr) + + _, err := os.Stat(cpuProfilePath) + assert.Assert(t, os.IsNotExist(err), "expected CPU profile file to be removed after start failure, got %v", err) +}