Fix otel span being created as a transaction

otel/propagator_test.go

@@ -41,9 +41,9 @@ func createTransactionAndMaybeSpan(transactionContext transactionTestContext, wi
 		// we "swap" span IDs from the transaction and the child span.
 		transaction.SpanID = span.SpanID
 		span.SpanID = SpanIDFromHex(transactionContext.spanID)
-		sentrySpanMap.Set(trace.SpanID(span.SpanID), span)
+		sentrySpanMap.Set(trace.SpanID(span.SpanID), span, trace.SpanID{})
 	}
-	sentrySpanMap.Set(trace.SpanID(transaction.SpanID), transaction)
+	sentrySpanMap.Set(trace.SpanID(transaction.SpanID), transaction, trace.SpanID{})
 
 	otelContext := trace.SpanContextConfig{
 		TraceID:    otelTraceIDFromHex(transactionContext.traceID),

otel/span_map.go

@@ -7,37 +7,111 @@ import (
 	otelTrace "go.opentelemetry.io/otel/trace"
 )
 
+type spanInfo struct {
+	span     *sentry.Span
+	finished bool
+	children map[otelTrace.SpanID]struct{}
+	parentID otelTrace.SpanID
+}
+
 // SentrySpanMap is a mapping between OpenTelemetry spans and Sentry spans.
 // It helps Sentry span processor and propagator to keep track of unfinished
 // Sentry spans and to establish parent-child links between spans.
 type SentrySpanMap struct {
-	spanMap map[otelTrace.SpanID]*sentry.Span
+	spanMap map[otelTrace.SpanID]*spanInfo
 	mu      sync.RWMutex
 }
 
 func (ssm *SentrySpanMap) Get(otelSpandID otelTrace.SpanID) (*sentry.Span, bool) {
 	ssm.mu.RLock()
 	defer ssm.mu.RUnlock()
-	span, ok := ssm.spanMap[otelSpandID]
-	return span, ok
+	info, ok := ssm.spanMap[otelSpandID]
+	if !ok {
+		return nil, false
+	}
+	return info.span, true
 }
 
-func (ssm *SentrySpanMap) Set(otelSpandID otelTrace.SpanID, sentrySpan *sentry.Span) {
+func (ssm *SentrySpanMap) Set(otelSpandID otelTrace.SpanID, sentrySpan *sentry.Span, parentID otelTrace.SpanID) {
 	ssm.mu.Lock()
 	defer ssm.mu.Unlock()
-	ssm.spanMap[otelSpandID] = sentrySpan
+
+	info := &spanInfo{
+		span:     sentrySpan,
+		finished: false,
+		children: make(map[otelTrace.SpanID]struct{}),
+		parentID: parentID,
+	}
+	ssm.spanMap[otelSpandID] = info
+
+	if parentID != (otelTrace.SpanID{}) {
+		if parentInfo, ok := ssm.spanMap[parentID]; ok {
+			parentInfo.children[otelSpandID] = struct{}{}
+		}
+	}
 }
 
-func (ssm *SentrySpanMap) Delete(otelSpandID otelTrace.SpanID) {
+func (ssm *SentrySpanMap) MarkFinished(otelSpandID otelTrace.SpanID) {
 	ssm.mu.Lock()
 	defer ssm.mu.Unlock()
-	delete(ssm.spanMap, otelSpandID)
+
+	info, ok := ssm.spanMap[otelSpandID]
+	if !ok {
+		return
+	}
+
+	info.finished = true
+	ssm.tryCleanupSpan(otelSpandID)
+}
+
+// tryCleanupSpan deletes a parent and all children only if the whole subtree is marked finished.
+// Must be called with lock held.
+func (ssm *SentrySpanMap) tryCleanupSpan(spanID otelTrace.SpanID) {
+	info, ok := ssm.spanMap[spanID]
+	if !ok || !info.finished {
+		return
+	}
+
+	if !info.span.IsTransaction() {
+		parentID := info.parentID
+		if parentID != (otelTrace.SpanID{}) {
+			if parentInfo, parentExists := ssm.spanMap[parentID]; parentExists && !parentInfo.finished {
+				return
+			}
+		}
+	}
+
+	// We need to have a lookup first to see if every child is marked as finished to actually cleanup everything.
+	// There probably is a better way to do this
+	for childID := range info.children {
+		if childInfo, exists := ssm.spanMap[childID]; exists && !childInfo.finished {
+			return
+		}
+	}
+
+	parentID := info.parentID
+	if parentID != (otelTrace.SpanID{}) {
+		if parentInfo, ok := ssm.spanMap[parentID]; ok {
+			delete(parentInfo.children, spanID)
+		}
+	}
+
+	for childID := range info.children {
+		if childInfo, exists := ssm.spanMap[childID]; exists && childInfo.finished {
+			ssm.tryCleanupSpan(childID)
+		}
+	}
+
+	delete(ssm.spanMap, spanID)
+	if parentID != (otelTrace.SpanID{}) {
+		ssm.tryCleanupSpan(parentID)
+	}
 }
 
 func (ssm *SentrySpanMap) Clear() {
 	ssm.mu.Lock()
 	defer ssm.mu.Unlock()
-	ssm.spanMap = make(map[otelTrace.SpanID]*sentry.Span)
+	ssm.spanMap = make(map[otelTrace.SpanID]*spanInfo)
 }
 
 func (ssm *SentrySpanMap) Len() int {
@@ -46,4 +120,4 @@ func (ssm *SentrySpanMap) Len() int {
 	return len(ssm.spanMap)
 }
 
-var sentrySpanMap = SentrySpanMap{spanMap: make(map[otelTrace.SpanID]*sentry.Span)}
+var sentrySpanMap = SentrySpanMap{spanMap: make(map[otelTrace.SpanID]*spanInfo)}

otel/span_processor.go

@@ -21,7 +21,7 @@ func NewSentrySpanProcessor() otelSdkTrace.SpanProcessor {
 		return sentrySpanProcessorInstance
 	}
 	sentry.AddGlobalEventProcessor(linkTraceContextToErrorEvent)
-	sentrySpanProcessorInstance := &sentrySpanProcessor{}
+	sentrySpanProcessorInstance = &sentrySpanProcessor{}
 	return sentrySpanProcessorInstance
 }
 
@@ -42,7 +42,7 @@ func (ssp *sentrySpanProcessor) OnStart(parent context.Context, s otelSdkTrace.R
 		span.SpanID = sentry.SpanID(otelSpanID)
 		span.StartTime = s.StartTime()
 
-		sentrySpanMap.Set(otelSpanID, span)
+		sentrySpanMap.Set(otelSpanID, span, otelParentSpanID)
 	} else {
 		traceParentContext := getTraceParentContext(parent)
 		transaction := sentry.StartTransaction(
@@ -58,7 +58,7 @@ func (ssp *sentrySpanProcessor) OnStart(parent context.Context, s otelSdkTrace.R
 			transaction.SetDynamicSamplingContext(dynamicSamplingContext)
 		}
 
-		sentrySpanMap.Set(otelSpanID, transaction)
+		sentrySpanMap.Set(otelSpanID, transaction, otelParentSpanID)
 	}
 }
 
@@ -71,7 +71,7 @@ func (ssp *sentrySpanProcessor) OnEnd(s otelSdkTrace.ReadOnlySpan) {
 	}
 
 	if utils.IsSentryRequestSpan(sentrySpan.Context(), s) {
-		sentrySpanMap.Delete(otelSpanId)
+		sentrySpanMap.MarkFinished(otelSpanId)
 		return
 	}
 
@@ -84,7 +84,7 @@ func (ssp *sentrySpanProcessor) OnEnd(s otelSdkTrace.ReadOnlySpan) {
 	sentrySpan.EndTime = s.EndTime()
 	sentrySpan.Finish()
 
-	sentrySpanMap.Delete(otelSpanId)
+	sentrySpanMap.MarkFinished(otelSpanId)
 }
 
 // https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/sdk.md#shutdown-1

otel/span_processor_test.go

@@ -75,7 +75,7 @@ func TestSpanProcessorShutdown(t *testing.T) {
 
 	assertEqual(t, sentrySpanMap.Len(), 1)
 
-	spanProcessor.Shutdown(ctx)
+	_ = spanProcessor.Shutdown(ctx)
 
 	// The span map should be empty
 	assertEqual(t, sentrySpanMap.Len(), 0)
@@ -399,3 +399,59 @@ func TestParseSpanAttributesHttpServer(t *testing.T) {
 	assertEqual(t, sentrySpan.Op, "http.server")
 	assertEqual(t, sentrySpan.Source, sentry.TransactionSource(""))
 }
+
+func TestSpanBecomesChildOfFinishedSpan(t *testing.T) {
+	_, _, tracer := setupSpanProcessorTest()
+	ctx, otelRootSpan := tracer.Start(
+		emptyContextWithSentry(),
+		"rootSpan",
+	)
+	sentryTransaction, _ := sentrySpanMap.Get(otelRootSpan.SpanContext().SpanID())
+
+	ctx, childSpan1 := tracer.Start(
+		ctx,
+		"span name 1",
+	)
+	sentrySpan1, _ := sentrySpanMap.Get(childSpan1.SpanContext().SpanID())
+	childSpan1.End()
+
+	_, childSpan2 := tracer.Start(
+		ctx,
+		"span name 2",
+	)
+	sentrySpan2, _ := sentrySpanMap.Get(childSpan2.SpanContext().SpanID())
+	childSpan2.End()
+
+	otelRootSpan.End()
+
+	assertEqual(t, sentryTransaction.IsTransaction(), true)
+	assertEqual(t, sentrySpan1.IsTransaction(), false)
+	assertEqual(t, sentrySpan1.ParentSpanID, sentryTransaction.SpanID)
+	assertEqual(t, sentrySpan2.IsTransaction(), false)
+	assertEqual(t, sentrySpan2.ParentSpanID, sentrySpan1.SpanID)
+}
+
+func TestSpanWithFinishedParentShouldBeDeleted(t *testing.T) {
+	_, _, tracer := setupSpanProcessorTest()
+
+	ctx, parent := tracer.Start(context.Background(), "parent")
+	parentSpanID := parent.SpanContext().SpanID()
+	_, child := tracer.Start(ctx, "child")
+	childSpanID := child.SpanContext().SpanID()
+
+	_, parentExists := sentrySpanMap.Get(parentSpanID)
+	_, childExists := sentrySpanMap.Get(childSpanID)
+	assertEqual(t, parentExists, true)
+	assertEqual(t, childExists, true)
+
+	parent.End()
+	_, parentExists = sentrySpanMap.Get(parentSpanID)
+	assertEqual(t, parentExists, true)
+	_, childExists = sentrySpanMap.Get(childSpanID)
+	assertEqual(t, childExists, true)
+
+	child.End()
+	_, childExists = sentrySpanMap.Get(childSpanID)
+	assertEqual(t, childExists, false)
+	assertEqual(t, sentrySpanMap.Len(), 0)
+}