Skip to content

Commit

Permalink
Update Step Functions Parent ID Generation (#622)
Browse files Browse the repository at this point in the history 
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`
  • Loading branch information
@avedmala
avedmala authored Feb 26, 2025
1 parent b73e44d commit c379ff8
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 16 deletions.
4 changes: 2 additions & 2 deletions src/trace/context/extractor.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -690,7 +690,7 @@ describe("TraceContextExtractor", () => {
expect(traceContext).not.toBeNull();

expect(traceContext?.toTraceId()).toBe("1139193989631387307");
expect(traceContext?.toSpanId()).toBe("5892738536804826142");
expect(traceContext?.toSpanId()).toBe("7747304477664363642");
expect(traceContext?.sampleMode()).toBe("1");
expect(traceContext?.source).toBe("event");
});
Expand Down Expand Up @@ -1092,7 +1092,7 @@ describe("TraceContextExtractor", () => {

const sentMessage = sentSegment.toString();
expect(sentMessage).toEqual(
'{"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"}}}}',
'{"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"}}}}',
);
});

Expand Down
2 changes: 1 addition & 1 deletion src/trace/context/extractors/step-function.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ describe("StepFunctionEventTraceExtractor", () => {
expect(traceContext).not.toBeNull();

expect(traceContext?.toTraceId()).toBe("435175499815315247");
expect(traceContext?.toSpanId()).toBe("5063839446130725204");
expect(traceContext?.toSpanId()).toBe("8782364156266188026");
expect(traceContext?.sampleMode()).toBe("1");
expect(traceContext?.source).toBe("event");
});
Expand Down
11 changes: 7 additions & 4 deletions src/trace/step-function-service.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,7 @@ describe("StepFunctionContextService", () => {
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",
});
Expand All @@ -163,6 +164,7 @@ describe("StepFunctionContextService", () => {
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",
root_execution_id:
Expand All @@ -179,6 +181,7 @@ describe("StepFunctionContextService", () => {
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",
trace_id: "10593586103637578129",
Expand All @@ -203,7 +206,7 @@ describe("StepFunctionContextService", () => {
expect(spanContext).not.toBeNull();

expect(spanContext?.toTraceId()).toBe("1139193989631387307");
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
expect(spanContext?.toSpanId()).toBe("7747304477664363642");
expect(spanContext?.sampleMode()).toBe("1");
expect(spanContext?.source).toBe("event");
});
Expand All @@ -218,7 +221,7 @@ describe("StepFunctionContextService", () => {
expect(spanContext).not.toBeNull();

expect(spanContext?.toTraceId()).toBe("8676990472248253142");
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
expect(spanContext?.toSpanId()).toBe("7747304477664363642");
expect(spanContext?.sampleMode()).toBe("1");
expect(spanContext?.source).toBe("event");
});
Expand All @@ -233,7 +236,7 @@ describe("StepFunctionContextService", () => {
expect(spanContext).not.toBeNull();

expect(spanContext?.toTraceId()).toBe("10593586103637578129");
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
expect(spanContext?.toSpanId()).toBe("7747304477664363642");
expect(spanContext?.sampleMode()).toBe("1");
expect(spanContext?.source).toBe("event");
});
Expand All @@ -258,7 +261,7 @@ describe("StepFunctionContextService", () => {
expect(spanContext).not.toBeNull();

expect(spanContext?.toTraceId()).toBe("1139193989631387307");
expect(spanContext?.toSpanId()).toBe("5892738536804826142");
expect(spanContext?.toSpanId()).toBe("7747304477664363642");
expect(spanContext?.sampleMode()).toBe("1");
expect(spanContext?.source).toBe("event");
});
Expand Down
30 changes: 21 additions & 9 deletions src/trace/step-function-service.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import { Sha256 } from "@aws-crypto/sha256-js";
interface NestedStepFunctionContext {
execution_id: string;
redrive_count: string;
retry_count: string;
state_entered_time: string;
state_name: string;
root_execution_id: string;
Expand All @@ -15,6 +16,7 @@ interface NestedStepFunctionContext {
interface LambdaRootStepFunctionContext {
execution_id: string;
redrive_count: string;
retry_count: string;
state_entered_time: string;
state_name: string;
trace_id: string;
Expand All @@ -25,6 +27,7 @@ interface LambdaRootStepFunctionContext {
interface LegacyStepFunctionContext {
execution_id: string;
redrive_count: string;
retry_count: string;
state_entered_time: string;
state_name: string;
}
Expand Down Expand Up @@ -91,13 +94,14 @@ export class StepFunctionContextService {
// Extract the common context variables
const stateMachineContext = this.extractStateMachineContext(event);
if (stateMachineContext === null) return;
const { execution_id, redrive_count, state_entered_time, state_name } = stateMachineContext;
const { execution_id, redrive_count, retry_count, state_entered_time, state_name } = stateMachineContext;

if (typeof event["serverless-version"] === "string" && event["serverless-version"] === "v1") {
if (typeof event.RootExecutionId === "string") {
this.context = {
execution_id,
redrive_count,
retry_count,
state_entered_time,
state_name,
root_execution_id: event.RootExecutionId,
Expand All @@ -107,6 +111,7 @@ export class StepFunctionContextService {
this.context = {
execution_id,
redrive_count,
retry_count,
state_entered_time,
state_name,
trace_id: event["x-datadog-trace-id"],
Expand All @@ -115,7 +120,13 @@ export class StepFunctionContextService {
} as LambdaRootStepFunctionContext;
}
} else {
this.context = { execution_id, redrive_count, state_entered_time, state_name } as LegacyStepFunctionContext;
this.context = {
execution_id,
redrive_count,
retry_count,
state_entered_time,
state_name,
} as LegacyStepFunctionContext;
}
}

Expand All @@ -139,17 +150,16 @@ export class StepFunctionContextService {
return null;
}

const redrivePostfix = this.context.redrive_count === "0" ? "" : `#${this.context.redrive_count}`;
const countsSuffix =
this.context.retry_count !== "0" || this.context.redrive_count !== "0"
? `#${this.context.retry_count}#${this.context.redrive_count}`
: "";

const parentId = this.deterministicSha256HashToBigIntString(
this.context.execution_id +
"#" +
this.context.state_name +
"#" +
this.context.state_entered_time +
redrivePostfix,
`${this.context.execution_id}#${this.context.state_name}#${this.context.state_entered_time}${countsSuffix}`,
PARENT_ID,
);

const sampleMode = SampleMode.AUTO_KEEP;

try {
Expand Down Expand Up @@ -209,13 +219,15 @@ export class StepFunctionContextService {
private extractStateMachineContext(event: any): {
execution_id: string;
redrive_count: string;
retry_count: string;
state_entered_time: string;
state_name: string;
} | null {
if (this.isValidContextObject(event)) {
return {
execution_id: event.Execution.Id,
redrive_count: (event.Execution.RedriveCount ?? "0").toString(),
retry_count: (event.State.RetryCount ?? "0").toString(),
state_entered_time: event.State.EnteredTime,
state_name: event.State.Name,
};
Expand Down

0 comments on commit c379ff8

Please sign in to comment.