tracing: add ExtractTraceSpanID (#550)

* tracing: add ExtractTraceSpanID function

Signed-off-by: György Krajcsovits <[email protected]>

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,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

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)
+		})
+	}
+}