Extract Lambda Trace Context from the Event's Step Functions Context (#331)

Author: kimi-p

package.json

@@ -39,7 +39,8 @@
     "hot-shots": "8.5.0",
     "promise-retry": "^2.0.1",
     "serialize-error": "^8.1.0",
-    "shimmer": "^1.2.1"
+    "shimmer": "^1.2.1",
+    "ts-md5": "1.3.1"
   },
   "jest": {
     "verbose": true,

src/trace/context.spec.ts

@@ -22,6 +22,11 @@ import {
   readTraceFromSQSEvent,
   readTraceFromHTTPEvent,
   readTraceFromLambdaContext,
+  hexToBinary,
+  deterministicMd5HashInBinary,
+  deterministicMd5HashToBigIntString,
+  readTraceFromStepFunctionsContext,
+  StepFunctionContext,
 } from "./context";
 
 let sentSegment: any;
@@ -1052,6 +1057,38 @@ describe("extractTraceContext", () => {
     expect(sentSegment).toBeUndefined();
   });
 
+  it("returns trace read from step functions event with the extractor as the highest priority", () => {
+    const stepFunctionEvent = {
+      MyInput: "MyValue",
+      Execution: {
+        Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
+        Input: {
+          MyInput: "MyValue",
+        },
+        Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
+        RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
+        StartTime: "2022-12-08T21:08:17.924Z",
+      },
+      State: {
+        Name: "step-one",
+        EnteredTime: "2022-12-08T21:08:19.224Z",
+        RetryCount: 2,
+      },
+      StateMachine: {
+        Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
+        Name: "my-state-machine",
+      },
+    };
+
+    const result = extractTraceContext(stepFunctionEvent, {} as Context, undefined);
+    expect(result).toEqual({
+      parentID: "4602916161841036335",
+      sampleMode: 1,
+      traceID: "947965466153612645",
+      source: "event",
+    });
+  });
+
   it("skips adding datadog metadata to x-ray when x-ray trace isn't sampled", () => {
     jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
     process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0";
@@ -1109,3 +1146,83 @@ describe("extractTraceContext", () => {
     `);
   });
 });
+
+describe.each([
+  ["0", "0000"],
+  ["1", "0001"],
+  ["2", "0010"],
+  ["3", "0011"],
+  ["4", "0100"],
+  ["5", "0101"],
+  ["6", "0110"],
+  ["7", "0111"],
+  ["8", "1000"],
+  ["9", "1001"],
+  ["a", "1010"],
+  ["b", "1011"],
+  ["c", "1100"],
+  ["d", "1101"],
+  ["e", "1110"],
+  ["f", "1111"],
+])(`test hexToBinary`, (hex, expected) => {
+  test(`${hex} to binary returns ${expected}`, () => {
+    expect(hexToBinary(hex)).toBe(expected);
+  });
+});
+
+describe("test_deterministicMd5HashInBinary", () => {
+  it("test same hashing is generated as logs-backend for a random string", () => {
+    const actual = deterministicMd5HashInBinary("some_testing_random_string");
+    expect(actual).toEqual("0001111100111110001000110110011110010111000110001001001111110001");
+  });
+
+  it("test same hashing is generated as logs-backend for an execution id", () => {
+    const actual = deterministicMd5HashInBinary(
+      "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d041f4",
+    );
+    expect(actual).toEqual("0010010000101100100000101011111101111100110110001110111100111101");
+  });
+
+  it("test same hashing is generated as logs-backend for another execution id", () => {
+    const actual = deterministicMd5HashInBinary(
+      "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111",
+    );
+    expect(actual).toEqual("0010001100110000011011011111010000100111100000110000100100101010");
+  });
+
+  it("test same hashing is generated as logs-backend for execution id # state name # entered time", () => {
+    const actual = deterministicMd5HashInBinary(
+      "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111#step-one#2022-12-08T21:08:19.224Z",
+    );
+    expect(actual).toEqual("0110111110000000010011011001111101110011100111000000011010100001");
+  });
+
+  it("test hashing different strings would generate different hashes", () => {
+    const times = 20;
+    for (let i = 0; i < times; i++) {
+      for (let j = i + 1; j < times; j++) {
+        expect(deterministicMd5HashInBinary(i.toString())).not.toMatch(deterministicMd5HashInBinary(j.toString()));
+      }
+    }
+  });
+
+  it("test always leading with 0", () => {
+    for (let i = 0; i < 20; i++) {
+      expect(deterministicMd5HashInBinary(i.toString()).substring(0, 1)).toMatch("0");
+    }
+  });
+});
+
+describe("test_deterministicMd5HashToBigIntString", () => {
+  it("test same hashing number is generated as logs-backend for a random string", () => {
+    const actual = deterministicMd5HashToBigIntString("some_testing_random_string");
+    expect(actual).toEqual("2251275791555400689");
+  });
+
+  it("test same hashing number is generated as logs-backend for execution id # state name # entered time", () => {
+    const actual = deterministicMd5HashToBigIntString(
+      "arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111#step-one#2022-12-08T21:08:19.224Z",
+    );
+    expect(actual).toEqual("8034507082463708833");
+  });
+});

src/trace/context.ts

@@ -28,6 +28,7 @@ import {
 } from "./constants";
 import { TraceExtractor } from "./listener";
 import { eventSubTypes, parseEventSourceSubType } from "./trigger";
+import { Md5 } from "ts-md5";
 
 export interface XRayTraceHeader {
   traceID: string;
@@ -55,6 +56,48 @@ export interface StepFunctionContext {
   "step_function.state_retry_count": number;
 }
 
+export function readTraceFromStepFunctionsContext(stepFunctionContext: StepFunctionContext): TraceContext | undefined {
+  return {
+    traceID: deterministicMd5HashToBigIntString(stepFunctionContext["step_function.execution_id"]),
+    parentID: deterministicMd5HashToBigIntString(
+      stepFunctionContext["step_function.execution_id"] +
+        "#" +
+        stepFunctionContext["step_function.state_name"] +
+        "#" +
+        stepFunctionContext["step_function.state_entered_time"],
+    ),
+    sampleMode: SampleMode.AUTO_KEEP.valueOf(),
+    source: Source.Event,
+  };
+}
+
+export function hexToBinary(hex: string) {
+  // convert hex to binary and padding with 0 in the front to fill 128 bits
+  return parseInt(hex, 16).toString(2).padStart(4, "0");
+}
+
+export function deterministicMd5HashInBinary(s: string): string {
+  // Md5 here is used here because we don't need a cryptographically secure hashing method but to generate the same trace/span ids as the backend does
+  const hex = Md5.hashStr(s);
+
+  let binary = "";
+  for (let i = 0; i < hex.length; i++) {
+    const ch = hex.charAt(i);
+    binary = binary + hexToBinary(ch);
+  }
+
+  const res = "0" + binary.substring(1, 64);
+  if (res === "0".repeat(64)) {
+    return "1";
+  }
+  return res;
+}
+
+export function deterministicMd5HashToBigIntString(s: string): string {
+  const binaryString = deterministicMd5HashInBinary(s);
+  return BigInt("0b" + binaryString).toString();
+}
+
 /**
  * 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.
@@ -95,6 +138,12 @@ export function extractTraceContext(
         logError("couldn't add step function metadata to xray", error as Error);
       }
     }
+    if (trace === undefined) {
+      trace = readTraceFromStepFunctionsContext(stepFuncContext);
+      if (trace !== undefined) {
+        return trace;
+      }
+    }
   }
 
   if (trace !== undefined) {

tsconfig.json

@@ -4,7 +4,11 @@
     // "incremental": true,                   /* Enable incremental compilation */
     "target": "es5" /* Specify ECMAScript target version: 'ES3' (default), 'ES5', 'ES2015', 'ES2016', 'ES2017', 'ES2018', 'ES2019' or 'ESNEXT'. */,
     "module": "commonjs" /* Specify module code generation: 'none', 'commonjs', 'amd', 'system', 'umd', 'es2015', or 'ESNext'. */,
-    "lib": ["es2015"] /* Specify library files to be included in the compilation. */,
+    "lib": [
+      "es2015",
+      "es2020.bigint"
+    ]
+    /* Specify library files to be included in the compilation. */,
     // "allowJs": true,                       /* Allow javascript files to be compiled. */
     // "checkJs": true,                       /* Report errors in .js files. */
     // "jsx": "preserve",                     /* Specify JSX code generation: 'preserve', 'react-native', or 'react'. */

yarn.lock

@@ -3378,6 +3378,11 @@ ts-jest@^27.0.1:
     semver "7.x"
     yargs-parser "20.x"
 
+[email protected]:
+  version "1.3.1"
+  resolved "https://registry.yarnpkg.com/ts-md5/-/ts-md5-1.3.1.tgz#f5b860c0d5241dd9bb4e909dd73991166403f511"
+  integrity sha512-DiwiXfwvcTeZ5wCE0z+2A9EseZsztaiZtGrtSaY5JOD7ekPnR/GoIVD5gXZAlK9Na9Kvpo9Waz5rW64WKAWApg==
+
 tslib@^1.13.0, tslib@^1.8.1:
   version "1.14.1"
   resolved "https://registry.yarnpkg.com/tslib/-/tslib-1.14.1.tgz#cf2d38bdc34a134bcaf1091c41f6619e2f672d00"