pprof: support limit the max number for dumping goroutines.

api/next.txt

@@ -0,0 +1 @@
+pkg runtime/pprof, func SetMaxDumpGoroutineNum(int)

src/runtime/mprof.go

@@ -755,6 +755,17 @@ func runtime_goroutineProfileWithLabels(p []StackRecord, labels []unsafe.Pointer
 
 // labels may be nil. If labels is non-nil, it must have the same length as p.
 func goroutineProfileWithLabels(p []StackRecord, labels []unsafe.Pointer) (n int, ok bool) {
+	return goroutineProfileWithLabelsAndLimit(p, labels, 0)
+}
+
+//go:linkname runtime_goroutineProfileWithLabelsAndLimit runtime/pprof.runtime_goroutineProfileWithLabelsAndLimit
+func runtime_goroutineProfileWithLabelsAndLimit(p []StackRecord, labels []unsafe.Pointer, max int) (n int, ok bool) {
+	return goroutineProfileWithLabelsAndLimit(p, labels, max)
+}
+
+// limit the max goroutine, sample them when return there are more than max goroutines.
+// 0 means not limitation.
+func goroutineProfileWithLabelsAndLimit(p []StackRecord, labels []unsafe.Pointer, max int) (m int, ok bool) {
 	if labels != nil && len(labels) != len(p) {
 		labels = nil
 	}
@@ -769,14 +780,19 @@ func goroutineProfileWithLabels(p []StackRecord, labels []unsafe.Pointer) (n int
 	stopTheWorld("profile")
 
 	// World is stopped, no locking required.
-	n = 1
+	n := 1
 	forEachGRace(func(gp1 *g) {
 		if isOK(gp1) {
 			n++
 		}
 	})
 
-	if n <= len(p) {
+	m = n
+	if max > 0 && n > max {
+		m = max
+	}
+
+	if m <= len(p) {
 		ok = true
 		r, lbl := p, labels
 
@@ -794,12 +810,22 @@ func goroutineProfileWithLabels(p []StackRecord, labels []unsafe.Pointer) (n int
 			lbl = lbl[1:]
 		}
 
+		got, i := 1, 0
+		rand := fastrand()
+
 		// Save other goroutines.
 		forEachGRace(func(gp1 *g) {
 			if !isOK(gp1) {
 				return
 			}
 
+			i++
+			// make sure we will have m goroutines.
+			if m < n && int(rand%uint32(n-i)) >= (m-got) {
+				return
+			}
+			got++
+
 			if len(r) == 0 {
 				// Should be impossible, but better to return a
 				// truncated profile than to crash the entire process.
@@ -819,7 +845,7 @@ func goroutineProfileWithLabels(p []StackRecord, labels []unsafe.Pointer) (n int
 	}
 
 	startTheWorld()
-	return n, ok
+	return m, ok
 }
 
 // GoroutineProfile returns n, the number of records in the active goroutine stack profile.

src/runtime/pprof/pprof.go

@@ -230,6 +230,14 @@ func Lookup(name string) *Profile {
 	return profiles.m[name]
 }
 
+// no limit by default
+var maxDumpGoroutineNum = 0
+
+// SetMaxDumpGoroutineNum set the max number for dumping goroutine
+func SetMaxDumpGoroutineNum(limit int) {
+	maxDumpGoroutineNum = limit
+}
+
 // Profiles returns a slice of all the known profiles, sorted by name.
 func Profiles() []*Profile {
 	lockProfiles()
@@ -664,7 +672,7 @@ func writeThreadCreate(w io.Writer, debug int) error {
 	// Until https://golang.org/issues/6104 is addressed, wrap
 	// ThreadCreateProfile because there's no point in tracking labels when we
 	// don't get any stack-traces.
-	return writeRuntimeProfile(w, debug, "threadcreate", func(p []runtime.StackRecord, _ []unsafe.Pointer) (n int, ok bool) {
+	return writeRuntimeProfile(w, debug, "threadcreate", func(p []runtime.StackRecord, _ []unsafe.Pointer, _ int) (n int, ok bool) {
 		return runtime.ThreadCreateProfile(p)
 	})
 }
@@ -674,15 +682,15 @@ func countGoroutine() int {
 	return runtime.NumGoroutine()
 }
 
-// runtime_goroutineProfileWithLabels is defined in runtime/mprof.go
-func runtime_goroutineProfileWithLabels(p []runtime.StackRecord, labels []unsafe.Pointer) (n int, ok bool)
+// runtime_goroutineProfileWithLabelsAndLimit is defined in runtime/mprof.go
+func runtime_goroutineProfileWithLabelsAndLimit(p []runtime.StackRecord, labels []unsafe.Pointer, max int) (n int, ok bool)
 
 // writeGoroutine writes the current runtime GoroutineProfile to w.
 func writeGoroutine(w io.Writer, debug int) error {
 	if debug >= 2 {
 		return writeGoroutineStacks(w)
 	}
-	return writeRuntimeProfile(w, debug, "goroutine", runtime_goroutineProfileWithLabels)
+	return writeRuntimeProfile(w, debug, "goroutine", runtime_goroutineProfileWithLabelsAndLimit)
 }
 
 func writeGoroutineStacks(w io.Writer) error {
@@ -706,7 +714,7 @@ func writeGoroutineStacks(w io.Writer) error {
 	return err
 }
 
-func writeRuntimeProfile(w io.Writer, debug int, name string, fetch func([]runtime.StackRecord, []unsafe.Pointer) (int, bool)) error {
+func writeRuntimeProfile(w io.Writer, debug int, name string, fetch func([]runtime.StackRecord, []unsafe.Pointer, int) (int, bool)) error {
 	// Find out how many records there are (fetch(nil)),
 	// allocate that many records, and get the data.
 	// There's a race—more records might be added between
@@ -715,14 +723,15 @@ func writeRuntimeProfile(w io.Writer, debug int, name string, fetch func([]runti
 	// The loop should only execute one iteration in the common case.
 	var p []runtime.StackRecord
 	var labels []unsafe.Pointer
-	n, ok := fetch(nil, nil)
+	max := maxDumpGoroutineNum
+	n, ok := fetch(nil, nil, max)
 	for {
 		// Allocate room for a slightly bigger profile,
 		// in case a few more entries have been added
 		// since the call to ThreadProfile.
 		p = make([]runtime.StackRecord, n+10)
 		labels = make([]unsafe.Pointer, n+10)
-		n, ok = fetch(p, labels)
+		n, ok = fetch(p, labels, max)
 		if ok {
 			p = p[0:n]
 			break