1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
|
package internal
import (
"context"
"os"
"runtime"
"runtime/pprof"
"runtime/trace"
"sync"
"time"
"ior/internal/flags"
)
// profilingControl manages optional CPU, memory, and execution-trace profiling
// for a single tracing run.
type profilingControl struct {
done chan struct{}
enabled bool
cpuProfile *os.File
memProfile *os.File
stopExecTrace func()
stopOnce sync.Once
}
// setupProfiling starts profiling if cfg.PprofEnable is set and returns a
// control handle. The caller must wait on control.done after the trace ends.
// started is non-nil in TUI mode; nil in plain/flamegraph mode.
func setupProfiling(ctx context.Context, cfg flags.Config, started chan<- struct{}) (*profilingControl, error) {
control := &profilingControl{
done: make(chan struct{}),
stopExecTrace: func() {},
}
if !cfg.PprofEnable {
close(control.done)
return control, nil
}
control.enabled = true
cpuProfilePath, memProfilePath, execTracePath, execTraceDuration := profilingFilesForMode(started != nil)
cpuProfile, memProfile, err := openProfilingFiles(cpuProfilePath, memProfilePath)
if err != nil {
return nil, err
}
control.cpuProfile = cpuProfile
control.memProfile = memProfile
if execTracePath != "" {
if err := startExecTrace(ctx, execTracePath, execTraceDuration, control); err != nil {
_ = cpuProfile.Close()
_ = memProfile.Close()
return nil, err
}
}
if err := pprof.StartCPUProfile(cpuProfile); err != nil {
control.stopExecTrace()
_ = cpuProfile.Close()
_ = memProfile.Close()
return nil, err
}
return control, nil
}
// openProfilingFiles creates the CPU and memory profile output files. On
// error any successfully opened file is closed before returning.
func openProfilingFiles(cpuPath, memPath string) (*os.File, *os.File, error) {
cpuProfile, err := os.Create(cpuPath)
if err != nil {
return nil, nil, err
}
memProfile, err := os.Create(memPath)
if err != nil {
_ = cpuProfile.Close()
return nil, nil, err
}
return cpuProfile, memProfile, nil
}
// startExecTrace creates the execution-trace output file, starts the runtime
// tracer, and wires a goroutine that stops it on context cancellation or after
// execTraceDuration, whichever comes first.
func startExecTrace(ctx context.Context, tracePath string, execTraceDuration time.Duration, control *profilingControl) error {
execTraceProfile, err := os.Create(tracePath)
if err != nil {
return err
}
if err := trace.Start(execTraceProfile); err != nil {
_ = execTraceProfile.Close()
return err
}
var stopOnce sync.Once
control.stopExecTrace = func() {
stopOnce.Do(func() {
trace.Stop()
_ = execTraceProfile.Close()
})
}
go func() {
timer := time.NewTimer(execTraceDuration)
defer timer.Stop()
select {
case <-ctx.Done():
case <-timer.C:
}
control.stopExecTrace()
}()
return nil
}
func (p *profilingControl) stop(logln func(...any)) {
p.stopOnce.Do(func() {
if !p.enabled {
return
}
logln("Stopping profiling and writing profile files")
pprof.StopCPUProfile()
runtime.GC()
// Log any failure writing the heap profile (e.g. full disk, permission
// denied) so it is not silently swallowed.
if err := pprof.WriteHeapProfile(p.memProfile); err != nil {
logln("ERROR: failed to write heap profile:", err)
}
p.stopExecTrace()
_ = p.cpuProfile.Close()
_ = p.memProfile.Close()
close(p.done)
})
}
// profilingFilesForMode returns the file paths and exec-trace duration to use
// depending on whether the binary is running in TUI mode or plain/flamegraph mode.
func profilingFilesForMode(tuiMode bool) (cpuProfilePath, memProfilePath, execTracePath string, execTraceDuration time.Duration) {
if tuiMode {
return "ior-tui-cpu.prof", "ior-tui-mem.prof", "ior-tui-trace.out", 10 * time.Second
}
return "ior.cpuprofile", "ior.memprofile", "", 0
}
|
