Skip to content

Commit c379ff8

Browse files
avedmalaavedmala
authored
Update Step Functions Parent ID Generation (#622)
Updates the parent ID generation using upstream Step Functions context object to be able to handle retries Old behavior would result in the same parent ID for all retries, this new behavior will use the `retry_count` when it's non-zero. This way we don't break backwards compatibility with old layers that don't use this value in the common case of `redrive_count:0`
1 parent b73e44d commit c379ff8

File tree

4 files changed

+31
-16
lines changed

4 files changed

+31
-16
lines changed

‎src/trace/context/extractor.spec.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -690,7 +690,7 @@ describe("TraceContextExtractor", () => {
690690
expect(traceContext).not.toBeNull();
691691

692692
expect(traceContext?.toTraceId()).toBe("1139193989631387307");
693-
expect(traceContext?.toSpanId()).toBe("5892738536804826142");
693+
expect(traceContext?.toSpanId()).toBe("7747304477664363642");
694694
expect(traceContext?.sampleMode()).toBe("1");
695695
expect(traceContext?.source).toBe("event");
696696
});
@@ -1092,7 +1092,7 @@ describe("TraceContextExtractor", () => {
10921092

10931093
const sentMessage = sentSegment.toString();
10941094
expect(sentMessage).toEqual(
1095-
'{"format": "json", "version": 1}\n{"id":"11111","trace_id":"1-5e272390-8c398be037738dc042009320","parent_id":"94ae789b969f1cc5","name":"datadog-metadata","start_time":1487076708,"end_time":1487076708,"type":"subsegment","metadata":{"datadog":{"root_span_metadata":{"execution_id":"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf","redrive_count":"0","state_entered_time":"2022-12-08T21:08:19.224Z","state_name":"step-one"}}}}',
1095+
'{"format": "json", "version": 1}\n{"id":"11111","trace_id":"1-5e272390-8c398be037738dc042009320","parent_id":"94ae789b969f1cc5","name":"datadog-metadata","start_time":1487076708,"end_time":1487076708,"type":"subsegment","metadata":{"datadog":{"root_span_metadata":{"execution_id":"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf","redrive_count":"0","retry_count":"2","state_entered_time":"2022-12-08T21:08:19.224Z","state_name":"step-one"}}}}',
10961096
);
10971097
});
10981098

‎src/trace/context/extractors/step-function.spec.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ describe("StepFunctionEventTraceExtractor", () => {
7272
expect(traceContext).not.toBeNull();
7373

7474
expect(traceContext?.toTraceId()).toBe("435175499815315247");
75-
expect(traceContext?.toSpanId()).toBe("5063839446130725204");
75+
expect(traceContext?.toSpanId()).toBe("8782364156266188026");
7676
expect(traceContext?.sampleMode()).toBe("1");
7777
expect(traceContext?.source).toBe("event");
7878
});

‎src/trace/step-function-service.spec.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ describe("StepFunctionContextService", () => {
150150
execution_id:
151151
"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
152152
redrive_count: "0",
153+
retry_count: "2",
153154
state_entered_time: "2022-12-08T21:08:19.224Z",
154155
state_name: "step-one",
155156
});
@@ -163,6 +164,7 @@ describe("StepFunctionContextService", () => {
163164
execution_id:
164165
"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
165166
redrive_count: "0",
167+
retry_count: "2",
166168
state_entered_time: "2022-12-08T21:08:19.224Z",
167169
state_name: "step-one",
168170
root_execution_id:
@@ -179,6 +181,7 @@ describe("StepFunctionContextService", () => {
179181
execution_id:
180182
"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
181183
redrive_count: "0",
184+
retry_count: "2",
182185
state_entered_time: "2022-12-08T21:08:19.224Z",
183186
state_name: "step-one",
184187
trace_id: "10593586103637578129",
@@ -203,7 +206,7 @@ describe("StepFunctionContextService", () => {
203206
expect(spanContext).not.toBeNull();
204207

205208
expect(spanContext?.toTraceId()).toBe("1139193989631387307");
206-
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
209+
expect(spanContext?.toSpanId()).toBe("7747304477664363642");
207210
expect(spanContext?.sampleMode()).toBe("1");
208211
expect(spanContext?.source).toBe("event");
209212
});
@@ -218,7 +221,7 @@ describe("StepFunctionContextService", () => {
218221
expect(spanContext).not.toBeNull();
219222

220223
expect(spanContext?.toTraceId()).toBe("8676990472248253142");
221-
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
224+
expect(spanContext?.toSpanId()).toBe("7747304477664363642");
222225
expect(spanContext?.sampleMode()).toBe("1");
223226
expect(spanContext?.source).toBe("event");
224227
});
@@ -233,7 +236,7 @@ describe("StepFunctionContextService", () => {
233236
expect(spanContext).not.toBeNull();
234237

235238
expect(spanContext?.toTraceId()).toBe("10593586103637578129");
236-
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
239+
expect(spanContext?.toSpanId()).toBe("7747304477664363642");
237240
expect(spanContext?.sampleMode()).toBe("1");
238241
expect(spanContext?.source).toBe("event");
239242
});
@@ -258,7 +261,7 @@ describe("StepFunctionContextService", () => {
258261
expect(spanContext).not.toBeNull();
259262

260263
expect(spanContext?.toTraceId()).toBe("1139193989631387307");
261-
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
264+
expect(spanContext?.toSpanId()).toBe("7747304477664363642");
262265
expect(spanContext?.sampleMode()).toBe("1");
263266
expect(spanContext?.source).toBe("event");
264267
});

‎src/trace/step-function-service.ts‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { Sha256 } from "@aws-crypto/sha256-js";
66
interface NestedStepFunctionContext {
77
execution_id: string;
88
redrive_count: string;
9+
retry_count: string;
910
state_entered_time: string;
1011
state_name: string;
1112
root_execution_id: string;
@@ -15,6 +16,7 @@ interface NestedStepFunctionContext {
1516
interface LambdaRootStepFunctionContext {
1617
execution_id: string;
1718
redrive_count: string;
19+
retry_count: string;
1820
state_entered_time: string;
1921
state_name: string;
2022
trace_id: string;
@@ -25,6 +27,7 @@ interface LambdaRootStepFunctionContext {
2527
interface LegacyStepFunctionContext {
2628
execution_id: string;
2729
redrive_count: string;
30+
retry_count: string;
2831
state_entered_time: string;
2932
state_name: string;
3033
}
@@ -91,13 +94,14 @@ export class StepFunctionContextService {
9194
// Extract the common context variables
9295
const stateMachineContext = this.extractStateMachineContext(event);
9396
if (stateMachineContext === null) return;
94-
const { execution_id, redrive_count, state_entered_time, state_name } = stateMachineContext;
97+
const { execution_id, redrive_count, retry_count, state_entered_time, state_name } = stateMachineContext;
9598

9699
if (typeof event["serverless-version"] === "string" && event["serverless-version"] === "v1") {
97100
if (typeof event.RootExecutionId === "string") {
98101
this.context = {
99102
execution_id,
100103
redrive_count,
104+
retry_count,
101105
state_entered_time,
102106
state_name,
103107
root_execution_id: event.RootExecutionId,
@@ -107,6 +111,7 @@ export class StepFunctionContextService {
107111
this.context = {
108112
execution_id,
109113
redrive_count,
114+
retry_count,
110115
state_entered_time,
111116
state_name,
112117
trace_id: event["x-datadog-trace-id"],
@@ -115,7 +120,13 @@ export class StepFunctionContextService {
115120
} as LambdaRootStepFunctionContext;
116121
}
117122
} else {
118-
this.context = { execution_id, redrive_count, state_entered_time, state_name } as LegacyStepFunctionContext;
123+
this.context = {
124+
execution_id,
125+
redrive_count,
126+
retry_count,
127+
state_entered_time,
128+
state_name,
129+
} as LegacyStepFunctionContext;
119130
}
120131
}
121132

@@ -139,17 +150,16 @@ export class StepFunctionContextService {
139150
return null;
140151
}
141152

142-
const redrivePostfix = this.context.redrive_count === "0" ? "" : `#${this.context.redrive_count}`;
153+
const countsSuffix =
154+
this.context.retry_count !== "0" || this.context.redrive_count !== "0"
155+
? `#${this.context.retry_count}#${this.context.redrive_count}`
156+
: "";
143157

144158
const parentId = this.deterministicSha256HashToBigIntString(
145-
this.context.execution_id +
146-
"#" +
147-
this.context.state_name +
148-
"#" +
149-
this.context.state_entered_time +
150-
redrivePostfix,
159+
`${this.context.execution_id}#${this.context.state_name}#${this.context.state_entered_time}${countsSuffix}`,
151160
PARENT_ID,
152161
);
162+
153163
const sampleMode = SampleMode.AUTO_KEEP;
154164

155165
try {
@@ -209,13 +219,15 @@ export class StepFunctionContextService {
209219
private extractStateMachineContext(event: any): {
210220
execution_id: string;
211221
redrive_count: string;
222+
retry_count: string;
212223
state_entered_time: string;
213224
state_name: string;
214225
} | null {
215226
if (this.isValidContextObject(event)) {
216227
return {
217228
execution_id: event.Execution.Id,
218229
redrive_count: (event.Execution.RedriveCount ?? "0").toString(),
230+
retry_count: (event.State.RetryCount ?? "0").toString(),
219231
state_entered_time: event.State.EnteredTime,
220232
state_name: event.State.Name,
221233
};

0 commit comments

Comments
 (0)