Darcy.rayner/step function support poc

package.json

@@ -50,7 +50,8 @@
     },
     "collectCoverage": true,
     "coverageReporters": [
-      "lcovonly"
+      "lcovonly",
+      "text-summary"
     ],
     "testRegex": "(src\\/).*(\\.spec\\.ts)$",
     "testPathIgnorePatterns": [

src/trace/constants.ts

@@ -10,4 +10,5 @@ export const parentIDHeader = "x-datadog-parent-id";
 export const samplingPriorityHeader = "x-datadog-sampling-priority";
 export const xraySubsegmentName = "datadog-metadata";
 export const xraySubsegmentKey = "trace";
+export const xrayBaggageSubsegmentKey = "root_span_metadata";
 export const xraySubsegmentNamespace = "datadog";

src/trace/context.spec.ts

@@ -1,5 +1,5 @@
 import { LogLevel, setLogLevel } from "../utils";
-import { SampleMode } from "./constants";
+import { SampleMode, xrayBaggageSubsegmentKey, xraySubsegmentNamespace } from "./constants";
 import {
   convertToAPMParentID,
   convertToAPMTraceID,
@@ -8,14 +8,19 @@ import {
   extractTraceContext,
   readTraceContextFromXray,
   readTraceFromEvent,
+  readStepFunctionContextFromEvent,
 } from "./context";
 
 let currentSegment: any;
 
 jest.mock("aws-xray-sdk-core", () => {
   return {
-    captureFunc: () => {
-      throw Error("Unimplemented");
+    captureFunc: (subsegmentName: string, callback: (segment: any) => void) => {
+      if (currentSegment) {
+        callback(currentSegment);
+      } else {
+        throw Error("Unimplemented");
+      }
     },
     getSegment: () => {
       if (currentSegment === undefined) {
@@ -219,6 +224,137 @@ describe("readTraceFromEvent", () => {
     expect(result).toBeUndefined();
   });
 });
+
+describe("readStepFunctionContextFromEvent", () => {
+  const stepFunctionEvent = {
+    dd: {
+      Execution: {
+        Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
+        StartTime: "2019-09-30T20:28:24.236Z",
+      },
+      State: {
+        Name: "step-one",
+        RetryCount: 2,
+      },
+      StateMachine: {
+        Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
+        Name: "my-state-machine",
+      },
+    },
+  } as const;
+  it("reads a trace from an execution id", () => {
+    const result = readStepFunctionContextFromEvent(stepFunctionEvent);
+    expect(result).toEqual({
+      "step_function.execution_id": "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
+      "step_function.retry_count": 2,
+      "step_function.state_machine_arn":
+        "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
+      "step_function.state_machine_name": "my-state-machine",
+      "step_function.step_name": "step-one",
+    });
+  });
+  it("returns undefined when event isn't an object", () => {
+    const result = readStepFunctionContextFromEvent("event");
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when event is missing datadogContext property", () => {
+    const result = readStepFunctionContextFromEvent({});
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when datadogContext is missing Execution property", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {},
+    });
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when Execution is missing Name field", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {
+        ...stepFunctionEvent.dd,
+        Execution: {},
+      },
+    });
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when Name isn't a string", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {
+        ...stepFunctionEvent.dd,
+        Execution: {
+          Name: 12345,
+        },
+      },
+    });
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when State isn't defined", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {
+        ...stepFunctionEvent.dd,
+        State: undefined,
+      },
+    });
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when try retry count isn't a number", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {
+        ...stepFunctionEvent.dd,
+        State: {
+          ...stepFunctionEvent.dd.State,
+          RetryCount: "1",
+        },
+      },
+    });
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when try step name isn't a string", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {
+        ...stepFunctionEvent.dd,
+        State: {
+          ...stepFunctionEvent.dd.State,
+          Name: 1,
+        },
+      },
+    });
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when StateMachine is undefined", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {
+        ...stepFunctionEvent.dd,
+        StateMachine: undefined,
+      },
+    });
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when StateMachineId isn't a string", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {
+        ...stepFunctionEvent.dd,
+        StateMachine: {
+          ...stepFunctionEvent.dd.StateMachine,
+          Id: 1,
+        },
+      },
+    });
+    expect(result).toBeUndefined();
+  });
+  it("returns undefined when StateMachineName isn't a string", () => {
+    const result = readStepFunctionContextFromEvent({
+      dd: {
+        ...stepFunctionEvent.dd,
+        StateMachine: {
+          ...stepFunctionEvent.dd.StateMachine,
+          Name: 1,
+        },
+      },
+    });
+    expect(result).toBeUndefined();
+  });
+});
+
 describe("extractTraceContext", () => {
   it("returns trace read from header as highest priority", () => {
     currentSegment = {
@@ -252,4 +388,51 @@ describe("extractTraceContext", () => {
       traceID: "4110911582297405557",
     });
   });
+  it("returns trace read from env if no headers present", () => {
+    currentSegment = {
+      id: "0b11cc4230d3e09e",
+      trace_id: "1-5ce31dc2-2c779014b90ce44db5e03875",
+    };
+
+    const result = extractTraceContext({});
+    expect(result).toEqual({
+      parentID: "797643193680388254",
+      sampleMode: SampleMode.USER_KEEP,
+      traceID: "4110911582297405557",
+    });
+  });
+
+  it("adds step function metadata to xray", () => {
+    const stepFunctionEvent = {
+      dd: {
+        Execution: {
+          Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
+          StartTime: "2019-09-30T20:28:24.236Z",
+        },
+        State: {
+          Name: "step-one",
+          RetryCount: 2,
+        },
+        StateMachine: {
+          Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
+          Name: "my-state-machine",
+        },
+      },
+    } as const;
+    const addMetadata = jest.fn();
+    currentSegment = { addMetadata };
+    extractTraceContext(stepFunctionEvent);
+    expect(addMetadata).toHaveBeenCalledWith(
+      xrayBaggageSubsegmentKey,
+      {
+        "step_function.execution_id": "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
+        "step_function.retry_count": 2,
+        "step_function.state_machine_arn":
+          "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
+        "step_function.state_machine_name": "my-state-machine",
+        "step_function.step_name": "step-one",
+      },
+      xraySubsegmentNamespace,
+    );
+  });
 });

src/trace/context.ts

@@ -1,12 +1,13 @@
 import { captureFunc, getSegment } from "aws-xray-sdk-core";
 import { BigNumber } from "bignumber.js";
 
-import { logError } from "../utils";
+import { logDebug, logError } from "../utils";
 import {
   parentIDHeader,
   SampleMode,
   samplingPriorityHeader,
   traceIDHeader,
+  xrayBaggageSubsegmentKey,
   xraySubsegmentKey,
   xraySubsegmentName,
   xraySubsegmentNamespace,
@@ -24,12 +25,28 @@ export interface TraceContext {
   sampleMode: SampleMode;
 }
 
+export interface StepFunctionContext {
+  "step_function.retry_count": number;
+  "step_function.execution_id": string;
+  "step_function.state_machine_name": string;
+  "step_function.state_machine_arn": string;
+  "step_function.step_name": string;
+}
+
 /**
  * Reads the trace context from either an incoming lambda event, or the current xray segment.
  * @param event An incoming lambda event. This must have incoming trace headers in order to be read.
  */
 export function extractTraceContext(event: any) {
   const trace = readTraceFromEvent(event);
+  const stepFuncContext = readStepFunctionContextFromEvent(event);
+  if (stepFuncContext) {
+    try {
+      addStepFunctionContextToXray(stepFuncContext);
+    } catch (error) {
+      logError("couldn't add step function metadata to xray", { innerError: error });
+    }
+  }
   if (trace !== undefined) {
     try {
       addTraceContextToXray(trace);
@@ -54,6 +71,12 @@ export function addTraceContextToXray(traceContext: TraceContext) {
   });
 }
 
+export function addStepFunctionContextToXray(context: StepFunctionContext) {
+  captureFunc(xraySubsegmentName, (segment) => {
+    segment.addMetadata(xrayBaggageSubsegmentKey, context, xraySubsegmentNamespace);
+  });
+}
+
 export function readTraceFromEvent(event: any): TraceContext | undefined {
   if (typeof event !== "object") {
     return;
@@ -94,6 +117,7 @@ export function readTraceFromEvent(event: any): TraceContext | undefined {
 export function readTraceContextFromXray() {
   try {
     const segment = getSegment();
+    logDebug(`Setting X-Ray parent trace to segment with ${segment.id} and trace ${segment.trace_id}`);
     const traceHeader = {
       parentID: segment.id,
       sampled: segment.notTraced ? 0 : 1,
@@ -106,6 +130,55 @@ export function readTraceContextFromXray() {
   return undefined;
 }
 
+export function readStepFunctionContextFromEvent(event: any): StepFunctionContext | undefined {
+  if (typeof event !== "object") {
+    return;
+  }
+  const { dd } = event;
+  if (typeof dd !== "object") {
+    return;
+  }
+  const execution = dd.Execution;
+  if (typeof execution !== "object") {
+    return;
+  }
+  const executionID = execution.Name;
+  if (typeof executionID !== "string") {
+    return;
+  }
+  const state = dd.State;
+  if (typeof state !== "object") {
+    return;
+  }
+  const retryCount = state.RetryCount;
+  if (typeof retryCount !== "number") {
+    return;
+  }
+  const stepName = state.Name;
+  if (typeof stepName !== "string") {
+    return;
+  }
+  const stateMachine = dd.StateMachine;
+  if (typeof stateMachine !== "object") {
+    return;
+  }
+  const stateMachineArn = stateMachine.Id;
+  if (typeof stateMachineArn !== "string") {
+    return;
+  }
+  const stateMachineName = stateMachine.Name;
+  if (typeof stateMachineName !== "string") {
+    return;
+  }
+  return {
+    "step_function.execution_id": executionID,
+    "step_function.retry_count": retryCount,
+    "step_function.state_machine_arn": stateMachineArn,
+    "step_function.state_machine_name": stateMachineName,
+    "step_function.step_name": stepName,
+  };
+}
+
 export function convertTraceContext(traceHeader: XRayTraceHeader): TraceContext | undefined {
   const sampleMode = convertToSampleMode(traceHeader.sampled);
   const traceID = convertToAPMTraceID(traceHeader.traceID);

src/trace/listener.ts

@@ -1,7 +1,7 @@
 import { Context } from "aws-lambda";
 import Tracer, { SpanContext, SpanOptions, TraceOptions } from "dd-trace";
 
-import { extractTraceContext } from "./context";
+import { extractTraceContext, readStepFunctionContextFromEvent, StepFunctionContext } from "./context";
 import { patchHttp, unpatchHttp } from "./patch-http";
 import { TraceContextService } from "./trace-context-service";
 
@@ -18,6 +18,7 @@ export interface TraceConfig {
 export class TraceListener {
   private contextService = new TraceContextService();
   private context?: Context;
+  private stepFunctionContext?: StepFunctionContext;
 
   public get currentTraceHeaders() {
     return this.contextService.currentTraceHeaders;
@@ -30,8 +31,8 @@ export class TraceListener {
       patchHttp(this.contextService);
     }
     this.context = context;
-
     this.contextService.rootTraceContext = extractTraceContext(event);
+    this.stepFunctionContext = readStepFunctionContextFromEvent(event);
   }
 
   public async onCompleteInvocation() {
@@ -51,6 +52,12 @@ export class TraceListener {
         request_id: this.context.awsRequestId,
         resource_names: this.context.functionName,
       };
+      if (this.stepFunctionContext) {
+        options.tags = {
+          ...options.tags,
+          ...this.stepFunctionContext,
+        };
+      }
     }
 
     if (spanContext !== null) {