-
Notifications
You must be signed in to change notification settings - Fork 3.5k
/
profile.go
64 lines (55 loc) · 1.15 KB
/
profile.go
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
package profile
import (
"log"
"os"
"runtime"
"runtime/pprof"
)
type Config struct {
// CPU, if set, specifies the file name of the CPU profile to capture
CPU string
// Memory, if set, specifies the file name of the CPU profile to capture
Memory string
}
func (c *Config) noProfiles() bool {
return c.CPU == "" && c.Memory == ""
}
// Start starts a CPU and / or Memory profile if configured and returns a
// function that should be called to terminate the profiles.
func (c *Config) Start() func() {
if c.noProfiles() {
return func() { return }
}
var prof struct {
cpu *os.File
mem *os.File
}
if c.CPU != "" {
f, err := os.Create(c.CPU)
if err != nil {
log.Fatalf("cpuprofile: %v", err)
}
prof.cpu = f
_ = pprof.StartCPUProfile(prof.cpu)
}
if c.Memory != "" {
f, err := os.Create(c.Memory)
if err != nil {
log.Fatalf("memprofile: %v", err)
}
prof.mem = f
runtime.MemProfileRate = 4096
}
return func() {
if prof.cpu != nil {
pprof.StopCPUProfile()
_ = prof.cpu.Close()
prof.cpu = nil
}
if prof.mem != nil {
_ = pprof.Lookup("heap").WriteTo(prof.mem, 0)
_ = prof.mem.Close()
prof.mem = nil
}
}
}