From 275a6e5a32b924e82dffaa4fef1577065d03b2c3 Mon Sep 17 00:00:00 2001
From: Matt Quinn <matt.quinn@sentry.io>
Date: Tue, 22 Sep 2026 16:04:34 -0400
Subject: [PATCH] ref(flamegraph): More robust concurrency structure

`GetFlamegraphFromCandidates`'s design:

- create an unbuffered `results` channel
- spin up a goroutine to dispatch one job per candidate (possibly a transaction
  profile or a continuous profile) to `jobs` (actually the global `readJobs`
  channel). Each job includes a reference to `results` so the result can be
  written back.
- read from the `results` channel once per candidate and add its data to the
  flamegraph
- return the flamegraph

(This relies on an assumption that every job submitted to `readJobs` will write
its result to `results`, which is handled elsewhere).

There are two minor gaps:

1. Although the `unexpected result from storage` error is never expected to
  happen, it's an error and not a panic. But, it may result in a panic anyway: it
  will cause a `return`, which will trigger the deferred `close(results)`, which
  will panic the next time a `readJobs` worker finishes a job and tries to write
  to it. Even if we fixed the panic, we would still need to drain `results` to
  avoid blocking `readJobs` workers while they write their results.
2. The dispatch goroutine may not finish before `GetFlamegraphFromCandidates`
  finishes, for example when there are no candidates.

To fix this issue and clarify the flow:

- Completely drain `results` even if we get an unexpected error, so that we
  can't block the `readJobs` workers. We still return the error.
- Don't `close(results)`. It isn't necessary for signalling (since we always
  read one result per submitted candidate) and closing it risks a panic if any
  worker writes after the receiver returns.
- Wait for the dispatch goroutine to finish, ensuring both the goroutine and
  its dispatch span finish before `GetFlamegraphFromCandidates` returns.
---
 CHANGELOG.md                      |  1 +
 internal/flamegraph/flamegraph.go | 21 ++++++++++++++++-----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 20212849..9bcabfee 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@
 **Internal**:
 
 - Fix flamegraph endpoint span nesting. ([#677](https://github.com/getsentry/vroom/pull/677))
+- More robust concurrency structure for flamegraph generation. ([#678](https://github.com/getsentry/vroom/pull/678))
 
 ## 26.9.0
 
diff --git a/internal/flamegraph/flamegraph.go b/internal/flamegraph/flamegraph.go
index 809117f0..bc789a9f 100644
--- a/internal/flamegraph/flamegraph.go
+++ b/internal/flamegraph/flamegraph.go
@@ -371,13 +371,16 @@ func GetFlamegraphFromCandidates(
 	hub := sentry.GetHubFromContext(ctx)
 
 	results := make(chan storageutil.ReadJobResult)
-	defer close(results)
+	dispatchDone := make(chan struct{})
 
 	go func() {
 		dispatchSpan := sentry.StartSpan(ctx, "dispatch candidates")
 		dispatchSpan.SetData("transaction_candidates", len(transactionProfileCandidates))
 		dispatchSpan.SetData("continuous_candidates", len(continuousProfileCandidates))
-		defer dispatchSpan.Finish()
+		defer func() {
+			dispatchSpan.Finish()
+			close(dispatchDone)
+		}()
 
 		for _, candidate := range transactionProfileCandidates {
 			jobs <- profile.CallTreesReadJob{
@@ -412,10 +415,15 @@ func GetFlamegraphFromCandidates(
 	flamegraphSpan := sentry.StartSpan(ctx, "processing candidates")
 
 	numCandidates := len(transactionProfileCandidates) + len(continuousProfileCandidates)
+	var processingErr error
 
 	for i := 0; i < numCandidates; i++ {
 		res := <-results
 
+		if processingErr != nil {
+			continue
+		}
+
 		err := res.Error()
 		if err != nil {
 			if errors.Is(err, storageutil.ErrObjectNotFound) {
@@ -495,14 +503,17 @@ func GetFlamegraphFromCandidates(
 			}
 			chunkProfileSpan.Finish()
 		} else {
-			// This should never happen
-			flamegraphSpan.Finish()
-			return speedscope.Output{}, errors.New("unexpected result from storage")
+			processingErr = errors.New("unexpected result from storage")
 		}
 	}
 
+	<-dispatchDone
 	flamegraphSpan.Finish()
 
+	if processingErr != nil {
+		return speedscope.Output{}, processingErr
+	}
+
 	serializeSpan := sentry.StartSpan(ctx, "serialize")
 	defer serializeSpan.Finish()