diff --git a/go/vt/servenv/pprof.go b/go/vt/servenv/pprof.go index 7c3ec124e9b..e81f605d6fc 100644 --- a/go/vt/servenv/pprof.go +++ b/go/vt/servenv/pprof.go @@ -57,10 +57,11 @@ func (p profmode) filename() string { } type profile struct { - mode profmode - rate int - path string - quiet bool + mode profmode + rate int + path string + quiet bool + waitSig bool } func parseProfileFlag(pf string) (*profile, error) { @@ -126,6 +127,15 @@ func parseProfileFlag(pf string) (*profile, error) { if err != nil { return nil, fmt.Errorf("invalid quiet flag %q: %v", fields[1], err) } + case "waitSig": + if len(fields) == 1 { + p.waitSig = true + continue + } + p.waitSig, err = strconv.ParseBool(fields[1]) + if err != nil { + return nil, fmt.Errorf("invalid waitSig flag %q: %v", fields[1], err) + } default: return nil, fmt.Errorf("unknown flag: %q", fields[0]) } @@ -136,6 +146,16 @@ func parseProfileFlag(pf string) (*profile, error) { var profileStarted uint32 +func startCallback(start func()) func() { + return func() { + if atomic.CompareAndSwapUint32(&profileStarted, 0, 1) { + start() + } else { + log.Fatal("profile: Start() already called") + } + } +} + func stopCallback(stop func()) func() { return func() { if atomic.CompareAndSwapUint32(&profileStarted, 1, 0) { @@ -144,14 +164,11 @@ func stopCallback(stop func()) func() { } } -// start begins the configured profiling process and returns a cleanup function -// that must be executed before process termination to flush the profile to disk. +// init returns a start function that 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") - } - +func (prof *profile) init() (start func(), stop func()) { var ( path string err error @@ -181,16 +198,21 @@ func (prof *profile) start() func() { switch prof.mode { case profileCPU: - pprof.StartCPUProfile(f) - return stopCallback(func() { + start = startCallback(func() { + pprof.StartCPUProfile(f) + }) + stop = stopCallback(func() { pprof.StopCPUProfile() f.Close() }) + return start, stop case profileMemHeap, profileMemAllocs: old := runtime.MemProfileRate - runtime.MemProfileRate = prof.rate - return stopCallback(func() { + start = startCallback(func() { + runtime.MemProfileRate = prof.rate + }) + stop = stopCallback(func() { tt := "heap" if prof.mode == profileMemAllocs { tt = "allocs" @@ -199,49 +221,63 @@ func (prof *profile) start() func() { f.Close() runtime.MemProfileRate = old }) + return start, stop case profileMutex: - runtime.SetMutexProfileFraction(prof.rate) - return stopCallback(func() { + start = startCallback(func() { + runtime.SetMutexProfileFraction(prof.rate) + }) + stop = stopCallback(func() { if mp := pprof.Lookup("mutex"); mp != nil { mp.WriteTo(f, 0) } f.Close() runtime.SetMutexProfileFraction(0) }) + return start, stop case profileBlock: - runtime.SetBlockProfileRate(prof.rate) - return stopCallback(func() { + start = startCallback(func() { + runtime.SetBlockProfileRate(prof.rate) + }) + stop = stopCallback(func() { pprof.Lookup("block").WriteTo(f, 0) f.Close() runtime.SetBlockProfileRate(0) }) + return start, stop case profileThreads: - return stopCallback(func() { + start = startCallback(func() {}) + stop = stopCallback(func() { if mp := pprof.Lookup("threadcreate"); mp != nil { mp.WriteTo(f, 0) } f.Close() }) + return start, stop case profileTrace: - if err := trace.Start(f); err != nil { - log.Fatalf("pprof: could not start trace: %v", err) - } - return stopCallback(func() { + start = startCallback(func() { + if err := trace.Start(f); err != nil { + log.Fatalf("pprof: could not start trace: %v", err) + } + }) + stop = stopCallback(func() { trace.Stop() f.Close() }) + return start, stop case profileGoroutine: - return stopCallback(func() { + start = startCallback(func() {}) + stop = stopCallback(func() { if mp := pprof.Lookup("goroutine"); mp != nil { mp.WriteTo(f, 0) } f.Close() }) + return start, stop default: panic("unsupported profile mode") @@ -255,12 +291,20 @@ func init() { log.Fatal(err) } if prof != nil { - stop := prof.start() + ch := make(chan os.Signal, 1) + signal.Notify(ch, syscall.SIGUSR1) + start, stop := prof.init() + + if prof.waitSig { + go func() { + <-ch + start() + }() + } else { + start() + } go func() { - ch := make(chan os.Signal, 1) - signal.Notify(ch, syscall.SIGUSR1) - <-ch stop() }() diff --git a/go/vt/servenv/pprof_test.go b/go/vt/servenv/pprof_test.go index 23d9a00fbd7..8ccabc773ff 100644 --- a/go/vt/servenv/pprof_test.go +++ b/go/vt/servenv/pprof_test.go @@ -29,6 +29,8 @@ func TestParseProfileFlag(t *testing.T) { {"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}, + {"cpu,waitSig", &profile{mode: profileCPU, waitSig: true}, false}, + {"cpu,path=a/b,waitSig", &profile{mode: profileCPU, waitSig: true, path: "a/b"}, false}, } for _, tt := range tests { t.Run(tt.arg, func(t *testing.T) {