@cryptotaxi247 / kubo
Code Commits Issues Pulls Sessions CI/CD
master
master
go 294 lines 6.87 KB
Raw
1 package profile
2
3 import (
4 "archive/zip"
5 "bytes"
6 "context"
7 "encoding/json"
8 "fmt"
9 "io"
10 "os"
11 "runtime"
12 "runtime/pprof"
13 "runtime/trace"
14 "sync"
15 "time"
16
17 "github.com/ipfs/go-log/v2"
18 version "github.com/ipfs/kubo"
19 )
20
21 const (
22 CollectorGoroutinesStack = "goroutines-stack"
23 CollectorGoroutinesPprof = "goroutines-pprof"
24 CollectorVersion = "version"
25 CollectorHeap = "heap"
26 CollectorAllocs = "allocs"
27 CollectorBin = "bin"
28 CollectorCPU = "cpu"
29 CollectorMutex = "mutex"
30 CollectorBlock = "block"
31 CollectorTrace = "trace"
32 )
33
34 var (
35 logger = log.Logger("profile")
36 goos = runtime.GOOS
37 )
38
39 type collector struct {
40 outputFile string
41 isExecutable bool
42 collectFunc func(ctx context.Context, opts Options, writer io.Writer) error
43 enabledFunc func(opts Options) bool
44 }
45
46 func (p *collector) outputFileName() string {
47 fName := p.outputFile
48 if p.isExecutable {
49 if goos == "windows" {
50 fName += ".exe"
51 }
52 }
53 return fName
54 }
55
56 var collectors = map[string]collector{
57 CollectorGoroutinesStack: {
58 outputFile: "goroutines.stacks",
59 collectFunc: goroutineStacksText,
60 enabledFunc: func(opts Options) bool { return true },
61 },
62 CollectorGoroutinesPprof: {
63 outputFile: "goroutines.pprof",
64 collectFunc: goroutineStacksProto,
65 enabledFunc: func(opts Options) bool { return true },
66 },
67 CollectorVersion: {
68 outputFile: "version.json",
69 collectFunc: versionInfo,
70 enabledFunc: func(opts Options) bool { return true },
71 },
72 CollectorHeap: {
73 outputFile: "heap.pprof",
74 collectFunc: heapProfile,
75 enabledFunc: func(opts Options) bool { return true },
76 },
77 CollectorAllocs: {
78 outputFile: "allocs.pprof",
79 collectFunc: allocsProfile,
80 enabledFunc: func(opts Options) bool { return true },
81 },
82 CollectorBin: {
83 outputFile: "ipfs",
84 isExecutable: true,
85 collectFunc: binary,
86 enabledFunc: func(opts Options) bool { return true },
87 },
88 CollectorCPU: {
89 outputFile: "cpu.pprof",
90 collectFunc: profileCPU,
91 enabledFunc: func(opts Options) bool { return opts.ProfileDuration > 0 },
92 },
93 CollectorMutex: {
94 outputFile: "mutex.pprof",
95 collectFunc: mutexProfile,
96 enabledFunc: func(opts Options) bool { return opts.ProfileDuration > 0 && opts.MutexProfileFraction > 0 },
97 },
98 CollectorBlock: {
99 outputFile: "block.pprof",
100 collectFunc: blockProfile,
101 enabledFunc: func(opts Options) bool { return opts.ProfileDuration > 0 && opts.BlockProfileRate > 0 },
102 },
103 CollectorTrace: {
104 outputFile: "trace",
105 collectFunc: captureTrace,
106 enabledFunc: func(opts Options) bool { return opts.ProfileDuration > 0 },
107 },
108 }
109
110 type Options struct {
111 Collectors []string
112 ProfileDuration time.Duration
113 MutexProfileFraction int
114 BlockProfileRate time.Duration
115 }
116
117 func WriteProfiles(ctx context.Context, archive *zip.Writer, opts Options) error {
118 p := profiler{
119 archive: archive,
120 opts: opts,
121 }
122 return p.runProfile(ctx)
123 }
124
125 // profiler runs the collectors concurrently and writes the results to the zip archive.
126 type profiler struct {
127 archive *zip.Writer
128 opts Options
129 }
130
131 func (p *profiler) runProfile(ctx context.Context) error {
132 type profileResult struct {
133 fName string
134 buf *bytes.Buffer
135 err error
136 }
137
138 ctx, cancelFn := context.WithCancel(ctx)
139 defer cancelFn()
140
141 collectorsToRun := make([]collector, len(p.opts.Collectors))
142 for i, name := range p.opts.Collectors {
143 c, ok := collectors[name]
144 if !ok {
145 return fmt.Errorf("unknown collector '%s'", name)
146 }
147 collectorsToRun[i] = c
148 }
149
150 results := make(chan profileResult, len(p.opts.Collectors))
151 wg := sync.WaitGroup{}
152 for _, c := range collectorsToRun {
153 if !c.enabledFunc(p.opts) {
154 continue
155 }
156
157 fName := c.outputFileName()
158
159 wg.Add(1)
160 go func(c collector) {
161 defer wg.Done()
162 logger.Infow("collecting profile", "File", fName)
163 defer logger.Infow("profile done", "File", fName)
164 b := bytes.Buffer{}
165 err := c.collectFunc(ctx, p.opts, &b)
166 if err != nil {
167 select {
168 case results <- profileResult{err: fmt.Errorf("generating profile data for %q: %w", fName, err)}:
169 case <-ctx.Done():
170 return
171 }
172 }
173 select {
174 case results <- profileResult{buf: &b, fName: fName}:
175 case <-ctx.Done():
176 }
177 }(c)
178 }
179 go func() {
180 wg.Wait()
181 close(results)
182 }()
183
184 for res := range results {
185 if res.err != nil {
186 return res.err
187 }
188 out, err := p.archive.Create(res.fName)
189 if err != nil {
190 return fmt.Errorf("creating output file %q: %w", res.fName, err)
191 }
192 _, err = io.Copy(out, res.buf)
193 if err != nil {
194 return fmt.Errorf("compressing result %q: %w", res.fName, err)
195 }
196 }
197
198 return nil
199 }
200
201 func goroutineStacksText(ctx context.Context, _ Options, w io.Writer) error {
202 return WriteAllGoroutineStacks(w)
203 }
204
205 func goroutineStacksProto(ctx context.Context, _ Options, w io.Writer) error {
206 return pprof.Lookup("goroutine").WriteTo(w, 0)
207 }
208
209 func heapProfile(ctx context.Context, _ Options, w io.Writer) error {
210 return pprof.Lookup("heap").WriteTo(w, 0)
211 }
212
213 func allocsProfile(ctx context.Context, _ Options, w io.Writer) error {
214 return pprof.Lookup("allocs").WriteTo(w, 0)
215 }
216
217 func versionInfo(ctx context.Context, _ Options, w io.Writer) error {
218 return json.NewEncoder(w).Encode(version.GetVersionInfo())
219 }
220
221 func binary(ctx context.Context, _ Options, w io.Writer) error {
222 var (
223 path string
224 err error
225 )
226 if goos == "linux" {
227 pid := os.Getpid()
228 path = fmt.Sprintf("/proc/%d/exe", pid)
229 } else {
230 path, err = os.Executable()
231 if err != nil {
232 return fmt.Errorf("finding binary path: %w", err)
233 }
234 }
235 fi, err := os.Open(path)
236 if err != nil {
237 return fmt.Errorf("opening binary %q: %w", path, err)
238 }
239 _, err = io.Copy(w, fi)
240 _ = fi.Close()
241 if err != nil {
242 return fmt.Errorf("copying binary %q: %w", path, err)
243 }
244 return nil
245 }
246
247 func mutexProfile(ctx context.Context, opts Options, w io.Writer) error {
248 prev := runtime.SetMutexProfileFraction(opts.MutexProfileFraction)
249 defer runtime.SetMutexProfileFraction(prev)
250 err := waitOrCancel(ctx, opts.ProfileDuration)
251 if err != nil {
252 return err
253 }
254 return pprof.Lookup("mutex").WriteTo(w, 2)
255 }
256
257 func blockProfile(ctx context.Context, opts Options, w io.Writer) error {
258 runtime.SetBlockProfileRate(int(opts.BlockProfileRate.Nanoseconds()))
259 defer runtime.SetBlockProfileRate(0)
260 err := waitOrCancel(ctx, opts.ProfileDuration)
261 if err != nil {
262 return err
263 }
264 return pprof.Lookup("block").WriteTo(w, 2)
265 }
266
267 func profileCPU(ctx context.Context, opts Options, w io.Writer) error {
268 err := pprof.StartCPUProfile(w)
269 if err != nil {
270 return err
271 }
272 defer pprof.StopCPUProfile()
273 return waitOrCancel(ctx, opts.ProfileDuration)
274 }
275
276 func captureTrace(ctx context.Context, opts Options, w io.Writer) error {
277 err := trace.Start(w)
278 if err != nil {
279 return err
280 }
281 defer trace.Stop()
282 return waitOrCancel(ctx, opts.ProfileDuration)
283 }
284
285 func waitOrCancel(ctx context.Context, d time.Duration) error {
286 timer := time.NewTimer(d)
287 defer timer.Stop()
288 select {
289 case <-timer.C:
290 return nil
291 case <-ctx.Done():
292 return ctx.Err()
293 }
294 }