From 5580a9a9e28ef53b91e950f4e10222b05d861d69 Mon Sep 17 00:00:00 2001
From: Damien Mathieu <42@dmathieu.com>
Date: Fri, 22 Sep 2023 19:24:23 +0200
Subject: [PATCH] apmotel: disallow closing spans twice (#1512)

---
 module/apmotel/span.go      |  5 ++++
 module/apmotel/span_test.go | 50 +++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/module/apmotel/span.go b/module/apmotel/span.go
index 3e9f2f981..7e358b5d2 100644
--- a/module/apmotel/span.go
+++ b/module/apmotel/span.go
@@ -51,6 +51,7 @@ type span struct {
 
 	provider *tracerProvider
 
+	ended       bool
 	startTime   time.Time
 	attributes  []attribute.KeyValue
 	events      []event
@@ -65,6 +66,9 @@ type span struct {
 func (s *span) End(options ...trace.SpanEndOption) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	if s.ended {
+		return
+	}
 
 	config := trace.NewSpanEndConfig(options...)
 
@@ -90,6 +94,7 @@ func (s *span) End(options ...trace.SpanEndOption) {
 	for iter := s.provider.resource.Iter(); iter.Next(); {
 		s.attributes = append(s.attributes, iter.Attribute())
 	}
+	s.ended = true
 
 	if s.span != nil {
 		s.setSpanAttributes()
diff --git a/module/apmotel/span_test.go b/module/apmotel/span_test.go
index b2ed60d8d..47c3a3749 100644
--- a/module/apmotel/span_test.go
+++ b/module/apmotel/span_test.go
@@ -579,6 +579,56 @@ func TestSpanEnd(t *testing.T) {
 	}
 }
 
+func TestSpanEndTwice(t *testing.T) {
+	for _, tt := range []struct {
+		name    string
+		getSpan func(context.Context, trace.Tracer) trace.Span
+
+		expectedSpansCount        int
+		expectedTransactionsCount int
+	}{
+		{
+			name: "with a transaction",
+			getSpan: func(ctx context.Context, tracer trace.Tracer) trace.Span {
+				ctx, s := tracer.Start(ctx, "transaction")
+				return s
+			},
+
+			expectedTransactionsCount: 1,
+		},
+		{
+			name: "with a span",
+			getSpan: func(ctx context.Context, tracer trace.Tracer) trace.Span {
+				ctx, _ = tracer.Start(ctx, "transaction")
+				ctx, s := tracer.Start(ctx, "span")
+				return s
+			},
+
+			expectedSpansCount: 1,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			apmTracer, recorder := transporttest.NewRecorderTracer()
+			tp, err := NewTracerProvider(
+				WithAPMTracer(apmTracer),
+			)
+			assert.NoError(t, err)
+			tracer := newTracer(tp.(*tracerProvider))
+
+			ctx := context.Background()
+			s := tt.getSpan(ctx, tracer)
+
+			s.End()
+			assert.NotPanics(t, func() { s.End() })
+
+			apmTracer.Flush(nil)
+			payloads := recorder.Payloads()
+			assert.Equal(t, tt.expectedSpansCount, len(payloads.Spans))
+			assert.Equal(t, tt.expectedTransactionsCount, len(payloads.Transactions))
+		})
+	}
+}
+
 func TestSpanAddEvent(t *testing.T) {
 	tp, err := NewTracerProvider()
 	assert.NoError(t, err)