profile.go raw
1 // Package profile provides a simple way to manage runtime/pprof
2 // profiling of your Go application.
3 package profile
4
5 import (
6 "io/ioutil"
7 "log"
8 "os"
9 "os/signal"
10 "path/filepath"
11 "runtime"
12 "runtime/pprof"
13 "runtime/trace"
14 "sync/atomic"
15
16 "github.com/felixge/fgprof"
17 )
18
19 const (
20 cpuMode = iota
21 memMode
22 mutexMode
23 blockMode
24 traceMode
25 threadCreateMode
26 goroutineMode
27 clockMode
28 )
29
30 // Profile represents an active profiling session.
31 type Profile struct {
32 // quiet suppresses informational messages during profiling.
33 quiet bool
34
35 // noShutdownHook controls whether the profiling package should
36 // hook SIGINT to write profiles cleanly.
37 noShutdownHook bool
38
39 // mode holds the type of profiling that will be made
40 mode int
41
42 // path holds the base path where various profiling files are written.
43 // If blank, the base path will be generated by ioutil.TempDir.
44 path string
45
46 // memProfileRate holds the rate for the memory profile.
47 memProfileRate int
48
49 // memProfileType holds the profile type for memory
50 // profiles. Allowed values are `heap` and `allocs`.
51 memProfileType string
52
53 // closer holds a cleanup function that run after each profile
54 closer func()
55
56 // stopped records if a call to profile.Stop has been made
57 stopped uint32
58 }
59
60 // NoShutdownHook controls whether the profiling package should
61 // hook SIGINT to write profiles cleanly.
62 // Programs with more sophisticated signal handling should set
63 // this to true and ensure the Stop() function returned from Start()
64 // is called during shutdown.
65 func NoShutdownHook(p *Profile) { p.noShutdownHook = true }
66
67 // Quiet suppresses informational messages during profiling.
68 func Quiet(p *Profile) { p.quiet = true }
69
70 // CPUProfile enables cpu profiling.
71 // It disables any previous profiling settings.
72 func CPUProfile(p *Profile) { p.mode = cpuMode }
73
74 // DefaultMemProfileRate is the default memory profiling rate.
75 // See also http://golang.org/pkg/runtime/#pkg-variables
76 const DefaultMemProfileRate = 4096
77
78 // MemProfile enables memory profiling.
79 // It disables any previous profiling settings.
80 func MemProfile(p *Profile) {
81 p.memProfileRate = DefaultMemProfileRate
82 p.mode = memMode
83 }
84
85 // MemProfileRate enables memory profiling at the preferred rate.
86 // It disables any previous profiling settings.
87 func MemProfileRate(rate int) func(*Profile) {
88 return func(p *Profile) {
89 p.memProfileRate = rate
90 p.mode = memMode
91 }
92 }
93
94 // MemProfileHeap changes which type of memory profiling to profile
95 // the heap.
96 func MemProfileHeap(p *Profile) {
97 p.memProfileType = "heap"
98 p.mode = memMode
99 }
100
101 // MemProfileAllocs changes which type of memory to profile
102 // allocations.
103 func MemProfileAllocs(p *Profile) {
104 p.memProfileType = "allocs"
105 p.mode = memMode
106 }
107
108 // MutexProfile enables mutex profiling.
109 // It disables any previous profiling settings.
110 func MutexProfile(p *Profile) { p.mode = mutexMode }
111
112 // BlockProfile enables block (contention) profiling.
113 // It disables any previous profiling settings.
114 func BlockProfile(p *Profile) { p.mode = blockMode }
115
116 // Trace profile enables execution tracing.
117 // It disables any previous profiling settings.
118 func TraceProfile(p *Profile) { p.mode = traceMode }
119
120 // ThreadcreationProfile enables thread creation profiling..
121 // It disables any previous profiling settings.
122 func ThreadcreationProfile(p *Profile) { p.mode = threadCreateMode }
123
124 // GoroutineProfile enables goroutine profiling.
125 // It disables any previous profiling settings.
126 func GoroutineProfile(p *Profile) { p.mode = goroutineMode }
127
128 // ClockProfile enables wall clock (fgprof) profiling.
129 // It disables any previous profiling settings.
130 func ClockProfile(p *Profile) { p.mode = clockMode }
131
132 // ProfilePath controls the base path where various profiling
133 // files are written. If blank, the base path will be generated
134 // by ioutil.TempDir.
135 func ProfilePath(path string) func(*Profile) {
136 return func(p *Profile) {
137 p.path = path
138 }
139 }
140
141 // Stop stops the profile and flushes any unwritten data.
142 func (p *Profile) Stop() {
143 if !atomic.CompareAndSwapUint32(&p.stopped, 0, 1) {
144 // someone has already called close
145 return
146 }
147 p.closer()
148 atomic.StoreUint32(&started, 0)
149 }
150
151 // started is non zero if a profile is running.
152 var started uint32
153
154 // Start starts a new profiling session.
155 // The caller should call the Stop method on the value returned
156 // to cleanly stop profiling.
157 func Start(options ...func(*Profile)) interface {
158 Stop()
159 } {
160 if !atomic.CompareAndSwapUint32(&started, 0, 1) {
161 log.Fatal("profile: Start() already called")
162 }
163
164 var prof Profile
165 for _, option := range options {
166 option(&prof)
167 }
168
169 path, err := func() (string, error) {
170 if p := prof.path; p != "" {
171 return p, os.MkdirAll(p, 0777)
172 }
173 return ioutil.TempDir("", "profile")
174 }()
175
176 if err != nil {
177 log.Fatalf("profile: could not create initial output directory: %v", err)
178 }
179
180 logf := func(format string, args ...interface{}) {
181 if !prof.quiet {
182 log.Printf(format, args...)
183 }
184 }
185
186 if prof.memProfileType == "" {
187 prof.memProfileType = "heap"
188 }
189
190 switch prof.mode {
191 case cpuMode:
192 fn := filepath.Join(path, "cpu.pprof")
193 f, err := os.Create(fn)
194 if err != nil {
195 log.Fatalf("profile: could not create cpu profile %q: %v", fn, err)
196 }
197 logf("profile: cpu profiling enabled, %s", fn)
198 pprof.StartCPUProfile(f)
199 prof.closer = func() {
200 pprof.StopCPUProfile()
201 f.Close()
202 logf("profile: cpu profiling disabled, %s", fn)
203 }
204
205 case memMode:
206 fn := filepath.Join(path, "mem.pprof")
207 f, err := os.Create(fn)
208 if err != nil {
209 log.Fatalf("profile: could not create memory profile %q: %v", fn, err)
210 }
211 old := runtime.MemProfileRate
212 runtime.MemProfileRate = prof.memProfileRate
213 logf("profile: memory profiling enabled (rate %d), %s", runtime.MemProfileRate, fn)
214 prof.closer = func() {
215 pprof.Lookup(prof.memProfileType).WriteTo(f, 0)
216 f.Close()
217 runtime.MemProfileRate = old
218 logf("profile: memory profiling disabled, %s", fn)
219 }
220
221 case mutexMode:
222 fn := filepath.Join(path, "mutex.pprof")
223 f, err := os.Create(fn)
224 if err != nil {
225 log.Fatalf("profile: could not create mutex profile %q: %v", fn, err)
226 }
227 runtime.SetMutexProfileFraction(1)
228 logf("profile: mutex profiling enabled, %s", fn)
229 prof.closer = func() {
230 if mp := pprof.Lookup("mutex"); mp != nil {
231 mp.WriteTo(f, 0)
232 }
233 f.Close()
234 runtime.SetMutexProfileFraction(0)
235 logf("profile: mutex profiling disabled, %s", fn)
236 }
237
238 case blockMode:
239 fn := filepath.Join(path, "block.pprof")
240 f, err := os.Create(fn)
241 if err != nil {
242 log.Fatalf("profile: could not create block profile %q: %v", fn, err)
243 }
244 runtime.SetBlockProfileRate(1)
245 logf("profile: block profiling enabled, %s", fn)
246 prof.closer = func() {
247 pprof.Lookup("block").WriteTo(f, 0)
248 f.Close()
249 runtime.SetBlockProfileRate(0)
250 logf("profile: block profiling disabled, %s", fn)
251 }
252
253 case threadCreateMode:
254 fn := filepath.Join(path, "threadcreation.pprof")
255 f, err := os.Create(fn)
256 if err != nil {
257 log.Fatalf("profile: could not create thread creation profile %q: %v", fn, err)
258 }
259 logf("profile: thread creation profiling enabled, %s", fn)
260 prof.closer = func() {
261 if mp := pprof.Lookup("threadcreate"); mp != nil {
262 mp.WriteTo(f, 0)
263 }
264 f.Close()
265 logf("profile: thread creation profiling disabled, %s", fn)
266 }
267
268 case traceMode:
269 fn := filepath.Join(path, "trace.out")
270 f, err := os.Create(fn)
271 if err != nil {
272 log.Fatalf("profile: could not create trace output file %q: %v", fn, err)
273 }
274 if err := trace.Start(f); err != nil {
275 log.Fatalf("profile: could not start trace: %v", err)
276 }
277 logf("profile: trace enabled, %s", fn)
278 prof.closer = func() {
279 trace.Stop()
280 logf("profile: trace disabled, %s", fn)
281 }
282
283 case goroutineMode:
284 fn := filepath.Join(path, "goroutine.pprof")
285 f, err := os.Create(fn)
286 if err != nil {
287 log.Fatalf("profile: could not create goroutine profile %q: %v", fn, err)
288 }
289 logf("profile: goroutine profiling enabled, %s", fn)
290 prof.closer = func() {
291 if mp := pprof.Lookup("goroutine"); mp != nil {
292 mp.WriteTo(f, 0)
293 }
294 f.Close()
295 logf("profile: goroutine profiling disabled, %s", fn)
296 }
297
298 case clockMode:
299 fn := filepath.Join(path, "clock.pprof")
300 f, err := os.Create(fn)
301 if err != nil {
302 log.Fatalf("profile: could not create clock profile %q: %v", fn, err)
303 }
304 logf("profile: clock profiling enabled, %s", fn)
305 stop := fgprof.Start(f, fgprof.FormatPprof)
306 prof.closer = func() {
307 stop()
308 f.Close()
309 logf("profile: clock profiling disabled, %s", fn)
310 }
311 }
312
313 if !prof.noShutdownHook {
314 go func() {
315 c := make(chan os.Signal, 1)
316 signal.Notify(c, os.Interrupt)
317 <-c
318
319 log.Println("profile: caught interrupt, stopping profiles")
320 prof.Stop()
321
322 os.Exit(0)
323 }()
324 }
325
326 return &prof
327 }
328