Merge pull request #7616 from frouioui/start-pprof-on-signal

Addition of waitSig pprof argument to start recording on USR1

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()
 			}()

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) {