Merge pull request #39 from salaboy/main

Adding support for getting the TraceParent context

Author: JoshVanL

backend/activity.go

@@ -49,7 +49,6 @@ func (p *activityProcessor) ProcessWorkItem(ctx context.Context, awi *ActivityWo
 	if ts == nil {
 		return fmt.Errorf("%v: invalid TaskScheduled event", awi.InstanceID)
 	}
-
 	// Create span as child of spanContext found in TaskScheduledEvent
 	ctx, err := helpers.ContextFromTraceContext(ctx, ts.ParentTraceContext)
 	if err != nil {
@@ -66,6 +65,9 @@ func (p *activityProcessor) ProcessWorkItem(ctx context.Context, awi *ActivityWo
 		}()
 	}
 
+	// set the parent trace context to be the newly created activity span
+	ts.ParentTraceContext = helpers.TraceContextFromSpan(span)
+
 	// Execute the activity and get its result
 	result, err := p.executor.ExecuteActivity(ctx, awi.InstanceID, awi.NewEvent)
 	if err != nil {
@@ -75,7 +77,6 @@ func (p *activityProcessor) ProcessWorkItem(ctx context.Context, awi *ActivityWo
 		}
 		return err
 	}
-
 	awi.Result = result
 	return nil
 }

backend/executor.go

@@ -177,6 +177,7 @@ func (executor *grpcExecutor) ExecuteActivity(ctx context.Context, iid api.Insta
 		OrchestrationInstance: &protos.OrchestrationInstance{InstanceId: string(iid)},
 		TaskId:                e.EventId,
 		TaskExecutionId:       task.TaskExecutionId,
+		ParentTraceContext:    task.ParentTraceContext,
 	}
 	workItem := &protos.WorkItem{
 		Request: &protos.WorkItem_ActivityRequest{

client/worker_grpc.go

@@ -15,6 +15,7 @@ import (
 	"google.golang.org/protobuf/types/known/wrapperspb"
 
 	"github.com/dapr/durabletask-go/api"
+	"github.com/dapr/durabletask-go/api/helpers"
 	"github.com/dapr/durabletask-go/api/protos"
 	"github.com/dapr/durabletask-go/backend"
 	"github.com/dapr/durabletask-go/task"
@@ -174,7 +175,12 @@ func (c *TaskHubGrpcClient) processActivityWorkItem(
 	executor backend.Executor,
 	req *protos.ActivityRequest,
 ) {
-	var tc *protos.TraceContext = nil // TODO: How to populate trace context?
+	var ptc *protos.TraceContext = req.ParentTraceContext
+	ctx, err := helpers.ContextFromTraceContext(ctx, ptc)
+	if err != nil {
+		c.logger.Warn("%v: failed to parse trace context: %v", req.Name, err)
+	}
+
 	event := &protos.HistoryEvent{
 		EventId:   req.TaskId,
 		Timestamp: timestamppb.New(time.Now()),
@@ -184,7 +190,7 @@ func (c *TaskHubGrpcClient) processActivityWorkItem(
 				Version:            req.Version,
 				Input:              req.Input,
 				TaskExecutionId:    req.TaskExecutionId,
-				ParentTraceContext: tc,
+				ParentTraceContext: ptc,
 			},
 		},
 	}

samples/distributedtracing/distributedtracing.go

@@ -18,6 +18,8 @@ import (
 	"github.com/dapr/durabletask-go/task"
 )
 
+var tracer = otel.Tracer("distributedtracing-example")
+
 func main() {
 	// Tracing can be configured independently of the orchestration code.
 	tp, err := ConfigureZipkinTracing()
@@ -136,6 +138,16 @@ func DoWorkActivity(ctx task.ActivityContext) (any, error) {
 		return "", err
 	}
 
+	_, childSpan := tracer.Start(ctx.Context(), "activity-subwork")
+	// Simulate doing some sub work
+	select {
+	case <-time.After(2 * time.Second):
+		// Ok
+	case <-ctx.Context().Done():
+		return nil, ctx.Context().Err()
+	}
+	childSpan.End()
+
 	// Simulate doing work
 	select {
 	case <-time.After(duration):

tests/grpc/grpc_test.go

@@ -21,11 +21,14 @@ import (
 	"github.com/dapr/durabletask-go/backend/sqlite"
 	"github.com/dapr/durabletask-go/client"
 	"github.com/dapr/durabletask-go/task"
+	"github.com/dapr/durabletask-go/tests/utils"
+	"go.opentelemetry.io/otel"
 )
 
 var (
 	grpcClient *client.TaskHubGrpcClient
 	ctx        = context.Background()
+	tracer     = otel.Tracer("grpc-test")
 )
 
 // TestMain is the entry point for the test suite. We use this to set up a gRPC server and client instance
@@ -489,3 +492,51 @@ func Test_Grpc_SubOrchestratorRetries(t *testing.T) {
 	// With 3 max attempts there will be two retries with 10 millis delay before each
 	require.GreaterOrEqual(t, metadata.LastUpdatedAt.AsTime(), metadata.CreatedAt.AsTime().Add(2*10*time.Millisecond))
 }
+
+func Test_SingleActivity_TaskSpan(t *testing.T) {
+	// Registration
+	r := task.NewTaskRegistry()
+	r.AddOrchestratorN("SingleActivity_TestSpan", func(ctx *task.OrchestrationContext) (any, error) {
+		var input string
+		if err := ctx.GetInput(&input); err != nil {
+			return nil, err
+		}
+		var output string
+		err := ctx.CallActivity("SayHello", task.WithActivityInput(input)).Await(&output)
+		return output, err
+	})
+	r.AddActivityN("SayHello", func(ctx task.ActivityContext) (any, error) {
+		var name string
+		if err := ctx.GetInput(&name); err != nil {
+			return nil, err
+		}
+		_, childSpan := tracer.Start(ctx.Context(), "activityChild_TestSpan")
+		childSpan.End()
+		return fmt.Sprintf("Hello, %s!", name), nil
+	})
+
+	exporter := utils.InitTracing()
+	cancelListener := startGrpcListener(t, r)
+	defer cancelListener()
+
+	// Run the orchestration
+	id, err := grpcClient.ScheduleNewOrchestration(ctx, "SingleActivity_TestSpan", api.WithInput("世界"), api.WithStartTime(time.Now()))
+	if assert.NoError(t, err) {
+		metadata, err := grpcClient.WaitForOrchestrationCompletion(ctx, id)
+		if assert.NoError(t, err) {
+			assert.Equal(t, protos.OrchestrationStatus_ORCHESTRATION_STATUS_COMPLETED, metadata.RuntimeStatus)
+			assert.Equal(t, `"Hello, 世界!"`, metadata.Output.Value)
+		}
+	}
+
+	// Validate the exported OTel traces
+	spans := exporter.GetSpans().Snapshots()
+	utils.AssertSpanSequence(t, spans,
+		utils.AssertOrchestratorCreated("SingleActivity_TestSpan", id),
+		utils.AssertSpan("activityChild_TestSpan"),
+		utils.AssertActivity("SayHello", id, 0),
+		utils.AssertOrchestratorExecuted("SingleActivity_TestSpan", id, "COMPLETED"),
+	)
+	// assert child-parent relationship
+	assert.Equal(t, spans[1].Parent().SpanID(), spans[2].SpanContext().SpanID())
+}