bytesprofile: add Profile.Collect

Add a method to bytesprofile.Profile to return the collected statistics in a
structured format. This can be used to return results from a SQL built-in in
CockroachDB.

Author: jbowens

internal/bytesprofile/bytesprofile.go

@@ -7,6 +7,7 @@ package bytesprofile
 import (
 	"cmp"
 	"fmt"
+	"iter"
 	"maps"
 	"runtime"
 	"slices"
@@ -56,6 +57,23 @@ func (p *Profile) Record(bytes int64) {
 	p.samples[stack] = curr
 }
 
+func (p *Profile) all() iter.Seq2[stack, aggSamples] {
+	return func(yield func(stack, aggSamples) bool) {
+		p.mu.Lock()
+		defer p.mu.Unlock()
+		// Sort the stacks by bytes in descending order.
+		uniqueStacks := slices.SortedFunc(maps.Keys(p.samples), func(a, b stack) int {
+			return -cmp.Compare(p.samples[a].bytes, p.samples[b].bytes)
+		})
+		for _, stack := range uniqueStacks {
+			samples := p.samples[stack]
+			if !yield(stack, samples) {
+				break
+			}
+		}
+	}
+}
+
 // TODO(jackson): We could add the ability to export the profile to a pprof file
 // (which internally is just a protocol buffer). Ideally the Go standard library
 // would provide facilities for this (e.g., golang/go#18454). The runtime/pprof
@@ -64,20 +82,15 @@ func (p *Profile) Record(bytes int64) {
 
 // String returns a string representation of the stacks captured by the profile.
 func (p *Profile) String() string {
-	p.mu.Lock()
-	defer p.mu.Unlock()
-	// Sort the stacks by bytes in descending order.
-	uniqueStacks := slices.SortedFunc(maps.Keys(p.samples), func(a, b stack) int {
-		return -cmp.Compare(p.samples[a].bytes, p.samples[b].bytes)
-	})
 	var sb strings.Builder
-	for i, stack := range uniqueStacks {
+	var i int
+	for stack, stats := range p.all() {
 		if i > 0 {
 			sb.WriteString("\n")
 		}
 		fmt.Fprintf(&sb, "%d: Count: %d (%s), Bytes: %d (%s)\n", i,
-			p.samples[stack].count, humanize.Count.Int64(p.samples[stack].count),
-			p.samples[stack].bytes, humanize.Bytes.Int64(p.samples[stack].bytes))
+			stats.count, humanize.Count.Int64(stats.count),
+			stats.bytes, humanize.Bytes.Int64(stats.bytes))
 		frames := runtime.CallersFrames(stack.trimmed())
 		for {
 			frame, more := frames.Next()
@@ -86,6 +99,39 @@ func (p *Profile) String() string {
 				break
 			}
 		}
+		i++
 	}
 	return sb.String()
 }
+
+// StackStats contains the stack trace and statistics for a given stack.
+type StackStats struct {
+	Stack string
+	Bytes int64
+	Count int64
+}
+
+// Collect returns a slice of StackStats, sorted by bytes in descending order.
+func (p *Profile) Collect() []StackStats {
+	var stats []StackStats
+	var sb strings.Builder
+	for stack, samples := range p.all() {
+		sb.Reset()
+		frames := runtime.CallersFrames(stack.trimmed())
+		for {
+			frame, more := frames.Next()
+			fmt.Fprintf(&sb, "  %s\n   %s:%d", frame.Function, frame.File, frame.Line)
+			if !more {
+				break
+			}
+			sb.WriteString("\n")
+		}
+
+		stats = append(stats, StackStats{
+			Stack: sb.String(),
+			Bytes: samples.bytes,
+			Count: samples.count,
+		})
+	}
+	return stats
+}

internal/bytesprofile/bytesprofile_test.go

@@ -21,5 +21,12 @@ func TestBytesProfile(t *testing.T) {
 	s := p.String()
 	require.Contains(t, s, "0: Count: 100 (100), Bytes: 4950 (4.8KB)")
 	require.Contains(t, s, "1: Count: 10 (10), Bytes: 495 (495B)")
+
+	stacks := p.Collect()
+	require.Len(t, stacks, 2)
+	require.Equal(t, stacks[0].Bytes, int64(4950))
+	require.Equal(t, stacks[0].Count, int64(100))
+	require.Equal(t, stacks[1].Bytes, int64(495))
+	require.Equal(t, stacks[1].Count, int64(10))
 	t.Log(s)
 }