diff --git a/CHANGELOG.md b/CHANGELOG.md index acdba14..845f0d4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,6 +40,8 @@ ### Fixes +- Make the Durable Task activity execution span active while user activity code runs, so + spans created by activities inherit the correct parent across awaits and concurrent executions. - Preserve received structured failure `properties` in task errors, retry callbacks, client state/history, entity failures, and testing. Forward them with rethrown task failure chains. Empty wire maps remain `undefined`; arbitrary JavaScript `Error` fields are not collected. diff --git a/packages/durabletask-js/src/worker/task-hub-grpc-worker.ts b/packages/durabletask-js/src/worker/task-hub-grpc-worker.ts index 9247cc6..d6e7854 100644 --- a/packages/durabletask-js/src/worker/task-hub-grpc-worker.ts +++ b/packages/durabletask-js/src/worker/task-hub-grpc-worker.ts @@ -38,6 +38,7 @@ import { setSpanError, setSpanOk, endSpan, + getOtelApi, } from "../tracing"; /** Default timeout in milliseconds for graceful shutdown. */ @@ -1351,13 +1352,19 @@ export class TaskHubGrpcWorker { ); } else { const executor = new ActivityExecutor(this._registry, this._logger); - const result = await executor.execute( - instanceId, - req.getName(), - req.getTaskid(), - req.getInput()?.getValue() ?? "", - req.getVersion()?.getValue(), - ); + const executeActivity = () => + executor.execute( + instanceId, + req.getName(), + req.getTaskid(), + req.getInput()?.getValue() ?? "", + req.getVersion()?.getValue(), + ); + const otel = getOtelApi(); + const result = + otel && activitySpan + ? await otel.context.with(otel.trace.setSpan(otel.context.active(), activitySpan), executeActivity) + : await executeActivity(); res.setResult(new StringValue().setValue(result ?? "")); setSpanOk(activitySpan); } diff --git a/packages/durabletask-js/test/worker-tracing.spec.ts b/packages/durabletask-js/test/worker-tracing.spec.ts index ea4f659..015548b 100644 --- a/packages/durabletask-js/test/worker-tracing.spec.ts +++ b/packages/durabletask-js/test/worker-tracing.spec.ts @@ -12,16 +12,20 @@ */ import * as pb from "../src/proto/orchestrator_service_pb"; +import * as stubs from "../src/proto/orchestrator_service_grpc_pb"; import { newOrchestratorStartedEvent, newExecutionStartedEvent, } from "../src/utils/pb-helper.util"; import { OrchestrationExecutor } from "../src/worker/orchestration-executor"; import { Registry } from "../src/worker/registry"; +import { TaskHubGrpcWorker } from "../src/worker/task-hub-grpc-worker"; import { NoOpLogger } from "../src/types/logger.type"; import { TOrchestrator } from "../src/types/orchestrator.type"; import { OrchestrationContext } from "../src/task/context/orchestration-context"; import { ActivityContext } from "../src/task/context/activity-context"; +import * as traceContextUtils from "../src/tracing/trace-context-utils"; +import { AsyncLocalStorage } from "async_hooks"; import { startSpanForOrchestrationExecution, @@ -50,17 +54,48 @@ let exporter: InMemorySpanExporter; let provider: BasicTracerProvider; let previousTracerProvider: otel.TracerProvider | undefined; +class TestAsyncContextManager implements otel.ContextManager { + private readonly storage = new AsyncLocalStorage(); + + active(): otel.Context { + return this.storage.getStore() ?? otel.ROOT_CONTEXT; + } + + with ReturnType>( + context: otel.Context, + fn: F, + thisArg?: ThisParameterType, + ...args: A + ): ReturnType { + return this.storage.run(context, () => fn.call(thisArg, ...args)); + } + + bind(_context: otel.Context, target: T): T { + return target; + } + + enable(): this { + return this; + } + + disable(): this { + this.storage.disable(); + return this; + } +} + beforeAll(() => { exporter = new InMemorySpanExporter(); provider = new BasicTracerProvider(); provider.addSpanProcessor(new SimpleSpanProcessor(exporter)); previousTracerProvider = otel.trace.getTracerProvider(); - provider.register(); + provider.register({ contextManager: new TestAsyncContextManager() }); }); afterAll(async () => { await provider.shutdown(); + otel.context.disable(); if (previousTracerProvider) { otel.trace.setGlobalTracerProvider(previousTracerProvider); @@ -71,6 +106,38 @@ beforeEach(() => { exporter.reset(); }); +function createActivityRequest( + name: string, + taskId: number, + traceId: string, + parentSpanId: string, +): pb.ActivityRequest { + return new pb.ActivityRequest() + .setName(name) + .setTaskid(taskId) + .setOrchestrationinstance(new pb.OrchestrationInstance().setInstanceid(`${name}-instance`)) + .setParenttracecontext(createPbTraceContext(`00-${traceId}-${parentSpanId}-01`)); +} + +function createActivityStub(): { + stub: stubs.TaskHubSidecarServiceClient; + response: () => pb.ActivityResponse | undefined; +} { + let capturedResponse: pb.ActivityResponse | undefined; + const stub = { + completeActivityTask: ( + response: pb.ActivityResponse, + _metadata: unknown, + callback: (error: Error | null, response: pb.CompleteTaskResponse) => void, + ) => { + capturedResponse = response; + callback(null, new pb.CompleteTaskResponse()); + }, + } as unknown as stubs.TaskHubSidecarServiceClient; + + return { stub, response: () => capturedResponse }; +} + /** * Helper that mirrors the worker's _executeOrchestratorInternal tracing flow. * This replicates the exact sequence of tracing calls made by the worker so @@ -432,3 +499,159 @@ describe("Worker Tracing - Activity Span Lifecycle", () => { expect(actSpan!.endTime).toBeDefined(); }); }); + +describe("Worker Tracing - Activity Active Context", () => { + const tracer = otel.trace.getTracer("worker-tracing-user-code"); + + it("parents user spans to the activity execution span before and after await", async () => { + const worker = new TaskHubGrpcWorker({ logger: testLogger }); + worker.addNamedActivity("parentedActivity", async () => { + const beforeAwait = tracer.startSpan("user-before-await"); + beforeAwait.end(); + + await Promise.resolve(); + + const afterAwait = tracer.startSpan("user-after-await"); + afterAwait.end(); + return "done"; + }); + const request = createActivityRequest( + "parentedActivity", + 1, + "11111111111111111111111111111111", + "1111111111111111", + ); + const { stub } = createActivityStub(); + const callerSpan = tracer.startSpan("success-caller"); + const callerContext = otel.trace.setSpan(otel.ROOT_CONTEXT, callerSpan); + + await otel.context.with(callerContext, async () => { + expect(otel.trace.getSpan(otel.context.active())).toBe(callerSpan); + await (worker as any)._executeActivityInternal(request, "completion-token", stub); + expect(otel.trace.getSpan(otel.context.active())).toBe(callerSpan); + }); + callerSpan.end(); + + const spans = exporter.getFinishedSpans(); + const activitySpan = spans.find((span) => span.name === "activity:parentedActivity")!; + const beforeAwait = spans.find((span) => span.name === "user-before-await")!; + const afterAwait = spans.find((span) => span.name === "user-after-await")!; + + expect(beforeAwait.parentSpanId).toBe(activitySpan.spanContext().spanId); + expect(afterAwait.parentSpanId).toBe(activitySpan.spanContext().spanId); + expect(beforeAwait.spanContext().traceId).toBe(activitySpan.spanContext().traceId); + expect(afterAwait.spanContext().traceId).toBe(activitySpan.spanContext().traceId); + }); + + it("keeps concurrent activity contexts isolated", async () => { + const worker = new TaskHubGrpcWorker({ logger: testLogger }); + let started = 0; + let releaseActivities!: () => void; + let reportBothStarted!: () => void; + const activitiesStarted = new Promise((resolve) => { + reportBothStarted = resolve; + }); + const release = new Promise((resolve) => { + releaseActivities = resolve; + }); + + for (const name of ["activityA", "activityB"]) { + worker.addNamedActivity(name, async () => { + started++; + if (started === 2) { + reportBothStarted(); + } + await release; + const child = tracer.startSpan(`user-${name}`); + child.end(); + return name; + }); + } + + const first = createActivityRequest("activityA", 1, "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", "aaaaaaaaaaaaaaaa"); + const second = createActivityRequest("activityB", 2, "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", "bbbbbbbbbbbbbbbb"); + const firstStub = createActivityStub(); + const secondStub = createActivityStub(); + + const firstExecution = (worker as any)._executeActivityInternal(first, "first-completion-token", firstStub.stub); + const secondExecution = (worker as any)._executeActivityInternal( + second, + "second-completion-token", + secondStub.stub, + ); + await activitiesStarted; + releaseActivities(); + await Promise.all([firstExecution, secondExecution]); + + const spans = exporter.getFinishedSpans(); + for (const name of ["activityA", "activityB"]) { + const activitySpan = spans.find((span) => span.name === `activity:${name}`)!; + const childSpan = spans.find((span) => span.name === `user-${name}`)!; + expect(childSpan.parentSpanId).toBe(activitySpan.spanContext().spanId); + expect(childSpan.spanContext().traceId).toBe(activitySpan.spanContext().traceId); + } + }); + + it("restores the caller context and preserves failure tracing when activity code throws", async () => { + const worker = new TaskHubGrpcWorker({ logger: testLogger }); + worker.addNamedActivity("throwingActivity", async () => { + await Promise.resolve(); + const child = tracer.startSpan("user-throwing-activity"); + child.end(); + throw new Error("activity exploded"); + }); + const request = createActivityRequest( + "throwingActivity", + 3, + "33333333333333333333333333333333", + "3333333333333333", + ); + const activityStub = createActivityStub(); + const callerSpan = tracer.startSpan("caller"); + const callerContext = otel.trace.setSpan(otel.ROOT_CONTEXT, callerSpan); + + await otel.context.with(callerContext, async () => { + expect(otel.trace.getSpan(otel.context.active())).toBe(callerSpan); + await (worker as any)._executeActivityInternal(request, "failure-completion-token", activityStub.stub); + expect(otel.trace.getSpan(otel.context.active())).toBe(callerSpan); + }); + callerSpan.end(); + + const spans = exporter.getFinishedSpans(); + const activitySpan = spans.find((span) => span.name === "activity:throwingActivity")!; + const childSpan = spans.find((span) => span.name === "user-throwing-activity")!; + + expect(childSpan.parentSpanId).toBe(activitySpan.spanContext().spanId); + expect(activitySpan.status).toEqual({ + code: otel.SpanStatusCode.ERROR, + message: "activity exploded", + }); + expect(activityStub.response()?.getFailuredetails()?.getErrormessage()).toBe("activity exploded"); + }); + + it("executes and completes activities when tracing is unavailable", async () => { + const worker = new TaskHubGrpcWorker({ logger: testLogger }); + worker.addNamedActivity("untracedActivity", async () => { + await Promise.resolve(); + return "untraced result"; + }); + const request = createActivityRequest( + "untracedActivity", + 4, + "44444444444444444444444444444444", + "4444444444444444", + ); + const activityStub = createActivityStub(); + const getOtelApi = jest.spyOn(traceContextUtils, "getOtelApi").mockReturnValue(undefined); + + try { + await (worker as any)._executeActivityInternal(request, "untraced-completion-token", activityStub.stub); + } finally { + getOtelApi.mockRestore(); + } + + expect(activityStub.response()?.getResult()?.getValue()).toBe(JSON.stringify("untraced result")); + expect(activityStub.response()?.getFailuredetails()).toBeUndefined(); + expect(exporter.getFinishedSpans()).toHaveLength(0); + }); +});