From 5407f87c8c1dae033b1b7c7bf513453ff4d78360 Mon Sep 17 00:00:00 2001
From: Viktar Dareuski <mr.darevski@gmail.com>
Date: Sun, 29 Sep 2024 17:04:51 +0300
Subject: [PATCH] Fix scope management in StartSpan and doFinish

---
 tracing.go      | 11 +++++++-
 tracing_test.go | 68 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/tracing.go b/tracing.go
index 0f5ade2e..dae0c635 100644
--- a/tracing.go
+++ b/tracing.go
@@ -195,6 +195,10 @@ func StartSpan(ctx context.Context, operation string, options ...SpanOption) *Sp
 	clientOptions := span.clientOptions()
 	if clientOptions.EnableTracing {
 		hub := hubFromContext(ctx)
+		if !span.IsTransaction() {
+			// Push a new scope to stack for non transaction span
+			hub.PushScope()
+		}
 		hub.Scope().SetSpan(&span)
 	}
 
@@ -355,6 +359,12 @@ func (s *Span) doFinish() {
 		s.EndTime = monotonicTimeSince(s.StartTime)
 	}
 
+	hub := hubFromContext(s.ctx)
+	if !s.IsTransaction() {
+		// Referenced to StartSpan function that pushes current non-transaction span to scope stack
+		defer hub.PopScope()
+	}
+
 	if !s.Sampled.Bool() {
 		return
 	}
@@ -370,7 +380,6 @@ func (s *Span) doFinish() {
 	// TODO(tracing): add breadcrumbs
 	// (see https://github.com/getsentry/sentry-python/blob/f6f3525f8812f609/sentry_sdk/tracing.py#L372)
 
-	hub := hubFromContext(s.ctx)
 	hub.CaptureEvent(event)
 }
 
diff --git a/tracing_test.go b/tracing_test.go
index 795c33d3..519b6f38 100644
--- a/tracing_test.go
+++ b/tracing_test.go
@@ -1019,3 +1019,71 @@ func TestSpanFinishConcurrentlyWithoutRaces(_ *testing.T) {
 
 	time.Sleep(50 * time.Millisecond)
 }
+
+func TestSpanScopeManagement(t *testing.T) {
+	// Initialize a test hub and client
+	transport := &TransportMock{}
+	client, err := NewClient(ClientOptions{
+		EnableTracing:    true,
+		TracesSampleRate: 1.0,
+		Transport:        transport,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	hub := NewHub(client, NewScope())
+
+	// Set the hub on the context
+	ctx := context.Background()
+	ctx = SetHubOnContext(ctx, hub)
+
+	// Start a parent span (transaction)
+	transaction := StartTransaction(ctx, "parent-operation")
+	defer transaction.Finish()
+
+	// Start a child span
+	childSpan := StartSpan(transaction.Context(), "child-operation")
+	// Finish the child span
+	defer childSpan.Finish()
+
+	subChildSpan := StartSpan(childSpan.Context(), "sub_child-operation")
+	subChildSpan.Finish()
+
+	// Capture an event after finishing the child span
+	// This event should be associated with the first child span
+	hub.CaptureMessage("Test event")
+
+	// Flush to ensure the event is sent
+	transport.Flush(time.Second)
+
+	// Verify that the event has the correct trace data
+	events := transport.Events()
+	if len(events) != 1 {
+		t.Fatalf("expected 2 event, got %d", len(events))
+	}
+	event := events[0]
+
+	// Extract the trace context from the event
+	traceCtx, ok := event.Contexts["trace"]
+	if !ok {
+		t.Fatalf("event does not have a trace context")
+	}
+
+	// Extract TraceID and SpanID from the trace context
+	traceID, ok := traceCtx["trace_id"].(TraceID)
+	if !ok {
+		t.Fatalf("trace_id not found")
+	}
+	spanID, ok := traceCtx["span_id"].(SpanID)
+	if !ok {
+		t.Fatalf("span_id not found")
+	}
+
+	// Verify that the IDs match the first child span IDs
+	if traceID != childSpan.TraceID {
+		t.Errorf("expected TraceID %s, got %s", transaction.TraceID, traceID)
+	}
+	if spanID != childSpan.SpanID {
+		t.Errorf("expected SpanID %s, got %s", transaction.SpanID, spanID)
+	}
+}