From 9c502e558d44098073b1fe950bf8b3851fd873ba Mon Sep 17 00:00:00 2001
From: Tony Xiao <txiao@sentry.io>
Date: Wed, 30 Oct 2024 14:32:03 -0400
Subject: [PATCH] instrument span

---
 internal/flamegraph/flamegraph.go      | 25 +++++++++++++++++++++----
 internal/flamegraph/flamegraph_test.go | 11 ++++++-----
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/internal/flamegraph/flamegraph.go b/internal/flamegraph/flamegraph.go
index f8a5de8..f2eee36 100644
--- a/internal/flamegraph/flamegraph.go
+++ b/internal/flamegraph/flamegraph.go
@@ -136,7 +136,7 @@ func GetFlamegraphFromProfiles(
 		countProfAggregated++
 	}
 
-	sp := toSpeedscope(flamegraphTree, 1000, projectID)
+	sp := toSpeedscope(ctx, flamegraphTree, 1000, projectID)
 	hub.Scope().SetTag("processed_profiles", strconv.Itoa(countProfAggregated))
 	return sp, nil
 }
@@ -226,6 +226,10 @@ type (
 		profiles          []utils.ExampleMetadata
 		endValue          uint64
 		maxSamples        int
+		// The total number of samples that were added to the flamegraph
+		// including the ones that were dropped due to them exceeding
+		// the max samples limit.
+		totalSamples int
 	}
 
 	flamegraphSample struct {
@@ -307,7 +311,16 @@ func (f *flamegraph) Pop() any {
 	return sample
 }
 
-func toSpeedscope(trees []*nodetree.Node, maxSamples int, projectID uint64) speedscope.Output {
+func toSpeedscope(
+	ctx context.Context,
+	trees []*nodetree.Node,
+	maxSamples int,
+	projectID uint64,
+) speedscope.Output {
+	s := sentry.StartSpan(ctx, "processing")
+	s.Description = "generating speedscope"
+	defer s.Finish()
+
 	fd := &flamegraph{
 		frames:           make([]speedscope.Frame, 0),
 		framesIndex:      make(map[string]int),
@@ -322,6 +335,9 @@ func toSpeedscope(trees []*nodetree.Node, maxSamples int, projectID uint64) spee
 		fd.visitCalltree(tree, &stack)
 	}
 
+	s.SetData("total_samples", fd.totalSamples)
+	s.SetData("final_samples", fd.Len())
+
 	aggProfiles := make([]interface{}, 1)
 	aggProfiles[0] = speedscope.SampledProfile{
 		Samples:           fd.samples,
@@ -422,6 +438,7 @@ func (f *flamegraph) addSample(
 	profileIDs map[string]struct{},
 	profiles map[utils.ExampleMetadata]struct{},
 ) {
+	f.totalSamples += 1
 	cp := make([]int, len(*stack))
 	copy(cp, *stack)
 	heap.Push(f, flamegraphSample{
@@ -541,7 +558,7 @@ func GetFlamegraphFromChunks(
 		countChunksAggregated++
 	}
 
-	sp := toSpeedscope(flamegraphTree, 1000, projectID)
+	sp := toSpeedscope(ctx, flamegraphTree, 1000, projectID)
 	if hub != nil {
 		hub.Scope().SetTag("processed_chunks", strconv.Itoa(countChunksAggregated))
 	}
@@ -682,7 +699,7 @@ func GetFlamegraphFromCandidates(
 	serializeSpan := span.StartChild("serialize")
 	defer serializeSpan.Finish()
 
-	sp := toSpeedscope(flamegraphTree, 1000, 0)
+	sp := toSpeedscope(ctx, flamegraphTree, 1000, 0)
 	if ma != nil {
 		fm := ma.ToMetrics()
 		sp.Metrics = &fm
diff --git a/internal/flamegraph/flamegraph_test.go b/internal/flamegraph/flamegraph_test.go
index b4fa003..2118bf2 100644
--- a/internal/flamegraph/flamegraph_test.go
+++ b/internal/flamegraph/flamegraph_test.go
@@ -1,6 +1,7 @@
 package flamegraph
 
 import (
+	"context"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -145,9 +146,9 @@ func TestFlamegraphAggregation(t *testing.T) {
 						SamplesExamples:   [][]int{{}, {}, {}, {}},
 						Type:              "sampled",
 						Unit:              "count",
-						Weights:           []uint64{2, 1, 1, 1},
-						SampleCounts:      []uint64{2, 1, 1, 1},
-						SampleDurationsNs: []uint64{20, 10, 10, 10},
+						Weights:           []uint64{1, 1, 1, 2},
+						SampleCounts:      []uint64{1, 1, 1, 2},
+						SampleDurationsNs: []uint64{10, 10, 10, 20},
 					},
 				},
 				Shared: speedscope.SharedData{
@@ -199,7 +200,7 @@ func TestFlamegraphAggregation(t *testing.T) {
 				addCallTreeToFlamegraph(&ft, callTrees[0], annotateWithProfileID(p.ID()))
 			}
 
-			if diff := testutil.Diff(toSpeedscope(ft, 10, 99), test.output, options); diff != "" {
+			if diff := testutil.Diff(toSpeedscope(context.TODO(), ft, 10, 99), test.output, options); diff != "" {
 				t.Fatalf("Result mismatch: got - want +\n%s", diff)
 			}
 		})
@@ -335,7 +336,7 @@ func TestAnnotatingWithExamples(t *testing.T) {
 			for _, example := range test.examples {
 				addCallTreeToFlamegraph(&ft, test.callTrees, annotateWithProfileExample(example))
 			}
-			if diff := testutil.Diff(toSpeedscope(ft, 10, 99), test.output, options); diff != "" {
+			if diff := testutil.Diff(toSpeedscope(context.TODO(), ft, 10, 99), test.output, options); diff != "" {
 				t.Fatalf("Result mismatch: got - want +\n%s", diff)
 			}
 		})