/
pprof.go
129 lines (114 loc) · 3.04 KB
/
pprof.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
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
package pprof
import (
"archive/tar"
"bytes"
"context"
"fmt"
"io"
"path"
"runtime"
"runtime/pprof"
"runtime/trace"
"time"
)
func SetGlobalProfiling(enabled bool) {
if enabled {
// Copy the rates used in 1.x.
runtime.MemProfileRate = 4096
runtime.SetBlockProfileRate(int(1 * time.Second))
runtime.SetMutexProfileFraction(1)
} else {
runtime.MemProfileRate = 0
runtime.SetBlockProfileRate(0)
runtime.SetMutexProfileFraction(0)
}
}
// collectAllProfiles generates a tarball containing:
// - goroutine profile
// - blocking profile
// - mutex profile
// - heap profile
// - allocations profile
// - (optionally) trace profile
// - (optionally) CPU profile
//
// All information is added to a tar archive and then compressed, before being
// returned to the requester as an archive file. Where profiles support debug
// parameters, the profile is collected with debug=1.
func collectAllProfiles(ctx context.Context, traceDuration time.Duration, cpuDuration time.Duration) (io.Reader, error) {
// prof describes a profile name and a debug value, or in the case of a CPU
// profile, the number of seconds to collect the profile for.
type prof struct {
Name string // name of profile
Duration time.Duration // duration of profile if applicable. currently only used by cpu and trace
}
var profiles = []prof{
{Name: "goroutine"},
{Name: "block"},
{Name: "mutex"},
{Name: "heap"},
{Name: "allocs"},
{Name: "threadcreate"},
}
if traceDuration > 0 {
profiles = append(profiles, prof{"trace", traceDuration})
}
if cpuDuration > 0 {
// We want to gather CPU profiles first, if enabled.
profiles = append([]prof{{"cpu", cpuDuration}}, profiles...)
}
tarball := &bytes.Buffer{}
buf := &bytes.Buffer{} // Temporary buffer for each profile/query result.
tw := tar.NewWriter(tarball)
// Collect and write out profiles.
for _, profile := range profiles {
switch profile.Name {
case "cpu":
if err := pprof.StartCPUProfile(buf); err != nil {
return nil, err
}
sleep(ctx, profile.Duration)
pprof.StopCPUProfile()
case "trace":
if err := trace.Start(buf); err != nil {
return nil, err
}
sleep(ctx, profile.Duration)
trace.Stop()
default:
prof := pprof.Lookup(profile.Name)
if prof == nil {
return nil, fmt.Errorf("unable to find profile %q", profile.Name)
}
if err := prof.WriteTo(buf, 0); err != nil {
return nil, err
}
}
// Write the profile file's header.
if err := tw.WriteHeader(&tar.Header{
Name: path.Join("profiles", profile.Name+".pb.gz"),
Mode: 0600,
Size: int64(buf.Len()),
}); err != nil {
return nil, err
}
// Write the profile file's data.
if _, err := tw.Write(buf.Bytes()); err != nil {
return nil, err
}
// Reset the buffer for the next profile.
buf.Reset()
}
// Close the tar writer.
if err := tw.Close(); err != nil {
return nil, err
}
return tarball, nil
}
// Adapted from net/http/pprof/pprof.go
func sleep(ctx context.Context, d time.Duration) {
select {
case <-time.After(d):
case <-ctx.Done():
}
}