Skip to content

Commit 0f1f682

Browse files
authored
Extract Lambda Trace Context from the Event's Step Functions Context (#331)
1 parent fcf19da commit 0f1f682

File tree

5 files changed

+178
-2
lines changed

5 files changed

+178
-2
lines changed

package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@
3939
"hot-shots": "8.5.0",
4040
"promise-retry": "^2.0.1",
4141
"serialize-error": "^8.1.0",
42-
"shimmer": "^1.2.1"
42+
"shimmer": "^1.2.1",
43+
"ts-md5": "1.3.1"
4344
},
4445
"jest": {
4546
"verbose": true,

src/trace/context.spec.ts

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,11 @@ import {
2222
readTraceFromSQSEvent,
2323
readTraceFromHTTPEvent,
2424
readTraceFromLambdaContext,
25+
hexToBinary,
26+
deterministicMd5HashInBinary,
27+
deterministicMd5HashToBigIntString,
28+
readTraceFromStepFunctionsContext,
29+
StepFunctionContext,
2530
} from "./context";
2631

2732
let sentSegment: any;
@@ -1052,6 +1057,38 @@ describe("extractTraceContext", () => {
10521057
expect(sentSegment).toBeUndefined();
10531058
});
10541059

1060+
it("returns trace read from step functions event with the extractor as the highest priority", () => {
1061+
const stepFunctionEvent = {
1062+
MyInput: "MyValue",
1063+
Execution: {
1064+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
1065+
Input: {
1066+
MyInput: "MyValue",
1067+
},
1068+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
1069+
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
1070+
StartTime: "2022-12-08T21:08:17.924Z",
1071+
},
1072+
State: {
1073+
Name: "step-one",
1074+
EnteredTime: "2022-12-08T21:08:19.224Z",
1075+
RetryCount: 2,
1076+
},
1077+
StateMachine: {
1078+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
1079+
Name: "my-state-machine",
1080+
},
1081+
};
1082+
1083+
const result = extractTraceContext(stepFunctionEvent, {} as Context, undefined);
1084+
expect(result).toEqual({
1085+
parentID: "4602916161841036335",
1086+
sampleMode: 1,
1087+
traceID: "947965466153612645",
1088+
source: "event",
1089+
});
1090+
});
1091+
10551092
it("skips adding datadog metadata to x-ray when x-ray trace isn't sampled", () => {
10561093
jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
10571094
process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0";
@@ -1109,3 +1146,83 @@ describe("extractTraceContext", () => {
11091146
`);
11101147
});
11111148
});
1149+
1150+
describe.each([
1151+
["0", "0000"],
1152+
["1", "0001"],
1153+
["2", "0010"],
1154+
["3", "0011"],
1155+
["4", "0100"],
1156+
["5", "0101"],
1157+
["6", "0110"],
1158+
["7", "0111"],
1159+
["8", "1000"],
1160+
["9", "1001"],
1161+
["a", "1010"],
1162+
["b", "1011"],
1163+
["c", "1100"],
1164+
["d", "1101"],
1165+
["e", "1110"],
1166+
["f", "1111"],
1167+
])(`test hexToBinary`, (hex, expected) => {
1168+
test(`${hex} to binary returns ${expected}`, () => {
1169+
expect(hexToBinary(hex)).toBe(expected);
1170+
});
1171+
});
1172+
1173+
describe("test_deterministicMd5HashInBinary", () => {
1174+
it("test same hashing is generated as logs-backend for a random string", () => {
1175+
const actual = deterministicMd5HashInBinary("some_testing_random_string");
1176+
expect(actual).toEqual("0001111100111110001000110110011110010111000110001001001111110001");
1177+
});
1178+
1179+
it("test same hashing is generated as logs-backend for an execution id", () => {
1180+
const actual = deterministicMd5HashInBinary(
1181+
"arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d041f4",
1182+
);
1183+
expect(actual).toEqual("0010010000101100100000101011111101111100110110001110111100111101");
1184+
});
1185+
1186+
it("test same hashing is generated as logs-backend for another execution id", () => {
1187+
const actual = deterministicMd5HashInBinary(
1188+
"arn:aws:states:sa-east-1:601427271234:express:DatadogStateMachine:acaf1a67-336a-e854-1599-2a627eb2dd8a:c8baf081-31f1-464d-971f-70cb17d01111",
1189+
);
1190+
expect(actual).toEqual("0010001100110000011011011111010000100111100000110000100100101010");
1191+
});
1192+
1193+
it("test same hashing is generated as logs-backend for execution id # state name # entered time", () => {
1194+
const actual = deterministicMd5HashInBinary(
1195+
"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",
1196+
);
1197+
expect(actual).toEqual("0110111110000000010011011001111101110011100111000000011010100001");
1198+
});
1199+
1200+
it("test hashing different strings would generate different hashes", () => {
1201+
const times = 20;
1202+
for (let i = 0; i < times; i++) {
1203+
for (let j = i + 1; j < times; j++) {
1204+
expect(deterministicMd5HashInBinary(i.toString())).not.toMatch(deterministicMd5HashInBinary(j.toString()));
1205+
}
1206+
}
1207+
});
1208+
1209+
it("test always leading with 0", () => {
1210+
for (let i = 0; i < 20; i++) {
1211+
expect(deterministicMd5HashInBinary(i.toString()).substring(0, 1)).toMatch("0");
1212+
}
1213+
});
1214+
});
1215+
1216+
describe("test_deterministicMd5HashToBigIntString", () => {
1217+
it("test same hashing number is generated as logs-backend for a random string", () => {
1218+
const actual = deterministicMd5HashToBigIntString("some_testing_random_string");
1219+
expect(actual).toEqual("2251275791555400689");
1220+
});
1221+
1222+
it("test same hashing number is generated as logs-backend for execution id # state name # entered time", () => {
1223+
const actual = deterministicMd5HashToBigIntString(
1224+
"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",
1225+
);
1226+
expect(actual).toEqual("8034507082463708833");
1227+
});
1228+
});

src/trace/context.ts

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import {
2828
} from "./constants";
2929
import { TraceExtractor } from "./listener";
3030
import { eventSubTypes, parseEventSourceSubType } from "./trigger";
31+
import { Md5 } from "ts-md5";
3132

3233
export interface XRayTraceHeader {
3334
traceID: string;
@@ -55,6 +56,48 @@ export interface StepFunctionContext {
5556
"step_function.state_retry_count": number;
5657
}
5758

59+
export function readTraceFromStepFunctionsContext(stepFunctionContext: StepFunctionContext): TraceContext | undefined {
60+
return {
61+
traceID: deterministicMd5HashToBigIntString(stepFunctionContext["step_function.execution_id"]),
62+
parentID: deterministicMd5HashToBigIntString(
63+
stepFunctionContext["step_function.execution_id"] +
64+
"#" +
65+
stepFunctionContext["step_function.state_name"] +
66+
"#" +
67+
stepFunctionContext["step_function.state_entered_time"],
68+
),
69+
sampleMode: SampleMode.AUTO_KEEP.valueOf(),
70+
source: Source.Event,
71+
};
72+
}
73+
74+
export function hexToBinary(hex: string) {
75+
// convert hex to binary and padding with 0 in the front to fill 128 bits
76+
return parseInt(hex, 16).toString(2).padStart(4, "0");
77+
}
78+
79+
export function deterministicMd5HashInBinary(s: string): string {
80+
// 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
81+
const hex = Md5.hashStr(s);
82+
83+
let binary = "";
84+
for (let i = 0; i < hex.length; i++) {
85+
const ch = hex.charAt(i);
86+
binary = binary + hexToBinary(ch);
87+
}
88+
89+
const res = "0" + binary.substring(1, 64);
90+
if (res === "0".repeat(64)) {
91+
return "1";
92+
}
93+
return res;
94+
}
95+
96+
export function deterministicMd5HashToBigIntString(s: string): string {
97+
const binaryString = deterministicMd5HashInBinary(s);
98+
return BigInt("0b" + binaryString).toString();
99+
}
100+
58101
/**
59102
* Reads the trace context from either an incoming lambda event, or the current xray segment.
60103
* @param event An incoming lambda event. This must have incoming trace headers in order to be read.
@@ -95,6 +138,12 @@ export function extractTraceContext(
95138
logError("couldn't add step function metadata to xray", error as Error);
96139
}
97140
}
141+
if (trace === undefined) {
142+
trace = readTraceFromStepFunctionsContext(stepFuncContext);
143+
if (trace !== undefined) {
144+
return trace;
145+
}
146+
}
98147
}
99148

100149
if (trace !== undefined) {

tsconfig.json

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,11 @@
44
// "incremental": true, /* Enable incremental compilation */
55
"target": "es5" /* Specify ECMAScript target version: 'ES3' (default), 'ES5', 'ES2015', 'ES2016', 'ES2017', 'ES2018', 'ES2019' or 'ESNEXT'. */,
66
"module": "commonjs" /* Specify module code generation: 'none', 'commonjs', 'amd', 'system', 'umd', 'es2015', or 'ESNext'. */,
7-
"lib": ["es2015"] /* Specify library files to be included in the compilation. */,
7+
"lib": [
8+
"es2015",
9+
"es2020.bigint"
10+
]
11+
/* Specify library files to be included in the compilation. */,
812
// "allowJs": true, /* Allow javascript files to be compiled. */
913
// "checkJs": true, /* Report errors in .js files. */
1014
// "jsx": "preserve", /* Specify JSX code generation: 'preserve', 'react-native', or 'react'. */

yarn.lock

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3378,6 +3378,11 @@ ts-jest@^27.0.1:
33783378
semver "7.x"
33793379
yargs-parser "20.x"
33803380

3381+
3382+
version "1.3.1"
3383+
resolved "https://registry.yarnpkg.com/ts-md5/-/ts-md5-1.3.1.tgz#f5b860c0d5241dd9bb4e909dd73991166403f511"
3384+
integrity sha512-DiwiXfwvcTeZ5wCE0z+2A9EseZsztaiZtGrtSaY5JOD7ekPnR/GoIVD5gXZAlK9Na9Kvpo9Waz5rW64WKAWApg==
3385+
33813386
tslib@^1.13.0, tslib@^1.8.1:
33823387
version "1.14.1"
33833388
resolved "https://registry.yarnpkg.com/tslib/-/tslib-1.14.1.tgz#cf2d38bdc34a134bcaf1091c41f6619e2f672d00"

0 commit comments

Comments
 (0)