From 89e6fcf4cc2d48d2d2990e1639a24cdfb557b89c Mon Sep 17 00:00:00 2001
From: George Krajcsovits <krajorama@users.noreply.github.com>
Date: Wed, 17 Jul 2024 16:40:23 +0200
Subject: [PATCH] tracing: add ExtractTraceSpanID (#550)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* tracing: add ExtractTraceSpanID function

Signed-off-by: György Krajcsovits <gyorgy.krajcsovits@grafana.com>
---
 CHANGELOG.md            |  1 +
 tracing/tracing.go      | 24 +++++++++++++++----
 tracing/tracing_test.go | 51 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 5 deletions(-)
 create mode 100644 tracing/tracing_test.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 41bfcda8e..ec010f23d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -214,6 +214,7 @@
 * [ENHANCEMENT] SpanProfiler: do less work on unsampled traces. #528
 * [ENHANCEMENT] Log Middleware: if the trace is not sampled, log its ID as `trace_id_unsampled` instead of `trace_id`. #529
 * [EHNANCEMENT] httpgrpc: httpgrpc Server can now use error message from special HTTP header when converting HTTP response to an error. This is useful when HTTP response body contains binary data that doesn't form valid utf-8 string, otherwise grpc would fail to marshal returned error. #531
+* [ENHANCEMENT] tracing: add ExtractTraceSpanID function.
 * [CHANGE] Backoff: added `Backoff.ErrCause()` which is like `Backoff.Err()` but returns the context cause if backoff is terminated because the context has been canceled. #538
 * [BUGFIX] spanlogger: Support multiple tenant IDs. #59
 * [BUGFIX] Memberlist: fixed corrupted packets when sending compound messages with more than 255 messages or messages bigger than 64KB. #85
diff --git a/tracing/tracing.go b/tracing/tracing.go
index 66b3a3cef..1882a081d 100644
--- a/tracing/tracing.go
+++ b/tracing/tracing.go
@@ -55,16 +55,30 @@ func NewFromEnv(serviceName string, options ...jaegercfg.Option) (io.Closer, err
 
 // ExtractTraceID extracts the trace id, if any from the context.
 func ExtractTraceID(ctx context.Context) (string, bool) {
+	if tid, _, ok := extractJaegerContext(ctx); ok {
+		return tid.String(), true
+	}
+	return "", false
+}
+
+// ExtractTraceSpanID extracts the trace id, span id if any from the context.
+func ExtractTraceSpanID(ctx context.Context) (string, string, bool) {
+	if tid, sid, ok := extractJaegerContext(ctx); ok {
+		return tid.String(), sid.String(), true
+	}
+	return "", "", false
+}
+
+func extractJaegerContext(ctx context.Context) (tid jaeger.TraceID, sid jaeger.SpanID, success bool) {
 	sp := opentracing.SpanFromContext(ctx)
 	if sp == nil {
-		return "", false
+		return
 	}
-	sctx, ok := sp.Context().(jaeger.SpanContext)
+	jsp, ok := sp.Context().(jaeger.SpanContext)
 	if !ok {
-		return "", false
+		return
 	}
-
-	return sctx.TraceID().String(), true
+	return jsp.TraceID(), jsp.SpanID(), true
 }
 
 // ExtractSampledTraceID works like ExtractTraceID but the returned bool is only
diff --git a/tracing/tracing_test.go b/tracing/tracing_test.go
new file mode 100644
index 000000000..c730366cc
--- /dev/null
+++ b/tracing/tracing_test.go
@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package tracing
+
+import (
+	"context"
+	"testing"
+
+	"github.com/opentracing/opentracing-go"
+	"github.com/stretchr/testify/require"
+	jaeger "github.com/uber/jaeger-client-go"
+)
+
+func TestExtractTraceSpanID(t *testing.T) {
+	spanCtx := jaeger.NewSpanContext(jaeger.TraceID{High: 1, Low: 2}, jaeger.SpanID(3), 0, true, nil)
+	tracer, closer := jaeger.NewTracer("test", jaeger.NewConstSampler(true), jaeger.NewNullReporter())
+	defer closer.Close()
+	span := tracer.StartSpan("test", opentracing.ChildOf(spanCtx))
+
+	testCases := map[string]struct {
+		ctx             context.Context
+		expectedOk      bool
+		expectedTraceID string
+		expectedSpanID  string
+	}{
+		"no trace": {
+			ctx:             context.Background(),
+			expectedOk:      false,
+			expectedTraceID: "",
+			expectedSpanID:  "",
+		},
+		"trace": {
+			ctx:             opentracing.ContextWithSpan(context.Background(), span),
+			expectedOk:      true,
+			expectedTraceID: "00000000000000010000000000000002",
+			expectedSpanID:  span.Context().(jaeger.SpanContext).SpanID().String(),
+		},
+	}
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
+			traceID, spanID, ok := ExtractTraceSpanID(tc.ctx)
+			require.Equal(t, tc.expectedOk, ok)
+			require.Equal(t, tc.expectedTraceID, traceID)
+			require.Equal(t, tc.expectedSpanID, spanID)
+
+			traceID, ok = ExtractTraceID(tc.ctx)
+			require.Equal(t, tc.expectedOk, ok)
+			require.Equal(t, tc.expectedTraceID, traceID)
+		})
+	}
+}