Skip to content

Commit e5f383c

Browse files
kuhekuhe
authored
feat(core/client): propagate w3c trace headers (#8156)
1 parent d0f426c commit e5f383c

4 files changed

Lines changed: 656 additions & 24 deletions

File tree

‎packages-internal/core/src/submodules/client/middleware-recursion-detection/configuration.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
import type { AbsoluteLocation, BuildHandlerOptions } from "@smithy/types";
22

33
/**
4+
* Used in conjunction with Lambda invoke store.
45
* @internal
56
*/
67
export const recursionDetectionMiddlewareOptions: BuildHandlerOptions & AbsoluteLocation = {
78
step: "build",
8-
tags: ["RECURSION_DETECTION"],
9+
tags: ["RECURSION_DETECTION", "TRACE_CONTEXT_PROPAGATION"],
910
name: "recursionDetectionMiddleware",
1011
override: true,
1112
priority: "low",

‎packages-internal/core/src/submodules/client/middleware-recursion-detection/middleware-recursion-detection.integ.spec.ts‎

Lines changed: 211 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,23 @@
1+
import { InvokeStore } from "@aws/lambda-invoke-store";
12
import { requireRequestsFrom } from "@aws-sdk/aws-util-test/src";
23
import { Lambda } from "@aws-sdk/client-lambda";
3-
import { describe, expect, test as it } from "vitest";
4+
import { afterEach, beforeEach, describe, expect, test as it, vi } from "vitest";
45

56
const ENV_LAMBDA_FUNCTION_NAME = "AWS_LAMBDA_FUNCTION_NAME";
67
const ENV_TRACE_ID = "_X_AMZN_TRACE_ID";
78

89
describe("middleware-recursion-detection", () => {
10+
const originEnv = process.env;
11+
12+
beforeEach(() => {
13+
vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue(undefined as any);
14+
});
15+
16+
afterEach(() => {
17+
vi.restoreAllMocks();
18+
process.env = originEnv;
19+
});
20+
921
describe(Lambda.name, () => {
1022
it("should create recursion detection headers", async () => {
1123
process.env[ENV_LAMBDA_FUNCTION_NAME] = "MyLambdaFunction";
@@ -27,5 +39,203 @@ describe("middleware-recursion-detection", () => {
2739

2840
expect.hasAssertions();
2941
});
42+
43+
it("should propagate W3C traceparent header from InvokeStore", async () => {
44+
const mockTraceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
45+
vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
46+
getXRayTraceId: () => undefined,
47+
getTraceparent: () => mockTraceparent,
48+
getTracestate: () => undefined,
49+
getBaggage: () => undefined,
50+
} as any);
51+
52+
const client = new Lambda({
53+
region: "us-west-2",
54+
});
55+
56+
requireRequestsFrom(client).toMatch({
57+
headers: {
58+
traceparent: new RegExp(`^${mockTraceparent.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
59+
},
60+
});
61+
62+
await client.invoke({
63+
FunctionName: "my-function",
64+
});
65+
66+
expect.hasAssertions();
67+
});
68+
69+
it("should propagate all W3C trace context headers from InvokeStore", async () => {
70+
const mockTraceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
71+
const mockTracestate = "congo=t61rcWkgMzE";
72+
const mockBaggage = "userId=alice,serverNode=DF%2028";
73+
vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
74+
getXRayTraceId: () => undefined,
75+
getTraceparent: () => mockTraceparent,
76+
getTracestate: () => mockTracestate,
77+
getBaggage: () => mockBaggage,
78+
} as any);
79+
80+
const client = new Lambda({
81+
region: "us-west-2",
82+
});
83+
84+
requireRequestsFrom(client).toMatch({
85+
headers: {
86+
traceparent: /^00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01$/,
87+
tracestate: /^congo=t61rcWkgMzE$/,
88+
baggage: /^userId=alice,serverNode=DF%2028$/,
89+
},
90+
});
91+
92+
await client.invoke({
93+
FunctionName: "my-function",
94+
});
95+
96+
expect.hasAssertions();
97+
});
98+
99+
it("should propagate both X-Amzn-Trace-Id and W3C headers when in Lambda", async () => {
100+
const mockXRayTraceId = "Root=1-abc-def;Parent=123;Sampled=1";
101+
const mockTraceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
102+
const mockTracestate = "congo=t61rcWkgMzE";
103+
process.env[ENV_LAMBDA_FUNCTION_NAME] = "MyLambdaFunction";
104+
105+
vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
106+
getXRayTraceId: () => mockXRayTraceId,
107+
getTraceparent: () => mockTraceparent,
108+
getTracestate: () => mockTracestate,
109+
getBaggage: () => undefined,
110+
} as any);
111+
112+
const client = new Lambda({
113+
region: "us-west-2",
114+
});
115+
116+
requireRequestsFrom(client).toMatch({
117+
headers: {
118+
"X-Amzn-Trace-Id": /^Root=1-abc-def;Parent=123;Sampled=1$/,
119+
traceparent: /^00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01$/,
120+
tracestate: /^congo=t61rcWkgMzE$/,
121+
},
122+
});
123+
124+
await client.invoke({
125+
FunctionName: "my-function",
126+
});
127+
128+
expect.hasAssertions();
129+
});
130+
131+
it("should NOT overwrite existing traceparent or add tracestate/baggage from InvokeStore when traceparent is already set", async () => {
132+
const existingTraceparent = "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01";
133+
const mockTraceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
134+
const mockTracestate = "congo=t61rcWkgMzE";
135+
const mockBaggage = "userId=alice,serverNode=DF%2028";
136+
vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
137+
getXRayTraceId: () => undefined,
138+
getTraceparent: () => mockTraceparent,
139+
getTracestate: () => mockTracestate,
140+
getBaggage: () => mockBaggage,
141+
} as any);
142+
143+
const client = new Lambda({
144+
region: "us-west-2",
145+
});
146+
147+
// Add middleware that sets traceparent before recursion detection runs.
148+
// recursionDetectionMiddleware runs at build step with low priority,
149+
// so we add at build step with high priority to run first.
150+
client.middlewareStack.add(
151+
(next) => async (args: any) => {
152+
args.request.headers["traceparent"] = existingTraceparent;
153+
return next(args);
154+
},
155+
{ step: "build", priority: "high", name: "testSetTraceparent", override: true }
156+
);
157+
158+
requireRequestsFrom(client).toMatch({
159+
headers: {
160+
traceparent: new RegExp(`^${existingTraceparent.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
161+
tracestate: (value: any) => expect(value).toBeUndefined(),
162+
baggage: (value: any) => expect(value).toBeUndefined(),
163+
},
164+
});
165+
166+
await client.invoke({
167+
FunctionName: "my-function",
168+
});
169+
170+
expect.hasAssertions();
171+
});
172+
173+
it("should preserve existing traceparent, tracestate, and baggage without overwriting from InvokeStore", async () => {
174+
const existingTraceparent = "00-abcdef1234567890abcdef1234567890-1234567890abcdef-01";
175+
const existingTracestate = "vendor1=opaqueValue1";
176+
const existingBaggage = "key1=value1";
177+
const mockTraceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
178+
const mockTracestate = "congo=t61rcWkgMzE";
179+
const mockBaggage = "userId=alice,serverNode=DF%2028";
180+
vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
181+
getXRayTraceId: () => undefined,
182+
getTraceparent: () => mockTraceparent,
183+
getTracestate: () => mockTracestate,
184+
getBaggage: () => mockBaggage,
185+
} as any);
186+
187+
const client = new Lambda({
188+
region: "us-west-2",
189+
});
190+
191+
client.middlewareStack.add(
192+
(next) => async (args: any) => {
193+
args.request.headers["traceparent"] = existingTraceparent;
194+
args.request.headers["tracestate"] = existingTracestate;
195+
args.request.headers["baggage"] = existingBaggage;
196+
return next(args);
197+
},
198+
{ step: "build", priority: "high", name: "testSetTraceHeaders", override: true }
199+
);
200+
201+
requireRequestsFrom(client).toMatch({
202+
headers: {
203+
traceparent: new RegExp(`^${existingTraceparent.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
204+
tracestate: new RegExp(`^${existingTracestate.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
205+
baggage: new RegExp(`^${existingBaggage.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}$`),
206+
},
207+
});
208+
209+
await client.invoke({
210+
FunctionName: "my-function",
211+
});
212+
213+
expect.hasAssertions();
214+
});
215+
216+
it("should NOT add W3C headers when InvokeStore has no traceparent", async () => {
217+
vi.spyOn(InvokeStore, "getInstanceAsync").mockResolvedValue({
218+
getXRayTraceId: () => undefined,
219+
getTraceparent: () => undefined,
220+
getTracestate: () => "congo=t61rcWkgMzE",
221+
getBaggage: () => "userId=alice",
222+
} as any);
223+
224+
const client = new Lambda({
225+
region: "us-west-2",
226+
});
227+
228+
requireRequestsFrom(client).toMatch({
229+
headers: {
230+
traceparent: (value: any) => expect(value).toBeUndefined(),
231+
},
232+
});
233+
234+
await client.invoke({
235+
FunctionName: "my-function",
236+
});
237+
238+
expect.hasAssertions();
239+
});
30240
});
31241
});

0 commit comments

Comments
 (0)