diff --git a/go/vt/servenv/pprof.go b/go/vt/servenv/pprof.go index 28c31dee347..8a696a336b1 100644 --- a/go/vt/servenv/pprof.go +++ b/go/vt/servenv/pprof.go @@ -18,27 +18,235 @@ package servenv import ( "flag" + "fmt" + "io/ioutil" "os" + "path/filepath" + "runtime" "runtime/pprof" + "runtime/trace" + "strconv" + "strings" + "sync/atomic" "vitess.io/vitess/go/vt/log" ) var ( - cpuProfile = flag.String("cpu_profile", "", "write cpu profile to file") + _ = flag.String("cpu_profile", "", "deprecated: use '-pprof=cpu' instead") + pprofFlag = flag.String("pprof", "", "enable profiling") ) -func init() { - OnInit(func() { - if *cpuProfile != "" { - f, err := os.Create(*cpuProfile) +type profmode string + +const ( + profileCPU profmode = "cpu" + profileMemHeap profmode = "mem_heap" + profileMemAllocs profmode = "mem_allocs" + profileMutex profmode = "mutex" + profileBlock profmode = "block" + profileTrace profmode = "trace" + profileThreads profmode = "threads" + profileGoroutine profmode = "goroutine" +) + +func (p profmode) filename() string { + return fmt.Sprintf("%s.pprof", string(p)) +} + +type profile struct { + mode profmode + rate int + path string + quiet bool +} + +func parseProfileFlag(pf string) (*profile, error) { + if pf == "" { + return nil, nil + } + + var p profile + + items := strings.Split(pf, ",") + switch items[0] { + case "cpu": + p.mode = profileCPU + case "mem", "mem=heap": + p.mode = profileMemHeap + p.rate = 4096 + case "mem=allocs": + p.mode = profileMemAllocs + p.rate = 4096 + case "mutex": + p.mode = profileMutex + p.rate = 1 + case "block": + p.mode = profileBlock + p.rate = 1 + case "trace": + p.mode = profileTrace + case "threads": + p.mode = profileThreads + case "goroutine": + p.mode = profileGoroutine + default: + return nil, fmt.Errorf("unknown profile mode: %q", items[0]) + } + + for _, kv := range items[1:] { + var err error + fields := strings.SplitN(kv, "=", 2) + + switch fields[0] { + case "rate": + if len(fields) == 1 { + return nil, fmt.Errorf("missing value for 'rate'") + } + p.rate, err = strconv.Atoi(fields[1]) if err != nil { - log.Fatalf("Failed to create profile file: %v", err) + return nil, fmt.Errorf("invalid profile rate %q: %v", fields[1], err) } - pprof.StartCPUProfile(f) - OnTerm(func() { - pprof.StopCPUProfile() - }) + + case "path": + if len(fields) == 1 { + return nil, fmt.Errorf("missing value for 'path'") + } + p.path = fields[1] + + case "quiet": + if len(fields) == 1 { + p.quiet = true + continue + } + + p.quiet, err = strconv.ParseBool(fields[1]) + if err != nil { + return nil, fmt.Errorf("invalid quiet flag %q: %v", fields[1], err) + } + default: + return nil, fmt.Errorf("unknown flag: %q", fields[0]) + } + } + + return &p, nil +} + +var profileStarted uint32 + +// start begins the configured profiling process and returns a cleanup function +// that must be executed before process termination to flush the profile to disk. +// Based on the profiling code in github.com/pkg/profile +func (prof *profile) start() func() { + if !atomic.CompareAndSwapUint32(&profileStarted, 0, 1) { + log.Fatal("profile: Start() already called") + } + + var ( + path string + err error + logf = func(format string, args ...interface{}) {} + ) + + if prof.path != "" { + path = prof.path + err = os.MkdirAll(path, 0777) + } else { + path, err = ioutil.TempDir("", "profile") + } + if err != nil { + log.Fatalf("pprof: could not create initial output directory: %v", err) + } + + if !prof.quiet { + logf = log.Infof + } + + fn := filepath.Join(path, prof.mode.filename()) + f, err := os.Create(fn) + if err != nil { + log.Fatalf("pprof: could not create profile %q: %v", fn, err) + } + logf("pprof: %s profiling enabled, %s", string(prof.mode), fn) + + switch prof.mode { + case profileCPU: + pprof.StartCPUProfile(f) + return func() { + pprof.StopCPUProfile() + f.Close() + } + + case profileMemHeap, profileMemAllocs: + old := runtime.MemProfileRate + runtime.MemProfileRate = prof.rate + return func() { + tt := "heap" + if prof.mode == profileMemAllocs { + tt = "allocs" + } + pprof.Lookup(tt).WriteTo(f, 0) + f.Close() + runtime.MemProfileRate = old + } + + case profileMutex: + runtime.SetMutexProfileFraction(prof.rate) + return func() { + if mp := pprof.Lookup("mutex"); mp != nil { + mp.WriteTo(f, 0) + } + f.Close() + runtime.SetMutexProfileFraction(0) + } + + case profileBlock: + runtime.SetBlockProfileRate(prof.rate) + return func() { + pprof.Lookup("block").WriteTo(f, 0) + f.Close() + runtime.SetBlockProfileRate(0) + } + + case profileThreads: + return func() { + if mp := pprof.Lookup("threadcreate"); mp != nil { + mp.WriteTo(f, 0) + } + f.Close() + } + + case profileTrace: + if err := trace.Start(f); err != nil { + log.Fatalf("pprof: could not start trace: %v", err) + } + return func() { + trace.Stop() + f.Close() + } + + case profileGoroutine: + return func() { + if mp := pprof.Lookup("goroutine"); mp != nil { + mp.WriteTo(f, 0) + } + f.Close() + } + + default: + panic("unsupported profile mode") + } +} + +func init() { + OnInit(func() { + prof, err := parseProfileFlag(*pprofFlag) + if err != nil { + log.Fatal(err) + } + if prof != nil { + stop := prof.start() + OnTerm(stop) } }) } diff --git a/go/vt/servenv/pprof_test.go b/go/vt/servenv/pprof_test.go new file mode 100644 index 00000000000..23d9a00fbd7 --- /dev/null +++ b/go/vt/servenv/pprof_test.go @@ -0,0 +1,45 @@ +package servenv + +import ( + "reflect" + "testing" +) + +func TestParseProfileFlag(t *testing.T) { + tests := []struct { + arg string + want *profile + wantErr bool + }{ + {"", nil, false}, + {"mem", &profile{mode: profileMemHeap, rate: 4096}, false}, + {"mem,rate=1234", &profile{mode: profileMemHeap, rate: 1234}, false}, + {"mem,rate", nil, true}, + {"mem,rate=foobar", nil, true}, + {"mem=allocs", &profile{mode: profileMemAllocs, rate: 4096}, false}, + {"mem=allocs,rate=420", &profile{mode: profileMemAllocs, rate: 420}, false}, + {"block", &profile{mode: profileBlock, rate: 1}, false}, + {"block,rate=4", &profile{mode: profileBlock, rate: 4}, false}, + {"cpu", &profile{mode: profileCPU}, false}, + {"cpu,quiet", &profile{mode: profileCPU, quiet: true}, false}, + {"cpu,quiet=true", &profile{mode: profileCPU, quiet: true}, false}, + {"cpu,quiet=false", &profile{mode: profileCPU, quiet: false}, false}, + {"cpu,quiet=foobar", nil, true}, + {"cpu,path=", &profile{mode: profileCPU, path: ""}, false}, + {"cpu,path", nil, true}, + {"cpu,path=a", &profile{mode: profileCPU, path: "a"}, false}, + {"cpu,path=a/b/c/d", &profile{mode: profileCPU, path: "a/b/c/d"}, false}, + } + for _, tt := range tests { + t.Run(tt.arg, func(t *testing.T) { + got, err := parseProfileFlag(tt.arg) + if (err != nil) != tt.wantErr { + t.Errorf("parseProfileFlag() error = %v, wantErr %v", err, tt.wantErr) + return + } + if !reflect.DeepEqual(got, tt.want) { + t.Errorf("parseProfileFlag() got = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/go/vt/servenv/servenv.go b/go/vt/servenv/servenv.go index 0303aef61ab..9695b48c1e5 100644 --- a/go/vt/servenv/servenv.go +++ b/go/vt/servenv/servenv.go @@ -33,7 +33,6 @@ import ( "net/url" "os" "os/signal" - "runtime" "strings" "sync" "syscall" @@ -56,11 +55,11 @@ var ( Port *int // Flags to alter the behavior of the library. - lameduckPeriod = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping") - onTermTimeout = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping") - memProfileRate = flag.Int("mem-profile-rate", 512*1024, "profile every n bytes allocated") - mutexProfileFraction = flag.Int("mutex-profile-fraction", 0, "profile every n mutex contention events (see runtime.SetMutexProfileFraction)") - catchSigpipe = flag.Bool("catch-sigpipe", false, "catch and ignore SIGPIPE on stdout and stderr if specified") + lameduckPeriod = flag.Duration("lameduck-period", 50*time.Millisecond, "keep running at least this long after SIGTERM before stopping") + onTermTimeout = flag.Duration("onterm_timeout", 10*time.Second, "wait no more than this for OnTermSync handlers before stopping") + _ = flag.Int("mem-profile-rate", 512*1024, "deprecated: use '-pprof=mem' instead") + _ = flag.Int("mutex-profile-fraction", 0, "deprecated: use '-pprof=mutex' instead") + catchSigpipe = flag.Bool("catch-sigpipe", false, "catch and ignore SIGPIPE on stdout and stderr if specified") // mutex used to protect the Init function mu sync.Mutex @@ -106,13 +105,6 @@ func Init() { log.Exitf("servenv.Init: running this as root makes no sense") } - runtime.MemProfileRate = *memProfileRate - - if *mutexProfileFraction != 0 { - log.Infof("setting mutex profile fraction to %v", *mutexProfileFraction) - runtime.SetMutexProfileFraction(*mutexProfileFraction) - } - // We used to set this limit directly, but you pretty much have to // use a root account to allow increasing a limit reliably. Dropping // privileges is also tricky. The best strategy is to make a shell