Merge pull request #523 from getsentry/txiao/fix/annotate-functions-t…

…o-the-right-thread fix(flamegraphs): Annotate functions to the right thread
getsentry · Oct 25, 2024 · 2c9c372 · 2c9c372
2 parents 1273586 + 040c129
commit 2c9c372
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -46,6 +46,7 @@
 - Fix metrics example list ([#505](https://github.com/getsentry/vroom/pull/505))
 - Increase readjob channel size  ([#512](https://github.com/getsentry/vroom/pull/512))
 - Return the duration measured by the profiler. ([#516](https://github.com/getsentry/vroom/pull/516), [#517](https://github.com/getsentry/vroom/pull/517))
+- Annotate functions to the right thread. ([#523](https://github.com/getsentry/vroom/pull/523))
 
 **Internal**:
 

diff --git a/internal/flamegraph/flamegraph.go b/internal/flamegraph/flamegraph.go
@@ -559,32 +559,34 @@ func GetFlamegraphFromCandidates(
 			chunkProfileSpan := span.StartChild("calltree")
 			chunkProfileSpan.Description = "continuous profile"
 
-			example := utils.NewExampleFromProfilerChunk(
-				result.Chunk.ProjectID,
-				result.Chunk.ProfilerID,
-				result.Chunk.ID,
-				result.TransactionID,
-				result.ThreadID,
-				result.Start,
-				result.End,
-			)
-			annotate := annotateWithProfileExample(example)
-
-			for _, callTree := range result.CallTrees {
+			for threadID, callTree := range result.CallTrees {
 				if result.Start > 0 && result.End > 0 {
 					interval := utils.Interval{
 						Start: result.Start,
 						End:   result.End,
 					}
 					callTree = sliceCallTree(&callTree, &[]utils.Interval{interval})
 				}
+
+				example := utils.NewExampleFromProfilerChunk(
+					result.Chunk.ProjectID,
+					result.Chunk.ProfilerID,
+					result.Chunk.ID,
+					result.TransactionID,
+					&threadID,
+					result.Start,
+					result.End,
+				)
+				annotate := annotateWithProfileExample(example)
+
 				addCallTreeToFlamegraph(&flamegraphTree, callTree, annotate)
-			}
-			// if metrics aggregator is not null, while we're at it,
-			// compute the metrics as well
-			if ma != nil {
-				functions := metrics.CapAndFilterFunctions(metrics.ExtractFunctionsFromCallTrees(result.CallTrees), int(ma.MaxUniqueFunctions), true)
-				ma.AddFunctions(functions, example)
+
+				// if metrics aggregator is not null, while we're at it,
+				// compute the metrics as well
+				if ma != nil {
+					functions := metrics.CapAndFilterFunctions(metrics.ExtractFunctionsFromCallTreesForThread(callTree), int(ma.MaxUniqueFunctions), true)
+					ma.AddFunctions(functions, example)
+				}
 			}
 			chunkProfileSpan.Finish()
 		} else {

diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go
@@ -111,6 +111,18 @@ func quantile(values []uint64, q float64) (uint64, error) {
 	return values[index], nil
 }
 
+func ExtractFunctionsFromCallTreesForThread(
+	callTreesForThread []*nodetree.Node,
+) []nodetree.CallTreeFunction {
+	functions := make(map[uint32]nodetree.CallTreeFunction, 0)
+
+	for _, callTree := range callTreesForThread {
+		callTree.CollectFunctions(functions)
+	}
+
+	return mergeAndSortFunctions(functions)
+}
+
 func ExtractFunctionsFromCallTrees[T comparable](
 	callTrees map[T][]*nodetree.Node,
 ) []nodetree.CallTreeFunction {
@@ -122,6 +134,12 @@ func ExtractFunctionsFromCallTrees[T comparable](
 		}
 	}
 
+	return mergeAndSortFunctions(functions)
+}
+
+func mergeAndSortFunctions(
+	functions map[uint32]nodetree.CallTreeFunction,
+) []nodetree.CallTreeFunction {
 	functionsList := make([]nodetree.CallTreeFunction, 0, len(functions))
 	for _, function := range functions {
 		if function.SampleCount <= 1 {