feat(profiling): Annotate flamegraph with profile data

CHANGELOG.md

@@ -21,6 +21,7 @@
 - Pass thread id to calltree generation ([#492](https://github.com/getsentry/vroom/pull/492))
 - Dual mode metrics endpoint ([#493](https://github.com/getsentry/vroom/pull/493))
 - Add optional generation of metrics during flamegraph aggregation ([#494](https://github.com/getsentry/vroom/pull/494))
+- Annotate flamegraph with profile data ([#501](https://github.com/getsentry/vroom/pull/501))
 
 **Bug Fixes**:
 

internal/chunk/chunk_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/getsentry/vroom/internal/nodetree"
 	"github.com/getsentry/vroom/internal/platform"
 	"github.com/getsentry/vroom/internal/testutil"
+	"github.com/getsentry/vroom/internal/utils"
 )
 
 func TestCallTrees(t *testing.T) {
@@ -47,6 +48,7 @@ func TestCallTrees(t *testing.T) {
 						StartNS:       10_000_000,
 						Frame:         frame.Frame{Function: "function0"},
 						ProfileIDs:    make(map[string]struct{}),
+						Profiles:      make(map[utils.ExampleMetadata]struct{}),
 						Children: []*nodetree.Node{
 							{
 								DurationNS:    40_000_000,
@@ -58,6 +60,7 @@ func TestCallTrees(t *testing.T) {
 								SampleCount:   2,
 								Frame:         frame.Frame{Function: "function1"},
 								ProfileIDs:    make(map[string]struct{}),
+								Profiles:      make(map[utils.ExampleMetadata]struct{}),
 								Children: []*nodetree.Node{
 									{
 										DurationNS:    10_000_000,
@@ -69,6 +72,7 @@ func TestCallTrees(t *testing.T) {
 										StartNS:       40_000_000,
 										Frame:         frame.Frame{Function: "function2"},
 										ProfileIDs:    make(map[string]struct{}),
+										Profiles:      make(map[utils.ExampleMetadata]struct{}),
 									},
 								},
 							},
@@ -108,6 +112,7 @@ func TestCallTrees(t *testing.T) {
 						StartNS:       10_000_000,
 						Frame:         frame.Frame{Function: "function0"},
 						ProfileIDs:    make(map[string]struct{}),
+						Profiles:      make(map[utils.ExampleMetadata]struct{}),
 						Children: []*nodetree.Node{
 							{
 								DurationNS:    30_000_000,
@@ -119,6 +124,7 @@ func TestCallTrees(t *testing.T) {
 								StartNS:       10_000_000,
 								Frame:         frame.Frame{Function: "function1"},
 								ProfileIDs:    make(map[string]struct{}),
+								Profiles:      make(map[utils.ExampleMetadata]struct{}),
 							},
 						},
 					},
@@ -158,6 +164,7 @@ func TestCallTrees(t *testing.T) {
 						StartNS:       10_000_000,
 						Frame:         frame.Frame{Function: "function0"},
 						ProfileIDs:    make(map[string]struct{}),
+						Profiles:      make(map[utils.ExampleMetadata]struct{}),
 					},
 					{
 						DurationNS:    10_000_000,
@@ -169,6 +176,7 @@ func TestCallTrees(t *testing.T) {
 						StartNS:       20_000_000,
 						Frame:         frame.Frame{Function: "function1"},
 						ProfileIDs:    make(map[string]struct{}),
+						Profiles:      make(map[utils.ExampleMetadata]struct{}),
 					},
 				},
 			},

internal/chunk/readjob.go

@@ -15,18 +15,20 @@ type (
 		ProjectID      uint64
 		ProfilerID     string
 		ChunkID        string
+		TransactionID  string
 		ThreadID       *string
 		Start          uint64
 		End            uint64
 		Result         chan<- storageutil.ReadJobResult
 	}
 
 	ReadJobResult struct {
-		Err      error
-		Chunk    Chunk
-		ThreadID *string
-		Start    uint64
-		End      uint64
+		Err           error
+		Chunk         Chunk
+		TransactionID string
+		ThreadID      *string
+		Start         uint64
+		End           uint64
 	}
 )
 

internal/flamegraph/flamegraph.go

@@ -130,7 +130,7 @@ func GetFlamegraphFromProfiles(
 	for pair := range callTreesQueue {
 		profileID := pair.First
 		for _, callTree := range pair.Second {
-			addCallTreeToFlamegraph(&flamegraphTree, callTree, profileID)
+			addCallTreeToFlamegraph(&flamegraphTree, callTree, annotateWithProfileID(profileID))
 		}
 		countProfAggregated++
 	}
@@ -157,14 +157,26 @@ func sumNodesSampleCount(nodes []*nodetree.Node) int {
 	return c
 }
 
-func addCallTreeToFlamegraph(flamegraphTree *[]*nodetree.Node, callTree []*nodetree.Node, profileID string) {
+func annotateWithProfileID(profileID string) func(n *nodetree.Node) {
+	return func(n *nodetree.Node) {
+		n.ProfileIDs[profileID] = void
+	}
+}
+
+func annotateWithProfileExample(example utils.ExampleMetadata) func(n *nodetree.Node) {
+	return func(n *nodetree.Node) {
+		n.Profiles[example] = void
+	}
+}
+
+func addCallTreeToFlamegraph(flamegraphTree *[]*nodetree.Node, callTree []*nodetree.Node, annotate func(n *nodetree.Node)) {
 	for _, node := range callTree {
 		if existingNode := getMatchingNode(flamegraphTree, node); existingNode != nil {
 			existingNode.SampleCount += node.SampleCount
 			existingNode.DurationNS += node.DurationNS
-			addCallTreeToFlamegraph(&existingNode.Children, node.Children, profileID)
+			addCallTreeToFlamegraph(&existingNode.Children, node.Children, annotate)
 			if node.SampleCount > sumNodesSampleCount(node.Children) {
-				existingNode.ProfileIDs[profileID] = void
+				annotate(existingNode)
 			}
 		} else {
 			*flamegraphTree = append(*flamegraphTree, node)
@@ -173,40 +185,43 @@ func addCallTreeToFlamegraph(flamegraphTree *[]*nodetree.Node, callTree []*nodet
 			// to the right children along the branch yet,
 			// therefore we call a utility that walk the branch
 			// and does it
-			expandCallTreeWithProfileID(node, profileID)
+			expandCallTreeWithProfileID(node, annotate)
 		}
 	}
 }
 
-func expandCallTreeWithProfileID(node *nodetree.Node, profileID string) {
+func expandCallTreeWithProfileID(node *nodetree.Node, annotate func(n *nodetree.Node)) {
 	// leaf frames: we  must add the profileID
 	if node.Children == nil {
-		node.ProfileIDs[profileID] = void
+		annotate(node)
 	} else {
 		childrenSampleCount := 0
 		for _, child := range node.Children {
 			childrenSampleCount += child.SampleCount
-			expandCallTreeWithProfileID(child, profileID)
+			expandCallTreeWithProfileID(child, annotate)
 		}
 		// If the children's sample count is less than the current
 		// nodes sample count, it means there are some samples
 		// ending at the current node. In this case, this node
 		// should also contain the profile ID
 		if node.SampleCount > childrenSampleCount {
-			node.ProfileIDs[profileID] = void
+			annotate(node)
 		}
 	}
 }
 
 type flamegraph struct {
 	samples           [][]int
 	samplesProfileIDs [][]int
+	samplesProfiles   [][]int
 	sampleCounts      []uint64
 	sampleDurationsNs []uint64
 	frames            []speedscope.Frame
 	framesIndex       map[string]int
 	profilesIDsIndex  map[string]int
 	profilesIDs       []string
+	profilesIndex     map[utils.ExampleMetadata]int
+	profiles          []utils.ExampleMetadata
 	endValue          uint64
 	minFreq           int
 }
@@ -217,6 +232,7 @@ func toSpeedscope(trees []*nodetree.Node, minFreq int, projectID uint64) speedsc
 		framesIndex:      make(map[string]int),
 		minFreq:          minFreq,
 		profilesIDsIndex: make(map[string]int),
+		profilesIndex:    make(map[utils.ExampleMetadata]int),
 		samples:          make([][]int, 0),
 		sampleCounts:     make([]uint64, 0),
 	}
@@ -229,6 +245,7 @@ func toSpeedscope(trees []*nodetree.Node, minFreq int, projectID uint64) speedsc
 	aggProfiles[0] = speedscope.SampledProfile{
 		Samples:           fd.samples,
 		SamplesProfiles:   fd.samplesProfileIDs,
+		SamplesExamples:   fd.samplesProfiles,
 		Weights:           fd.sampleCounts,
 		SampleCounts:      fd.sampleCounts,
 		SampleDurationsNs: fd.sampleDurationsNs,
@@ -247,6 +264,7 @@ func toSpeedscope(trees []*nodetree.Node, minFreq int, projectID uint64) speedsc
 		Shared: speedscope.SharedData{
 			Frames:     fd.frames,
 			ProfileIDs: fd.profilesIDs,
+			Profiles:   fd.profiles,
 		},
 		Profiles: aggProfiles,
 	}
@@ -284,7 +302,13 @@ func (f *flamegraph) visitCalltree(node *nodetree.Node, currentStack *[]int) {
 
 	// base case (when we reach leaf frames)
 	if node.Children == nil {
-		f.addSample(currentStack, uint64(node.SampleCount), node.DurationNS, node.ProfileIDs)
+		f.addSample(
+			currentStack,
+			uint64(node.SampleCount),
+			node.DurationNS,
+			node.ProfileIDs,
+			node.Profiles,
+		)
 	} else {
 		totChildrenSampleCount := 0
 		var totChildrenDuration uint64
@@ -301,20 +325,33 @@ func (f *flamegraph) visitCalltree(node *nodetree.Node, currentStack *[]int) {
 		diffCount := node.SampleCount - totChildrenSampleCount
 		diffDuration := node.DurationNS - totChildrenDuration
 		if diffCount >= f.minFreq {
-			f.addSample(currentStack, uint64(diffCount), diffDuration, node.ProfileIDs)
+			f.addSample(
+				currentStack,
+				uint64(diffCount),
+				diffDuration,
+				node.ProfileIDs,
+				node.Profiles,
+			)
 		}
 	}
 	// pop last element before returning
 	*currentStack = (*currentStack)[:len(*currentStack)-1]
 }
 
-func (f *flamegraph) addSample(stack *[]int, count uint64, duration uint64, profileIDs map[string]struct{}) {
+func (f *flamegraph) addSample(
+	stack *[]int,
+	count uint64,
+	duration uint64,
+	profileIDs map[string]struct{},
+	profiles map[utils.ExampleMetadata]struct{},
+) {
 	cp := make([]int, len(*stack))
 	copy(cp, *stack)
 	f.samples = append(f.samples, cp)
 	f.sampleCounts = append(f.sampleCounts, count)
 	f.sampleDurationsNs = append(f.sampleDurationsNs, duration)
 	f.samplesProfileIDs = append(f.samplesProfileIDs, f.getProfileIDsIndices(profileIDs))
+	f.samplesProfiles = append(f.samplesProfiles, f.getProfilesIndices(profiles))
 	f.endValue += count
 }
 
@@ -332,6 +369,20 @@ func (f *flamegraph) getProfileIDsIndices(profileIDs map[string]struct{}) []int
 	return indices
 }
 
+func (f *flamegraph) getProfilesIndices(profiles map[utils.ExampleMetadata]struct{}) []int {
+	indices := make([]int, 0, len(profiles))
+	for i := range profiles {
+		if idx, ok := f.profilesIndex[i]; ok {
+			indices = append(indices, idx)
+		} else {
+			indices = append(indices, len(f.profiles))
+			f.profilesIndex[i] = len(f.profiles)
+			f.profiles = append(f.profiles, i)
+		}
+	}
+	return indices
+}
+
 func GetFlamegraphFromChunks(
 	ctx context.Context,
 	organizationID uint64,
@@ -388,9 +439,21 @@ func GetFlamegraphFromChunks(
 				continue
 			}
 			intervals := []utils.Interval{interval}
+
+			annotate := annotateWithProfileExample(
+				utils.NewExampleFromProfilerChunk(
+					result.Chunk.ProjectID,
+					result.Chunk.ProfilerID,
+					result.Chunk.ID,
+					result.TransactionID,
+					result.ThreadID,
+					result.Start,
+					result.End,
+				),
+			)
 			for _, callTree := range callTrees {
 				slicedTree := sliceCallTree(&callTree, &intervals)
-				addCallTreeToFlamegraph(&flamegraphTree, slicedTree, result.Chunk.ID)
+				addCallTreeToFlamegraph(&flamegraphTree, slicedTree, annotate)
 			}
 		}
 		countChunksAggregated++
@@ -472,9 +535,12 @@ func GetFlamegraphFromCandidates(
 				continue
 			}
 
+			annotate := annotateWithProfileExample(
+				utils.NewExampleFromProfileID(result.Profile.ProjectID(), result.Profile.ID()),
+			)
+
 			for _, callTree := range profileCallTrees {
-				// TODO: properly pass the profile ID around
-				addCallTreeToFlamegraph(&flamegraphTree, callTree, "")
+				addCallTreeToFlamegraph(&flamegraphTree, callTree, annotate)
 			}
 			// if metrics aggregator is not null, while we're at it,
 			// compute the metrics as well
@@ -492,6 +558,18 @@ func GetFlamegraphFromCandidates(
 				continue
 			}
 
+			annotate := annotateWithProfileExample(
+				utils.NewExampleFromProfilerChunk(
+					result.Chunk.ProjectID,
+					result.Chunk.ProfilerID,
+					result.Chunk.ID,
+					result.TransactionID,
+					result.ThreadID,
+					result.Start,
+					result.End,
+				),
+			)
+
 			for _, callTree := range chunkCallTrees {
 				if result.Start > 0 && result.End > 0 {
 					interval := utils.Interval{
@@ -500,8 +578,7 @@ func GetFlamegraphFromCandidates(
 					}
 					callTree = sliceCallTree(&callTree, &[]utils.Interval{interval})
 				}
-				// TODO: properly pass the profile ID around
-				addCallTreeToFlamegraph(&flamegraphTree, callTree, "")
+				addCallTreeToFlamegraph(&flamegraphTree, callTree, annotate)
 			}
 			// if metrics aggregator is not null, while we're at it,
 			// compute the metrics as well