tracing: add ExtractTraceSpanID

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

tracing/tracing.go

@@ -55,16 +55,31 @@ 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 sctx, ok := extractJaegerContext(ctx); ok {
+		return sctx.TraceID().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 sctx, ok := extractJaegerContext(ctx); ok {
+		return sctx.TraceID().String(), sctx.SpanID().String(), true
+	}
+	return "", "", false
+}
+
+func extractJaegerContext(ctx context.Context) (jsp jaeger.SpanContext, 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
+	success = true
+	return
 }
 
 // ExtractSampledTraceID works like ExtractTraceID but the returned bool is only

tracing/tracing_test.go

@@ -0,0 +1,47 @@
+// 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)
+		})
+	}
+}