diff --git a/contrib/datadog/tracing/interceptor.go b/contrib/datadog/tracing/interceptor.go index 954c79598..4faadde58 100644 --- a/contrib/datadog/tracing/interceptor.go +++ b/contrib/datadog/tracing/interceptor.go @@ -143,6 +143,20 @@ func (t *tracerImpl) ContextWithSpan(ctx context.Context, span interceptor.Trace return tracer.ContextWithSpan(ctx, span.(*tracerSpan).Span) } +// SpanFromWorkflowContext extracts the DataDog Span object from the workflow context. +func SpanFromWorkflowContext(ctx workflow.Context) (ddtrace.Span, bool) { + val := ctx.Value(activeSpanContextKey) + if val == nil { + return tracer.SpanFromContext(nil) + } + + if span, ok := val.(*tracerSpan); ok { + return span.Span, true + } + + return tracer.SpanFromContext(nil) +} + func genSpanID(idempotencyKey string) uint64 { h := fnv.New64() // Write() always writes all bytes and never fails; the count and error result are for implementing io.Writer. diff --git a/contrib/datadog/tracing/interceptor_test.go b/contrib/datadog/tracing/interceptor_test.go index caeff15a5..1d7943a01 100644 --- a/contrib/datadog/tracing/interceptor_test.go +++ b/contrib/datadog/tracing/interceptor_test.go @@ -22,6 +22,7 @@ package tracing import ( + "errors" "strings" "testing" @@ -31,6 +32,9 @@ import ( "go.temporal.io/sdk/interceptor" "go.temporal.io/sdk/internal/interceptortest" + "go.temporal.io/sdk/testsuite" + "go.temporal.io/sdk/worker" + "go.temporal.io/sdk/workflow" ) type testTracer struct { @@ -133,7 +137,6 @@ func Test_OnFinishOption(t *testing.T) { Tracer: impl, mt: mt, } - interceptortest.RunTestWorkflowWithError(t, trc) spans := trc.FinishedSpans() @@ -141,3 +144,41 @@ func Test_OnFinishOption(t *testing.T) { require.Len(t, spans, 1) require.Equal(t, "temporal.RunWorkflow", spans[0].Name) } + +func setCustomSpanTagWorkflow(ctx workflow.Context) error { + span, ok := SpanFromWorkflowContext(ctx) + + if !ok { + return errors.New("Did not find span in workflow context") + } + + span.SetTag("testTag", "testValue") + return nil +} + +func Test_SpanFromWorkflowContext(t *testing.T) { + // Start the mock tracer. + mt := mocktracer.Start() + defer mt.Stop() + + var suite testsuite.WorkflowTestSuite + env := suite.NewTestWorkflowEnvironment() + env.RegisterWorkflow(setCustomSpanTagWorkflow) + + impl := NewTracer(TracerOptions{}) + testTracer := testTracer{ + Tracer: impl, + mt: mt, + } + + // Set tracer interceptor + env.SetWorkerOptions(worker.Options{ + Interceptors: []interceptor.WorkerInterceptor{interceptor.NewTracingInterceptor(testTracer)}, + }) + + env.ExecuteWorkflow(setCustomSpanTagWorkflow) + + require.True(t, env.IsWorkflowCompleted()) + testSpan := mt.FinishedSpans()[0] + require.Equal(t, "testValue", testSpan.Tag("testTag")) +}