Improve trace propagation

Author: cschleiden

backend/redis/signal.go

@@ -18,9 +18,13 @@ func (rb *redisBackend) SignalWorkflow(ctx context.Context, instanceID string, e
 		return err
 	}
 
-	ctx = tracing.ExtractSpan(ctx, instanceState.Metadata)
+	ctx, err = (&tracing.TracingContextPropagator{}).Extract(ctx, instanceState.Metadata)
+	if err != nil {
+		rb.Logger().Error("extracting tracing context", log.ErrorKey, err)
+	}
+
 	a := event.Attributes.(*history.SignalReceivedAttributes)
-	_, span := rb.Tracer().Start(ctx, fmt.Sprintf("SignalWorkflow: %s", a.Name), trace.WithAttributes(
+	ctx, span := rb.Tracer().Start(ctx, fmt.Sprintf("SignalWorkflow: %s", a.Name), trace.WithAttributes(
 		attribute.String(log.InstanceIDKey, instanceID),
 		attribute.String(log.SignalNameKey, event.Attributes.(*history.SignalReceivedAttributes).Name),
 	))

backend/redis/workflow.go

@@ -275,7 +275,11 @@ func (rb *redisBackend) CompleteWorkflowTask(
 
 	if state == core.WorkflowInstanceStateFinished {
 		// Trace workflow completion
-		ctx = tracing.ExtractSpan(ctx, instanceState.Metadata)
+		ctx, err = (&tracing.TracingContextPropagator{}).Extract(ctx, instanceState.Metadata)
+		if err != nil {
+			rb.Logger().Error("extracting tracing context", log.ErrorKey, err)
+		}
+
 		_, span := rb.Tracer().Start(ctx, "WorkflowComplete",
 			trace.WithAttributes(
 				attribute.String(log.NamespaceKey+log.InstanceIDKey, instanceState.Instance.InstanceID),

go.mod

@@ -12,6 +12,8 @@ require (
 	github.com/redis/go-redis/v9 v9.0.2
 	github.com/stretchr/testify v1.8.3
 	go.opentelemetry.io/otel v1.16.0
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.16.0
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.16.0
 	go.opentelemetry.io/otel/exporters/stdout/stdouttrace v1.16.0
 	go.opentelemetry.io/otel/trace v1.16.0
 	golang.org/x/tools v0.6.0
@@ -24,6 +26,7 @@ require (
 	github.com/breml/bidichk v0.2.3 // indirect
 	github.com/curioswitch/go-reassign v0.2.0 // indirect
 	github.com/firefart/nonamedreturns v1.0.4 // indirect
+	github.com/grpc-ecosystem/grpc-gateway/v2 v2.7.0 // indirect
 	github.com/kkHAIKE/contextcheck v1.1.2 // indirect
 	github.com/lufeee/execinquery v1.2.1 // indirect
 	github.com/maratori/testableexamples v1.0.0 // indirect
@@ -34,12 +37,16 @@ require (
 	github.com/sivchari/nosnakecase v1.7.0 // indirect
 	github.com/stbenjam/no-sprintf-host-port v0.1.1 // indirect
 	github.com/timonwong/loggercheck v0.9.3 // indirect
+	go.opentelemetry.io/otel/exporters/otlp/internal/retry v1.16.0 // indirect
 	go.opentelemetry.io/otel/metric v1.16.0 // indirect
+	go.opentelemetry.io/proto/otlp v0.19.0 // indirect
 	go.uber.org/atomic v1.7.0 // indirect
 	go.uber.org/multierr v1.6.0 // indirect
 	go.uber.org/zap v1.17.0 // indirect
 	golang.org/x/exp v0.0.0-20220722155223-a9213eeb770e // indirect
 	golang.org/x/exp/typeparams v0.0.0-20220827204233-334a2380cb91 // indirect
+	google.golang.org/genproto v0.0.0-20230306155012-7f2fa6fef1f4 // indirect
+	google.golang.org/grpc v1.55.0 // indirect
 )
 
 require (

go.sum

@@ -6,6 +6,7 @@ import (
 	"github.com/cschleiden/go-workflows/internal/contextpropagation"
 	"github.com/cschleiden/go-workflows/internal/core"
 	"github.com/cschleiden/go-workflows/internal/sync"
+	"github.com/cschleiden/go-workflows/internal/workflowtracer"
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -15,52 +16,40 @@ var propagator propagation.TextMapPropagator = propagation.NewCompositeTextMapPr
 	propagation.Baggage{},
 )
 
-func InjectSpan(ctx context.Context, metadata *core.WorkflowMetadata) {
+func injectSpan(ctx context.Context, metadata *core.WorkflowMetadata) {
 	propagator.Inject(ctx, metadata)
 }
 
-func ExtractSpan(ctx context.Context, metadata *core.WorkflowMetadata) context.Context {
+func extractSpan(ctx context.Context, metadata *core.WorkflowMetadata) context.Context {
 	return propagator.Extract(ctx, metadata)
 }
 
-type traceContextKeyType int
-
-const currentSpanKey traceContextKeyType = iota
-
-func WorkflowContextWithSpan(ctx sync.Context, span trace.Span) sync.Context {
-	return sync.WithValue(ctx, currentSpanKey, span)
-}
-
-func SpanFromWorkflowContext(ctx sync.Context) trace.Span {
-	if span, ok := ctx.Value(currentSpanKey).(trace.Span); ok {
-		return span
-	}
-
-	panic("no span in context")
-}
-
 type TracingContextPropagator struct {
 }
 
 var _ contextpropagation.ContextPropagator = &TracingContextPropagator{}
 
 func (*TracingContextPropagator) Inject(ctx context.Context, metadata *core.WorkflowMetadata) error {
-	InjectSpan(ctx, metadata)
+	injectSpan(ctx, metadata)
 	return nil
 }
 
 func (*TracingContextPropagator) Extract(ctx context.Context, metadata *core.WorkflowMetadata) (context.Context, error) {
-	return ExtractSpan(ctx, metadata), nil
+	return extractSpan(ctx, metadata), nil
 }
 
 func (*TracingContextPropagator) InjectFromWorkflow(ctx sync.Context, metadata *core.WorkflowMetadata) error {
-	// Ignore
+	span := workflowtracer.SpanFromContext(ctx)
+	sctx := trace.ContextWithSpan(context.Background(), span)
+
+	injectSpan(sctx, metadata)
 
 	return nil
 }
 
 func (*TracingContextPropagator) ExtractToWorkflow(ctx sync.Context, metadata *core.WorkflowMetadata) (sync.Context, error) {
-	// Ignore
+	sctx := extractSpan(context.Background(), metadata)
+	span := trace.SpanFromContext(sctx)
 
-	return ctx, nil
+	return workflowtracer.ContextWithSpan(ctx, span), nil
 }

internal/tracing/tracing.go

@@ -15,11 +15,9 @@ import (
 	"github.com/cschleiden/go-workflows/internal/payload"
 	"github.com/cschleiden/go-workflows/internal/sync"
 	"github.com/cschleiden/go-workflows/internal/task"
-	"github.com/cschleiden/go-workflows/internal/tracing"
 	"github.com/cschleiden/go-workflows/internal/workflowstate"
 	"github.com/cschleiden/go-workflows/internal/workflowtracer"
 	"github.com/cschleiden/go-workflows/log"
-	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -53,6 +51,7 @@ type executor struct {
 	logger            log.Logger
 	tracer            trace.Tracer
 	lastSequenceID    int64
+	parentSpan        trace.Span
 }
 
 func NewExecutor(
@@ -85,6 +84,9 @@ func NewExecutor(
 		}
 	}
 
+	// Get span from the workflow context, set by the default context propagator
+	parentSpan := workflowtracer.SpanFromContext(wfCtx)
+
 	return &executor{
 		registry:          registry,
 		historyProvider:   historyProvider,
@@ -95,23 +97,11 @@ func NewExecutor(
 		clock:             clock,
 		logger:            logger,
 		tracer:            tracer,
+		parentSpan:        parentSpan,
 	}, nil
 }
 
 func (e *executor) ExecuteTask(ctx context.Context, t *task.Workflow) (*ExecutionResult, error) {
-	ctx = tracing.ExtractSpan(ctx, t.Metadata)
-	ctx, span := e.tracer.Start(ctx, "WorkflowTaskExecution", trace.WithAttributes(
-		attribute.String(log.InstanceIDKey, t.WorkflowInstance.InstanceID),
-		attribute.String(log.TaskIDKey, t.ID),
-		attribute.Int(log.NewEventsKey, len(t.NewEvents)),
-	))
-	defer span.End()
-
-	// Make the current span available to the tracer that we pass into the workflow execution. With caching
-	// the executor instance might be used for multiple workflow tasks and we want calls made in each task
-	// execution to be associated with the new span for the WorkflowTaskExecution.
-	e.workflowTracer.UpdateExecution(span)
-
 	logger := e.logger.With(
 		log.TaskIDKey, t.ID,
 		log.InstanceIDKey, t.WorkflowInstance.InstanceID)

internal/workflow/executor.go

@@ -32,9 +32,9 @@ func ContextWithSpan(ctx sync.Context, span trace.Span) sync.Context {
 	return sync.WithValue(ctx, spanKey, span)
 }
 
-func SpanFromContext(ctx sync.Context) *trace.Span {
+func SpanFromContext(ctx sync.Context) trace.Span {
 	if span, ok := ctx.Value(spanKey).(trace.Span); ok {
-		return &span
+		return span
 	}
 
 	return nil
@@ -51,21 +51,10 @@ func New(tracer trace.Tracer) *WorkflowTracer {
 	}
 }
 
-func (wt *WorkflowTracer) UpdateExecution(span trace.Span) {
-	wt.parentSpan = span
-}
-
 func (wt *WorkflowTracer) Start(ctx sync.Context, name string, opts ...trace.SpanStartOption) (sync.Context, Span) {
-	parentSpan := wt.parentSpan
-
-	// Try to get parent span from ctx, otherwise use the one from the tracer instance
-	if span := SpanFromContext(ctx); span != nil {
-		parentSpan = *span
-	}
-
 	state := workflowstate.WorkflowState(ctx)
 
-	sctx := trace.ContextWithSpan(context.Background(), parentSpan)
+	sctx := trace.ContextWithSpan(context.Background(), SpanFromContext(ctx))
 	opts = append(opts, trace.WithTimestamp(state.Time()))
 	sctx, span := wt.tracer.Start(sctx, name, opts...)
 

internal/workflowtracer/tracer.go

@@ -3,6 +3,8 @@ package log
 const (
 	NamespaceKey = "workflows"
 
+	ErrorKey = "error"
+
 	ActivityIDKey   = NamespaceKey + ".activity.id"
 	ActivityNameKey = NamespaceKey + ".activity.name"
 	InstanceIDKey   = NamespaceKey + ".instance.id"

log/fields.go

@@ -14,6 +14,8 @@ import (
 	"go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 
+	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
+	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
 	"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
 
 	"github.com/cschleiden/go-workflows/worker"
@@ -31,14 +33,25 @@ func main() {
 		attribute.String("environment", "sample"),
 	)
 
+	oclient := otlptracehttp.NewClient(
+		// otlptracehttp.WithEndpoint("localhost:8360"),
+		// otlptracehttp.WithURLPath("/traces/otlp/v0.9"),
+		otlptracehttp.WithEndpoint("localhost:4318"),
+		otlptracehttp.WithInsecure(),
+	)
+	exp, err := otlptrace.New(ctx, oclient)
+	if err != nil {
+		panic(err)
+	}
+
 	stdoutexp, err := stdouttrace.New(stdouttrace.WithPrettyPrint())
 	if err != nil {
 		panic(err)
 	}
 
 	tp := trace.NewTracerProvider(
 		trace.WithSyncer(stdoutexp),
-		// trace.WithBatcher(exp),
+		trace.WithBatcher(exp),
 		trace.WithResource(r),
 	)
 
@@ -75,7 +88,7 @@ func runWorkflow(ctx context.Context, c client.Client) {
 	}
 
 	// Test signal
-	time.Sleep(time.Millisecond * 500)
+	time.Sleep(time.Second * 5)
 	c.SignalWorkflow(ctx, wf.InstanceID, "test-signal", "")
 
 	result, err := client.GetWorkflowResult[int](ctx, c, wf, time.Second*120)

samples/tracing/tracing.go

@@ -21,14 +21,18 @@ func Workflow1(ctx workflow.Context, msg string, times int, inputs Inputs) (int,
 	logger.Debug("Entering Workflow1", "msg", msg, "times", times, "inputs", inputs)
 	defer logger.Debug("Leaving Workflow1")
 
-	ctx, span := workflow.Tracer(ctx).Start(ctx, "Workflow1 span", trace.WithAttributes(
+	tracer := workflow.Tracer(ctx)
+	ctx, span := tracer.Start(ctx, "Workflow1")
+	defer span.End()
+
+	_, customSpan := tracer.Start(ctx, "Workflow1 span", trace.WithAttributes(
 		// Add additional
 		attribute.String("msg", "hello world"),
 	))
 
 	// Do something
 
-	span.End()
+	customSpan.End()
 
 	workflow.ExecuteActivity[int](ctx, workflow.DefaultActivityOptions, Activity1, 35, 12).Get(ctx)
 
@@ -38,6 +42,10 @@ func Workflow1(ctx workflow.Context, msg string, times int, inputs Inputs) (int,
 
 	workflow.ExecuteActivity[int](ctx, workflow.DefaultActivityOptions, Activity1, 35, 12).Get(ctx)
 
+	_, span = workflow.Tracer(ctx).Start(ctx, "Wait-for-signal")
+	workflow.NewSignalChannel[string](ctx, "test-signal").Receive(ctx)
+	span.End()
+
 	r1, _ := workflow.ExecuteActivity[int](ctx, workflow.DefaultActivityOptions, Activity1, 35, 12).Get(ctx)
 
 	return r1, nil