Skip to content

Commit 6b94ed5

Browse files
DarcyRaynerDDDarcyRaynerDD
authored
Merge pull request #29 from DataDog/darcy.rayner/step-function-support-poc
Darcy.rayner/step function support poc
2 parents 5b2f1bb + c62b5bd commit 6b94ed5

File tree

6 files changed

+274
-8
lines changed

6 files changed

+274
-8
lines changed

‎package.json‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,8 @@
5050
},
5151
"collectCoverage": true,
5252
"coverageReporters": [
53-
"lcovonly"
53+
"lcovonly",
54+
"text-summary"
5455
],
5556
"testRegex": "(src\\/).*(\\.spec\\.ts)$",
5657
"testPathIgnorePatterns": [

‎src/trace/constants.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,5 @@ export const parentIDHeader = "x-datadog-parent-id";
1010
export const samplingPriorityHeader = "x-datadog-sampling-priority";
1111
export const xraySubsegmentName = "datadog-metadata";
1212
export const xraySubsegmentKey = "trace";
13+
export const xrayBaggageSubsegmentKey = "root_span_metadata";
1314
export const xraySubsegmentNamespace = "datadog";

‎src/trace/context.spec.ts‎

Lines changed: 186 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { LogLevel, setLogLevel } from "../utils";
2-
import { SampleMode } from "./constants";
2+
import { SampleMode, xrayBaggageSubsegmentKey, xraySubsegmentNamespace } from "./constants";
33
import {
44
convertToAPMParentID,
55
convertToAPMTraceID,
@@ -8,14 +8,19 @@ import {
88
extractTraceContext,
99
readTraceContextFromXray,
1010
readTraceFromEvent,
11+
readStepFunctionContextFromEvent,
1112
} from "./context";
1213

1314
let currentSegment: any;
1415

1516
jest.mock("aws-xray-sdk-core", () => {
1617
return {
17-
captureFunc: () => {
18-
throw Error("Unimplemented");
18+
captureFunc: (subsegmentName: string, callback: (segment: any) => void) => {
19+
if (currentSegment) {
20+
callback(currentSegment);
21+
} else {
22+
throw Error("Unimplemented");
23+
}
1924
},
2025
getSegment: () => {
2126
if (currentSegment === undefined) {
@@ -219,6 +224,137 @@ describe("readTraceFromEvent", () => {
219224
expect(result).toBeUndefined();
220225
});
221226
});
227+
228+
describe("readStepFunctionContextFromEvent", () => {
229+
const stepFunctionEvent = {
230+
dd: {
231+
Execution: {
232+
Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
233+
StartTime: "2019-09-30T20:28:24.236Z",
234+
},
235+
State: {
236+
Name: "step-one",
237+
RetryCount: 2,
238+
},
239+
StateMachine: {
240+
Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
241+
Name: "my-state-machine",
242+
},
243+
},
244+
} as const;
245+
it("reads a trace from an execution id", () => {
246+
const result = readStepFunctionContextFromEvent(stepFunctionEvent);
247+
expect(result).toEqual({
248+
"step_function.execution_id": "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
249+
"step_function.retry_count": 2,
250+
"step_function.state_machine_arn":
251+
"arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
252+
"step_function.state_machine_name": "my-state-machine",
253+
"step_function.step_name": "step-one",
254+
});
255+
});
256+
it("returns undefined when event isn't an object", () => {
257+
const result = readStepFunctionContextFromEvent("event");
258+
expect(result).toBeUndefined();
259+
});
260+
it("returns undefined when event is missing datadogContext property", () => {
261+
const result = readStepFunctionContextFromEvent({});
262+
expect(result).toBeUndefined();
263+
});
264+
it("returns undefined when datadogContext is missing Execution property", () => {
265+
const result = readStepFunctionContextFromEvent({
266+
dd: {},
267+
});
268+
expect(result).toBeUndefined();
269+
});
270+
it("returns undefined when Execution is missing Name field", () => {
271+
const result = readStepFunctionContextFromEvent({
272+
dd: {
273+
...stepFunctionEvent.dd,
274+
Execution: {},
275+
},
276+
});
277+
expect(result).toBeUndefined();
278+
});
279+
it("returns undefined when Name isn't a string", () => {
280+
const result = readStepFunctionContextFromEvent({
281+
dd: {
282+
...stepFunctionEvent.dd,
283+
Execution: {
284+
Name: 12345,
285+
},
286+
},
287+
});
288+
expect(result).toBeUndefined();
289+
});
290+
it("returns undefined when State isn't defined", () => {
291+
const result = readStepFunctionContextFromEvent({
292+
dd: {
293+
...stepFunctionEvent.dd,
294+
State: undefined,
295+
},
296+
});
297+
expect(result).toBeUndefined();
298+
});
299+
it("returns undefined when try retry count isn't a number", () => {
300+
const result = readStepFunctionContextFromEvent({
301+
dd: {
302+
...stepFunctionEvent.dd,
303+
State: {
304+
...stepFunctionEvent.dd.State,
305+
RetryCount: "1",
306+
},
307+
},
308+
});
309+
expect(result).toBeUndefined();
310+
});
311+
it("returns undefined when try step name isn't a string", () => {
312+
const result = readStepFunctionContextFromEvent({
313+
dd: {
314+
...stepFunctionEvent.dd,
315+
State: {
316+
...stepFunctionEvent.dd.State,
317+
Name: 1,
318+
},
319+
},
320+
});
321+
expect(result).toBeUndefined();
322+
});
323+
it("returns undefined when StateMachine is undefined", () => {
324+
const result = readStepFunctionContextFromEvent({
325+
dd: {
326+
...stepFunctionEvent.dd,
327+
StateMachine: undefined,
328+
},
329+
});
330+
expect(result).toBeUndefined();
331+
});
332+
it("returns undefined when StateMachineId isn't a string", () => {
333+
const result = readStepFunctionContextFromEvent({
334+
dd: {
335+
...stepFunctionEvent.dd,
336+
StateMachine: {
337+
...stepFunctionEvent.dd.StateMachine,
338+
Id: 1,
339+
},
340+
},
341+
});
342+
expect(result).toBeUndefined();
343+
});
344+
it("returns undefined when StateMachineName isn't a string", () => {
345+
const result = readStepFunctionContextFromEvent({
346+
dd: {
347+
...stepFunctionEvent.dd,
348+
StateMachine: {
349+
...stepFunctionEvent.dd.StateMachine,
350+
Name: 1,
351+
},
352+
},
353+
});
354+
expect(result).toBeUndefined();
355+
});
356+
});
357+
222358
describe("extractTraceContext", () => {
223359
it("returns trace read from header as highest priority", () => {
224360
currentSegment = {
@@ -252,4 +388,51 @@ describe("extractTraceContext", () => {
252388
traceID: "4110911582297405557",
253389
});
254390
});
391+
it("returns trace read from env if no headers present", () => {
392+
currentSegment = {
393+
id: "0b11cc4230d3e09e",
394+
trace_id: "1-5ce31dc2-2c779014b90ce44db5e03875",
395+
};
396+
397+
const result = extractTraceContext({});
398+
expect(result).toEqual({
399+
parentID: "797643193680388254",
400+
sampleMode: SampleMode.USER_KEEP,
401+
traceID: "4110911582297405557",
402+
});
403+
});
404+
405+
it("adds step function metadata to xray", () => {
406+
const stepFunctionEvent = {
407+
dd: {
408+
Execution: {
409+
Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
410+
StartTime: "2019-09-30T20:28:24.236Z",
411+
},
412+
State: {
413+
Name: "step-one",
414+
RetryCount: 2,
415+
},
416+
StateMachine: {
417+
Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
418+
Name: "my-state-machine",
419+
},
420+
},
421+
} as const;
422+
const addMetadata = jest.fn();
423+
currentSegment = { addMetadata };
424+
extractTraceContext(stepFunctionEvent);
425+
expect(addMetadata).toHaveBeenCalledWith(
426+
xrayBaggageSubsegmentKey,
427+
{
428+
"step_function.execution_id": "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
429+
"step_function.retry_count": 2,
430+
"step_function.state_machine_arn":
431+
"arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
432+
"step_function.state_machine_name": "my-state-machine",
433+
"step_function.step_name": "step-one",
434+
},
435+
xraySubsegmentNamespace,
436+
);
437+
});
255438
});

‎src/trace/context.ts‎

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
import { captureFunc, getSegment } from "aws-xray-sdk-core";
22
import { BigNumber } from "bignumber.js";
33

4-
import { logError } from "../utils";
4+
import { logDebug, logError } from "../utils";
55
import {
66
parentIDHeader,
77
SampleMode,
88
samplingPriorityHeader,
99
traceIDHeader,
10+
xrayBaggageSubsegmentKey,
1011
xraySubsegmentKey,
1112
xraySubsegmentName,
1213
xraySubsegmentNamespace,
@@ -24,12 +25,28 @@ export interface TraceContext {
2425
sampleMode: SampleMode;
2526
}
2627

28+
export interface StepFunctionContext {
29+
"step_function.retry_count": number;
30+
"step_function.execution_id": string;
31+
"step_function.state_machine_name": string;
32+
"step_function.state_machine_arn": string;
33+
"step_function.step_name": string;
34+
}
35+
2736
/**
2837
* Reads the trace context from either an incoming lambda event, or the current xray segment.
2938
* @param event An incoming lambda event. This must have incoming trace headers in order to be read.
3039
*/
3140
export function extractTraceContext(event: any) {
3241
const trace = readTraceFromEvent(event);
42+
const stepFuncContext = readStepFunctionContextFromEvent(event);
43+
if (stepFuncContext) {
44+
try {
45+
addStepFunctionContextToXray(stepFuncContext);
46+
} catch (error) {
47+
logError("couldn't add step function metadata to xray", { innerError: error });
48+
}
49+
}
3350
if (trace !== undefined) {
3451
try {
3552
addTraceContextToXray(trace);
@@ -54,6 +71,12 @@ export function addTraceContextToXray(traceContext: TraceContext) {
5471
});
5572
}
5673

74+
export function addStepFunctionContextToXray(context: StepFunctionContext) {
75+
captureFunc(xraySubsegmentName, (segment) => {
76+
segment.addMetadata(xrayBaggageSubsegmentKey, context, xraySubsegmentNamespace);
77+
});
78+
}
79+
5780
export function readTraceFromEvent(event: any): TraceContext | undefined {
5881
if (typeof event !== "object") {
5982
return;
@@ -94,6 +117,7 @@ export function readTraceFromEvent(event: any): TraceContext | undefined {
94117
export function readTraceContextFromXray() {
95118
try {
96119
const segment = getSegment();
120+
logDebug(`Setting X-Ray parent trace to segment with ${segment.id} and trace ${segment.trace_id}`);
97121
const traceHeader = {
98122
parentID: segment.id,
99123
sampled: segment.notTraced ? 0 : 1,
@@ -106,6 +130,55 @@ export function readTraceContextFromXray() {
106130
return undefined;
107131
}
108132

133+
export function readStepFunctionContextFromEvent(event: any): StepFunctionContext | undefined {
134+
if (typeof event !== "object") {
135+
return;
136+
}
137+
const { dd } = event;
138+
if (typeof dd !== "object") {
139+
return;
140+
}
141+
const execution = dd.Execution;
142+
if (typeof execution !== "object") {
143+
return;
144+
}
145+
const executionID = execution.Name;
146+
if (typeof executionID !== "string") {
147+
return;
148+
}
149+
const state = dd.State;
150+
if (typeof state !== "object") {
151+
return;
152+
}
153+
const retryCount = state.RetryCount;
154+
if (typeof retryCount !== "number") {
155+
return;
156+
}
157+
const stepName = state.Name;
158+
if (typeof stepName !== "string") {
159+
return;
160+
}
161+
const stateMachine = dd.StateMachine;
162+
if (typeof stateMachine !== "object") {
163+
return;
164+
}
165+
const stateMachineArn = stateMachine.Id;
166+
if (typeof stateMachineArn !== "string") {
167+
return;
168+
}
169+
const stateMachineName = stateMachine.Name;
170+
if (typeof stateMachineName !== "string") {
171+
return;
172+
}
173+
return {
174+
"step_function.execution_id": executionID,
175+
"step_function.retry_count": retryCount,
176+
"step_function.state_machine_arn": stateMachineArn,
177+
"step_function.state_machine_name": stateMachineName,
178+
"step_function.step_name": stepName,
179+
};
180+
}
181+
109182
export function convertTraceContext(traceHeader: XRayTraceHeader): TraceContext | undefined {
110183
const sampleMode = convertToSampleMode(traceHeader.sampled);
111184
const traceID = convertToAPMTraceID(traceHeader.traceID);

‎src/trace/listener.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { Context } from "aws-lambda";
22
import Tracer, { SpanContext, SpanOptions, TraceOptions } from "dd-trace";
33

4-
import { extractTraceContext } from "./context";
4+
import { extractTraceContext, readStepFunctionContextFromEvent, StepFunctionContext } from "./context";
55
import { patchHttp, unpatchHttp } from "./patch-http";
66
import { TraceContextService } from "./trace-context-service";
77

@@ -18,6 +18,7 @@ export interface TraceConfig {
1818
export class TraceListener {
1919
private contextService = new TraceContextService();
2020
private context?: Context;
21+
private stepFunctionContext?: StepFunctionContext;
2122

2223
public get currentTraceHeaders() {
2324
return this.contextService.currentTraceHeaders;
@@ -30,8 +31,8 @@ export class TraceListener {
3031
patchHttp(this.contextService);
3132
}
3233
this.context = context;
33-
3434
this.contextService.rootTraceContext = extractTraceContext(event);
35+
this.stepFunctionContext = readStepFunctionContextFromEvent(event);
3536
}
3637

3738
public async onCompleteInvocation() {
@@ -51,6 +52,12 @@ export class TraceListener {
5152
request_id: this.context.awsRequestId,
5253
resource_names: this.context.functionName,
5354
};
55+
if (this.stepFunctionContext) {
56+
options.tags = {
57+
...options.tags,
58+
...this.stepFunctionContext,
59+
};
60+
}
5461
}
5562

5663
if (spanContext !== null) {

0 commit comments

Comments
 (0)