From 85759da360158d800874a7046186d9165fb03f2e Mon Sep 17 00:00:00 2001
From: Todd Treece <360020+toddtreece@users.noreply.github.com>
Date: Fri, 6 Sep 2024 13:02:17 -0400
Subject: [PATCH] Tracing: Add Start helper (#93052)

---
 pkg/infra/tracing/tracing.go      |  7 ++++++
 pkg/infra/tracing/tracing_test.go | 36 +++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/pkg/infra/tracing/tracing.go b/pkg/infra/tracing/tracing.go
index 624f892214c..4c6bd6545a3 100644
--- a/pkg/infra/tracing/tracing.go
+++ b/pkg/infra/tracing/tracing.go
@@ -399,3 +399,10 @@ func Errorf(span trace.Span, format string, args ...any) error {
 	err := fmt.Errorf(format, args...)
 	return Error(span, err)
 }
+
+var instrumentationScope = "github.com/grafana/grafana/pkg/infra/tracing"
+
+// Start only creates an OpenTelemetry span if the incoming context already includes a span.
+func Start(ctx context.Context, name string, attributes ...attribute.KeyValue) (context.Context, trace.Span) {
+	return trace.SpanFromContext(ctx).TracerProvider().Tracer(instrumentationScope).Start(ctx, name, trace.WithAttributes(attributes...))
+}
diff --git a/pkg/infra/tracing/tracing_test.go b/pkg/infra/tracing/tracing_test.go
index b52b772d424..08880b46ca7 100644
--- a/pkg/infra/tracing/tracing_test.go
+++ b/pkg/infra/tracing/tracing_test.go
@@ -1,10 +1,13 @@
 package tracing
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/trace"
 )
 
 func TestInitSampler(t *testing.T) {
@@ -41,3 +44,36 @@ func TestInitSampler(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "RateLimitingSampler{100.25}", sampler.Description())
 }
+
+func TestStart(t *testing.T) {
+	name := "test-span"
+	attributes := []attribute.KeyValue{
+		attribute.String("test1", "1"),
+		attribute.Int("test2", 2),
+	}
+
+	t.Run("should return noop span if there is not currently a span in context", func(t *testing.T) {
+		ctx := context.Background()
+		_, span := Start(ctx, name, attributes...)
+		defer span.End()
+
+		require.NotNil(t, span)
+		require.False(t, span.SpanContext().IsValid())
+	})
+
+	t.Run("should return a span with a valid span context if there is currently a span in context", func(t *testing.T) {
+		spanCtx := trace.NewSpanContext(trace.SpanContextConfig{
+			TraceID:    trace.TraceID{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16},
+			SpanID:     trace.SpanID{1, 2, 3, 4, 5, 6, 7, 8},
+			TraceFlags: trace.FlagsSampled,
+		})
+
+		ctx := trace.ContextWithSpanContext(context.Background(), spanCtx)
+		_, childSpan := Start(ctx, name, attributes...)
+		defer childSpan.End()
+
+		require.NotNil(t, childSpan)
+		require.Equal(t, spanCtx.TraceID(), childSpan.SpanContext().TraceID())
+		require.True(t, childSpan.SpanContext().IsValid())
+	})
+}