Add trace support into task

backend/activity.go

@@ -68,6 +68,8 @@ func (p *activityProcessor) ProcessWorkItem(ctx context.Context, wi WorkItem) er
 		}()
 	}
 
+	// 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 {

backend/executor.go

@@ -148,6 +148,7 @@ func (executor *grpcExecutor) ExecuteActivity(ctx context.Context, iid api.Insta
 				Input:                 task.Input,
 				OrchestrationInstance: &protos.OrchestrationInstance{InstanceId: string(iid)},
 				TaskId:                e.EventId,
+				ParentTraceContext:    task.ParentTraceContext,
 			},
 		},
 	}

backend/runtimestate.go

@@ -158,7 +158,7 @@ func (s *OrchestrationRuntimeState) ApplyActions(actions []*protos.OrchestratorA
 			}
 		} else if createtimer := action.GetCreateTimer(); createtimer != nil {
 			s.AddEvent(helpers.NewTimerCreatedEvent(action.Id, createtimer.FireAt))
-			s.pendingTimers = append(s.pendingTimers, helpers.NewTimerFiredEvent(action.Id, createtimer.FireAt, currentTraceContext))
+			s.pendingTimers = append(s.pendingTimers, helpers.NewTimerFiredEvent(action.Id, createtimer.FireAt))
 		} else if scheduleTask := action.GetScheduleTask(); scheduleTask != nil {
 			scheduledEvent := helpers.NewTaskScheduledEvent(
 				action.Id,

client/worker_grpc.go

@@ -94,7 +94,7 @@ func (c *TaskHubGrpcClient) processActivityWorkItem(
 	executor backend.Executor,
 	req *protos.ActivityRequest,
 ) {
-	var tc *protos.TraceContext = nil // TODO: How to populate trace context?
+	var tc *protos.TraceContext = req.ParentTraceContext
 	event := helpers.NewTaskScheduledEvent(req.TaskId, req.Name, req.Version, req.Input, tc)
 	result, err := executor.ExecuteActivity(ctx, api.InstanceID(req.OrchestrationInstance.InstanceId), event)
 

internal/helpers/history.go

@@ -145,7 +145,6 @@ func NewTimerCreatedEvent(eventID int32, fireAt *timestamppb.Timestamp) *protos.
 func NewTimerFiredEvent(
 	timerID int32,
 	fireAt *timestamppb.Timestamp,
-	parentTraceContext *protos.TraceContext,
 ) *protos.HistoryEvent {
 	return &protos.HistoryEvent{
 		EventId:   -1,

samples/distributedtracing/distributedtracing.go

@@ -18,6 +18,8 @@ import (
 	"github.com/microsoft/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):

task/activity.go

@@ -2,7 +2,9 @@ package task
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/microsoft/durabletask-go/internal/helpers"
 	"github.com/microsoft/durabletask-go/internal/protos"
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
@@ -52,6 +54,10 @@ type activityContext struct {
 type Activity func(ctx ActivityContext) (any, error)
 
 func newTaskActivityContext(ctx context.Context, taskID int32, ts *protos.TaskScheduledEvent) *activityContext {
+	ctx, err := helpers.ContextFromTraceContext(ctx, ts.ParentTraceContext)
+	if err != nil {
+		fmt.Printf("%v: failed to parse trace context: %v", ts.Name, err)
+	}
 	return &activityContext{
 		TaskID:   taskID,
 		Name:     ts.Name,

tests/orchestrations_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel"
 
 	"github.com/microsoft/durabletask-go/api"
 	"github.com/microsoft/durabletask-go/backend"
@@ -19,6 +20,8 @@ import (
 	"github.com/microsoft/durabletask-go/task"
 )
 
+var tracer = otel.Tracer("orchestration-test")
+
 func Test_EmptyOrchestration(t *testing.T) {
 	// Registration
 	r := task.NewTaskRegistry()
@@ -207,6 +210,56 @@ func Test_SingleActivity(t *testing.T) {
 	)
 }
 
+func Test_SingleActivity_TaskSpan(t *testing.T) {
+	// Registration
+	r := task.NewTaskRegistry()
+	r.AddOrchestratorN("SingleActivity", 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")
+		childSpan.End()
+		return fmt.Sprintf("Hello, %s!", name), nil
+	})
+
+	// Initialization
+	ctx := context.Background()
+	exporter := initTracing()
+	client, worker := initTaskHubWorker(ctx, r)
+	defer worker.Shutdown(ctx)
+
+	// Run the orchestration
+	id, err := client.ScheduleNewOrchestration(ctx, "SingleActivity", api.WithInput("世界"))
+	if assert.NoError(t, err) {
+		metadata, err := client.WaitForOrchestrationCompletion(ctx, id)
+		if assert.NoError(t, err) {
+			assert.Equal(t, protos.OrchestrationStatus_ORCHESTRATION_STATUS_COMPLETED, metadata.RuntimeStatus)
+			assert.Equal(t, `"Hello, 世界!"`, metadata.SerializedOutput)
+		}
+	}
+
+	// Validate the exported OTel traces
+	spans := exporter.GetSpans().Snapshots()
+	assertSpanSequence(t, spans,
+		assertOrchestratorCreated("SingleActivity", id),
+		assertSpan("activityChild"),
+		assertActivity("SayHello", id, 0),
+		assertOrchestratorExecuted("SingleActivity", id, "COMPLETED"),
+	)
+	// assert child-parent relationship
+	assert.Equal(t, spans[1].Parent().SpanID(), spans[2].SpanContext().SpanID())
+}
+
 func Test_ActivityChain(t *testing.T) {
 	// Registration
 	r := task.NewTaskRegistry()