generated from mesosphere/golang-repository-template
/
profiling_options.go
104 lines (93 loc) · 2.37 KB
/
profiling_options.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
// Copyright 2022 D2iQ, Inc. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
package root
import (
"fmt"
"os"
"os/signal"
"runtime"
"runtime/pprof"
"github.com/spf13/pflag"
)
// ProfilingOptions contains settings for profiling.
type ProfilingOptions struct {
profileName string
profileOutput string
}
// NewProfilingOptions initializes ProfilingOptions with defaults.
func NewProfilingOptions() *ProfilingOptions {
return &ProfilingOptions{
profileName: "none",
profileOutput: "profile.pprof",
}
}
// AddFlags adds flags for setting profiling options to the provided FlagSet.
func (o *ProfilingOptions) AddFlags(flagSet *pflag.FlagSet) {
flagSet.StringVar(
&o.profileName, "profile", o.profileName,
"Name of profile to capture. One of (none|cpu|goroutine|threadcreate|heap|allocs|block|mutex)",
)
_ = flagSet.MarkHidden("profile")
flagSet.StringVar(&o.profileOutput, "profile-output", o.profileOutput, "Name of the file to write the profile to")
_ = flagSet.MarkHidden("profile-output")
}
// InitProfiling starts profiling.
func (o *ProfilingOptions) InitProfiling() error {
switch o.profileName {
case "none":
return nil
case "cpu":
f, err := os.Create(o.profileOutput)
if err != nil {
return err
}
err = pprof.StartCPUProfile(f)
if err != nil {
return err
}
// Block and mutex profiles need a call to Set{Block,Mutex}ProfileRate to
// output anything. We choose to sample all events.
case "block":
runtime.SetBlockProfileRate(1)
case "mutex":
runtime.SetMutexProfileFraction(1)
default:
// Check the profile name is valid.
if profile := pprof.Lookup(o.profileName); profile == nil {
return fmt.Errorf("unknown profile '%s'", o.profileName)
}
}
// If the command is interrupted before the end (ctrl-c), flush the
// profiling files
c := make(chan os.Signal, 1)
signal.Notify(c, os.Interrupt)
go func() {
<-c
_ = o.FlushProfiling()
os.Exit(0)
}()
return nil
}
// FlushProfiling stops profiling and writes remaining unwritten data.
func (o *ProfilingOptions) FlushProfiling() error {
switch o.profileName {
case "none":
return nil
case "cpu":
pprof.StopCPUProfile()
case "heap":
runtime.GC()
fallthrough
default:
profile := pprof.Lookup(o.profileName)
if profile == nil {
return nil
}
f, err := os.Create(o.profileOutput)
if err != nil {
return err
}
_ = profile.WriteTo(f, 0)
}
return nil
}