From 84a665e5493e824219f91afef7bcc89f08d9407a Mon Sep 17 00:00:00 2001
From: Tony Xiao <txiao@sentry.io>
Date: Thu, 24 Oct 2024 16:03:14 -0400
Subject: [PATCH 1/2] fix(flamegraphs): Annoate functions to the right thread

This was using the given thread id for when annotating all samples in the chunk.
We need to correctly assign the correct thread for each function.
---
 internal/flamegraph/flamegraph.go | 38 ++++++++++++++++---------------
 internal/metrics/metrics.go       | 18 +++++++++++++++
 2 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/internal/flamegraph/flamegraph.go b/internal/flamegraph/flamegraph.go
index f60778a..cb445b5 100644
--- a/internal/flamegraph/flamegraph.go
+++ b/internal/flamegraph/flamegraph.go
@@ -559,18 +559,7 @@ 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,
@@ -578,13 +567,26 @@ func GetFlamegraphFromCandidates(
 					}
 					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
index eb6cb2f..fe3bd3c 100644
--- 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 {

From 040c1294a3e0370acf3e49676f3422eeddc27e93 Mon Sep 17 00:00:00 2001
From: Tony Xiao <txiao@sentry.io>
Date: Thu, 24 Oct 2024 16:05:58 -0400
Subject: [PATCH 2/2] update changelog

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4e02dd9..dee9e84 100644
--- 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**: