diff --git a/.changeset/silly-owls-travel.md b/.changeset/silly-owls-travel.md new file mode 100644 index 0000000..3552ed4 --- /dev/null +++ b/.changeset/silly-owls-travel.md @@ -0,0 +1,5 @@ +--- +"@aws/lambda-invoke-store": minor +--- + +add support for traceparent, tracestate and baggage diff --git a/README.md b/README.md index 5b91122..056cf1a 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ concurrent executions. ## Features - **Invocation Isolation**: Safely store and retrieve data within a single Lambda invocation. -- **Protected Lambda Context**: Built-in protection for Lambda execution metadata (requestId, [traceId](https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-traces)) +- **Protected Lambda Context**: Built-in protection for Lambda execution metadata (requestId, [traceId](https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-traces), [W3C Trace Context](https://www.w3.org/TR/trace-context/), [W3C Baggage](https://www.w3.org/TR/baggage/)) - **Custom Data Storage**: Store any custom data within the invocation context - **Async/Await Support**: Full support for asynchronous operations with context preservation - **Type Safety**: Complete TypeScript type definitions @@ -123,6 +123,30 @@ Convenience method to get the current [X-Ray trace ID](https://docs.aws.amazon.c const traceId = invokeStore.getXRayTraceId(); // Returns undefined if not set or outside context ``` +### invokeStore.getTraceparent() + +Convenience method to get the [W3C `traceparent`](https://www.w3.org/TR/trace-context/#traceparent-header) header value. + +```typescript +const traceparent = invokeStore.getTraceparent(); // e.g. "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" +``` + +### invokeStore.getTracestate() + +Convenience method to get the [W3C `tracestate`](https://www.w3.org/TR/trace-context/#tracestate-header) header value. + +```typescript +const tracestate = invokeStore.getTracestate(); // e.g. "congo=t61rcWkgMzE,rojo=00f067aa0ba902b7" +``` + +### invokeStore.getBaggage() + +Convenience method to get the [W3C Baggage](https://www.w3.org/TR/baggage/) header value. + +```typescript +const baggage = invokeStore.getBaggage(); // e.g. "userId=alice,serverNode=DF28" +``` + ### invokeStore.hasContext() Checks if code is currently running within an invoke context. diff --git a/src/invoke-store.global.spec.ts b/src/invoke-store.global.spec.ts index cf165f5..3880fd4 100644 --- a/src/invoke-store.global.spec.ts +++ b/src/invoke-store.global.spec.ts @@ -120,6 +120,13 @@ describe.each([ getRequestId: vi.fn().mockReturnValue("mock-request-id"), getXRayTraceId: vi.fn(), getTenantId: vi.fn().mockReturnValue("my-test-tenant-id"), + getTraceparent: vi + .fn() + .mockReturnValue( + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + ), + getTracestate: vi.fn().mockReturnValue("congo=t61rcWkgMzE"), + getBaggage: vi.fn().mockReturnValue("userId=alice,serverNode=DF28"), hasContext: vi.fn(), }; @@ -136,6 +143,13 @@ describe.each([ expect(awaitedReimportedStore).toBe(mockInstance); expect(awaitedReimportedStore.getRequestId()).toBe("mock-request-id"); expect(awaitedReimportedStore.getTenantId()).toBe("my-test-tenant-id"); + expect(awaitedReimportedStore.getTraceparent()).toBe( + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + ); + expect(awaitedReimportedStore.getTracestate()).toBe("congo=t61rcWkgMzE"); + expect(awaitedReimportedStore.getBaggage()).toBe( + "userId=alice,serverNode=DF28", + ); }); }); diff --git a/src/invoke-store.spec.ts b/src/invoke-store.spec.ts index be2c674..7677a8a 100644 --- a/src/invoke-store.spec.ts +++ b/src/invoke-store.spec.ts @@ -43,6 +43,89 @@ describe.each([ }); }); + describe("getTraceparent, getTracestate, and getBaggage", () => { + it("should return w3c tracing headers when set in context", async () => { + // WHEN + const result = await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + [InvokeStoreBase.PROTECTED_KEYS.TRACEPARENT]: + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + [InvokeStoreBase.PROTECTED_KEYS.TRACESTATE]: "congo=t61rcWkgMzE", + [InvokeStoreBase.PROTECTED_KEYS.BAGGAGE]: + "userId=alice,serverNode=DF28", + }, + () => { + return { + traceparent: invokeStore.getTraceparent(), + tracestate: invokeStore.getTracestate(), + baggage: invokeStore.getBaggage(), + }; + }, + ); + + // THEN + expect(result.traceparent).toBe( + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + ); + expect(result.tracestate).toBe("congo=t61rcWkgMzE"); + expect(result.baggage).toBe("userId=alice,serverNode=DF28"); + }); + + it("should return undefined when w3c headers are not in context", async () => { + // WHEN + const result = await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + }, + () => { + return { + traceparent: invokeStore.getTraceparent(), + tracestate: invokeStore.getTracestate(), + baggage: invokeStore.getBaggage(), + }; + }, + ); + + // THEN + expect(result.traceparent).toBeUndefined(); + expect(result.tracestate).toBeUndefined(); + expect(result.baggage).toBeUndefined(); + }); + + it("should prevent modifying w3c tracing protected fields", async () => { + // WHEN & THEN + await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + [InvokeStoreBase.PROTECTED_KEYS.TRACEPARENT]: "original", + }, + () => { + expect(() => { + invokeStore.set( + InvokeStoreBase.PROTECTED_KEYS.TRACEPARENT, + "modified", + ); + }).toThrow(/Cannot modify protected Lambda context field/); + + expect(() => { + invokeStore.set( + InvokeStoreBase.PROTECTED_KEYS.TRACESTATE, + "modified", + ); + }).toThrow(/Cannot modify protected Lambda context field/); + + expect(() => { + invokeStore.set( + InvokeStoreBase.PROTECTED_KEYS.BAGGAGE, + "modified", + ); + }).toThrow(/Cannot modify protected Lambda context field/); + }, + ); + }); + }); + describe("custom properties", () => { it("should allow setting and getting custom properties", async () => { // WHEN diff --git a/src/invoke-store.ts b/src/invoke-store.ts index 572ea13..ca49d2a 100644 --- a/src/invoke-store.ts +++ b/src/invoke-store.ts @@ -14,6 +14,9 @@ const PROTECTED_KEYS = { REQUEST_ID: Symbol.for("_AWS_LAMBDA_REQUEST_ID"), X_RAY_TRACE_ID: Symbol.for("_AWS_LAMBDA_X_RAY_TRACE_ID"), TENANT_ID: Symbol.for("_AWS_LAMBDA_TENANT_ID"), + TRACEPARENT: Symbol.for("_AWS_LAMBDA_TRACEPARENT"), + TRACESTATE: Symbol.for("_AWS_LAMBDA_TRACESTATE"), + BAGGAGE: Symbol.for("_AWS_LAMBDA_BAGGAGE"), } as const; const NO_GLOBAL_AWS_LAMBDA = ["true", "1"].includes( @@ -57,6 +60,18 @@ export abstract class InvokeStoreBase { getTenantId(): string | undefined { return this.get(PROTECTED_KEYS.TENANT_ID); } + + getTraceparent(): string | undefined { + return this.get(PROTECTED_KEYS.TRACEPARENT); + } + + getTracestate(): string | undefined { + return this.get(PROTECTED_KEYS.TRACESTATE); + } + + getBaggage(): string | undefined { + return this.get(PROTECTED_KEYS.BAGGAGE); + } } /**