feat: Support StepFunction->StepFunction trace merging

examples/step-functions-typescript-stack/lib/cdk-step-functions-typescript-stack.ts

@@ -11,6 +11,25 @@ export class CdkStepFunctionsTypeScriptStack extends Stack {
 
     console.log("Creating Hello World TypeScript stack");
 
+    /* Set up the child state machine */
+
+    const passState = new sfn.Pass(this, "PassState");
+    const waitState = new sfn.Wait(this, "WaitState", { time: sfn.WaitTime.duration(Duration.seconds(1)) });
+    const successState = new sfn.Succeed(this, "SuccessState");
+
+    const childStateMachine = new sfn.StateMachine(this, "CdkTypeScriptTestChildStateMachine", {
+      definitionBody: sfn.DefinitionBody.fromChainable(passState.next(waitState).next(successState)),
+    });
+
+    /* Set up the parent state machine */
+
+    const invokeChildStateMachineTask = new tasks.StepFunctionsStartExecution(this, "InvokeChildStateMachineTask", {
+      stateMachine: childStateMachine,
+      input: sfn.TaskInput.fromObject(
+        DatadogStepFunctions.buildStepFunctionTaskInputToMergeTraces({ "custom-key": "custom-value" }),
+      ),
+    });
+
     const helloLambdaFunction = new lambda.Function(this, "hello-python", {
       runtime: lambda.Runtime.PYTHON_3_12,
       timeout: Duration.seconds(10),
@@ -24,15 +43,17 @@ export class CdkStepFunctionsTypeScriptStack extends Stack {
       handler: "hello.lambda_handler",
     });
 
-    const stateMachine = new sfn.StateMachine(this, "CdkTypeScriptTestStateMachine", {
-      definitionBody: sfn.DefinitionBody.fromChainable(
-        new tasks.LambdaInvoke(this, "MyLambdaTask", {
-          lambdaFunction: helloLambdaFunction,
-          payload: sfn.TaskInput.fromObject(DatadogStepFunctions.buildLambdaPayloadToMergeTraces()),
-        }).next(new sfn.Succeed(this, "GreetedWorld")),
-      ),
+    const lambdaTask = new tasks.LambdaInvoke(this, "MyLambdaTask", {
+      lambdaFunction: helloLambdaFunction,
+      payload: sfn.TaskInput.fromObject(DatadogStepFunctions.buildLambdaPayloadToMergeTraces()),
     });
 
+    const parentStateMachine = new sfn.StateMachine(this, "CdkTypeScriptTestStateMachine", {
+      definitionBody: sfn.DefinitionBody.fromChainable(lambdaTask.next(invokeChildStateMachineTask)),
+    });
+
+    /* Instrument the lambda functions and the state machines */
+
     console.log("Instrumenting Step Functions in TypeScript stack with Datadog");
 
     const datadogSfn = new DatadogStepFunctions(this, "DatadogSfn", {
@@ -42,7 +63,7 @@ export class CdkStepFunctionsTypeScriptStack extends Stack {
       forwarderArn: process.env.DD_FORWARDER_ARN,
       tags: "custom-tag-1:tag-value-1,custom-tag-2:tag-value-2",
     });
-    datadogSfn.addStateMachines([stateMachine]);
+    datadogSfn.addStateMachines([childStateMachine, parentStateMachine]);
 
     const datadogLambda = new DatadogLambda(this, "DatadogLambda", {
       pythonLayerVersion: 101,

src/datadog-step-functions.ts

@@ -14,7 +14,10 @@ import { Construct } from "constructs";
 import log from "loglevel";
 import { addForwarderForStateMachine } from "./forwarder";
 import { DatadogStepFunctionsProps } from "./index";
-import { buildStepFunctionLambdaTaskPayloadToMergeTraces } from "./span-link";
+import {
+  buildStepFunctionLambdaTaskPayloadToMergeTraces,
+  buildStepFunctionSfnExecutionTaskInputToMergeTraces,
+} from "./span-link";
 import { setTags } from "./tag";
 
 const unsupportedCaseErrorMessage =
@@ -26,6 +29,12 @@ export class DatadogStepFunctions extends Construct {
     return buildStepFunctionLambdaTaskPayloadToMergeTraces(payload);
   }
 
+  public static buildStepFunctionTaskInputToMergeTraces(input: { [key: string]: any } = {}): {
+    [key: string]: any;
+  } {
+    return buildStepFunctionSfnExecutionTaskInputToMergeTraces(input);
+  }
+
   scope: Construct;
   props: DatadogStepFunctionsProps;
   stack: Stack;

src/span-link.ts

@@ -31,3 +31,29 @@ Please open an issue in https://github.com/DataDog/datadog-cdk-constructs to dis
   payload["StateMachine.$"] = "$$.StateMachine";
   return payload;
 }
+
+/**
+ * Builds an input for a Step Function execution task, so the Step Function traces
+ * can be merged with downstream Step Function traces.
+ *
+ * This function modifies the provided input to include context fields necessary
+ * for trace merging purposes. If the input already contains CONTEXT or CONTEXT.$ field,
+ * an error is thrown to avoid conflicts.
+ *
+ * @param input - The user's input object. Defaults to an empty object.
+ * @returns The modified input object with necessary context added.
+ * @throws {ConflictError} If the input already contains `CONTEXT` or `CONTEXT.$` fields.
+ */
+
+export function buildStepFunctionSfnExecutionTaskInputToMergeTraces(input: { [key: string]: any } = {}): {
+  [key: string]: any;
+} {
+  if ("CONTEXT" in input || "CONTEXT.$" in input) {
+    throw new Error(`The StepFunction StartExecution task may be using custom CONTEXT field. Step Functions Context Object injection skipped. \
+Your Step Functions trace will not be merged with downstream Lambda traces. \
+Please open an issue in https://github.com/DataDog/datadog-cdk-constructs to discuss your workaround.`);
+  }
+
+  input["CONTEXT.$"] = `$$['Execution', 'State', 'StateMachine']`;
+  return input;
+}

test/span-link.spec.ts

@@ -1,4 +1,7 @@
-import { buildStepFunctionLambdaTaskPayloadToMergeTraces } from "../src/span-link";
+import {
+  buildStepFunctionLambdaTaskPayloadToMergeTraces,
+  buildStepFunctionSfnExecutionTaskInputToMergeTraces,
+} from "../src/span-link";
 
 describe("buildStepFunctionLambdaTaskPayloadToMergeTraces", () => {
   it("adds necessary fields to an empty payload", () => {
@@ -28,3 +31,28 @@ describe("buildStepFunctionLambdaTaskPayloadToMergeTraces", () => {
     );
   });
 });
+
+describe("buildStepFunctionSfnExecutionTaskInputToMergeTraces", () => {
+  it("adds necessary fields to an empty input", () => {
+    const result = buildStepFunctionSfnExecutionTaskInputToMergeTraces();
+    expect(result).toEqual({
+      "CONTEXT.$": `$$['Execution', 'State', 'StateMachine']`,
+    });
+  });
+
+  it("adds necessary fields to a non-empty input", () => {
+    const input = { "custom-key": "custom-value" };
+    const result = buildStepFunctionSfnExecutionTaskInputToMergeTraces(input);
+    expect(result).toEqual({
+      "custom-key": "custom-value",
+      "CONTEXT.$": `$$['Execution', 'State', 'StateMachine']`,
+    });
+  });
+
+  it("throws an error if input already contains CONTEXT field", () => {
+    const input = { CONTEXT: "value" };
+    expect(() => buildStepFunctionSfnExecutionTaskInputToMergeTraces(input)).toThrowError(
+      "The StepFunction StartExecution task may be using custom CONTEXT field. Step Functions Context Object injection skipped. Your Step Functions trace will not be merged with downstream Lambda traces. Please open an issue in https://github.com/DataDog/datadog-cdk-constructs to discuss your workaround.",
+    );
+  });
+});