diff --git a/.changeset/changelog.mjs b/.changeset/changelog.mjs index 474802a..b42de9d 100644 --- a/.changeset/changelog.mjs +++ b/.changeset/changelog.mjs @@ -14,7 +14,7 @@ export const getDependencyReleaseLine = (changesets, dependenciesUpdated) => { .join(", ")}]:`; const updatedDepsList = dependenciesUpdated.map( - (dependency) => ` - ${dependency.name}@${dependency.newVersion}` + (dependency) => ` - ${dependency.name}@${dependency.newVersion}`, ); return [changesetLink, ...updatedDepsList].join("\n"); @@ -22,9 +22,11 @@ export const getDependencyReleaseLine = (changesets, dependenciesUpdated) => { export const getReleaseLine = (changeset, _type) => { const { commit, summary } = changeset; - const [firstLine, ...futureLines] = summary.split("\n").map((l) => l.trimRight()); + const [firstLine, ...futureLines] = summary + .split("\n") + .map((l) => l.trimRight()); return `- ${firstLine} (${getGithubCommitWithLink(commit)})${ futureLines.length > 0 ? futureLines.map((l) => ` ${l}`).join("\n") : "" }`; -}; \ No newline at end of file +}; diff --git a/.changeset/goofy-planes-tease.md b/.changeset/goofy-planes-tease.md new file mode 100644 index 0000000..1ae1b57 --- /dev/null +++ b/.changeset/goofy-planes-tease.md @@ -0,0 +1,10 @@ +--- +"@aws/lambda-invoke-store": minor +--- + +Invoke Store is now accessible via `InvokeStore.getInstanceAsync()` instead of direct instantiation + +- Lazy loads `node:async_hooks` to improve startup performance +- Selects dynamic implementation based on Lambda environment: + - Single-context implementation for standard Lambda executions + - Multi-context implementation (using AsyncLocalStorage) diff --git a/README.md b/README.md index 47f15e7..5b91122 100644 --- a/README.md +++ b/README.md @@ -33,19 +33,20 @@ export const handler = async (event, context) => { // The RIC has already initialized the InvokeStore with requestId and X-Ray traceId // Access Lambda context data - console.log(`Processing request: ${InvokeStore.getRequestId()}`); + const invokeStore = await InvokeStore.getInstanceAsync(); + console.log(`Processing request: ${invokeStore.getRequestId()}`); // Store custom data - InvokeStore.set("userId", event.userId); + invokeStore.set("userId", event.userId); // Data persists across async operations await processData(event); // Retrieve custom data - const userId = InvokeStore.get("userId"); + const userId = invokeStore.get("userId"); return { - requestId: InvokeStore.getRequestId(), + requestId: invokeStore.getRequestId(), userId, }; }; @@ -53,89 +54,96 @@ export const handler = async (event, context) => { // Context is preserved in async operations async function processData(event) { // Still has access to the same invoke context - console.log(`Processing in same context: ${InvokeStore.getRequestId()}`); + const invokeStore = await InvokeStore.getInstanceAsync(); + console.log(`Processing in same context: ${invokeStore.getRequestId()}`); // Can set additional data - InvokeStore.set("processedData", { result: "success" }); + invokeStore.set("processedData", { result: "success" }); } ``` ## API Reference -### InvokeStore.getContext() +### InvokeStore.getInstanceAsync() +First, get an instance of the InvokeStore: +```typescript +const invokeStore = await InvokeStore.getInstanceAsync(); +``` + +### invokeStore.getContext() Returns the complete current context or `undefined` if outside a context. ```typescript -const context = InvokeStore.getContext(); +const context = invokeStore.getContext(); ``` -### InvokeStore.get(key) +### invokeStore.get(key) Gets a value from the current context. ```typescript -const requestId = InvokeStore.get(InvokeStore.PROTECTED_KEYS.REQUEST_ID); -const customValue = InvokeStore.get("customKey"); +const requestId = invokeStore.get(InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID); +const customValue = invokeStore.get("customKey"); ``` -### InvokeStore.set(key, value) +### invokeStore.set(key, value) Sets a custom value in the current context. Protected Lambda fields cannot be modified. ```typescript -InvokeStore.set("userId", "user-123"); -InvokeStore.set("timestamp", Date.now()); +invokeStore.set("userId", "user-123"); +invokeStore.set("timestamp", Date.now()); // This will throw an error: -// InvokeStore.set(InvokeStore.PROTECTED_KEYS.REQUEST_ID, 'new-id'); +// invokeStore.set(InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID, 'new-id'); ``` -### InvokeStore.getRequestId() +### invokeStore.getRequestId() Convenience method to get the current request ID. ```typescript -const requestId = InvokeStore.getRequestId(); // Returns '-' if outside context +const requestId = invokeStore.getRequestId(); // Returns '-' if outside context ``` -### InvokeStore.getTenantId() +### invokeStore.getTenantId() Convenience method to get the tenant ID. ```typescript -const requestId = InvokeStore.getTenantId(); +const requestId = invokeStore.getTenantId(); ``` -### InvokeStore.getXRayTraceId() +### invokeStore.getXRayTraceId() Convenience method to get the current [X-Ray trace ID](https://docs.aws.amazon.com/xray/latest/devguide/xray-concepts.html#xray-concepts-traces). This ID is used for distributed tracing across AWS services. ```typescript -const traceId = InvokeStore.getXRayTraceId(); // Returns undefined if not set or outside context +const traceId = invokeStore.getXRayTraceId(); // Returns undefined if not set or outside context ``` -### InvokeStore.hasContext() +### invokeStore.hasContext() Checks if code is currently running within an invoke context. ```typescript -if (InvokeStore.hasContext()) { +if (invokeStore.hasContext()) { // We're inside an invoke context } ``` -### InvokeStore.run(context, fn) +### invokeStore.run(context, fn) > **Note**: This method is primarily used by the Lambda Runtime Interface Client (RIC) to initialize the context for each invocation. Lambda function developers typically don't need to call this method directly. Runs a function within an invoke context. ```typescript -InvokeStore.run( +invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-123", - [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-456", // Optional X-Ray trace ID + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-123", + [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-456", // Optional X-Ray trace ID customField: "value", // Optional custom fields }, () => { diff --git a/package.json b/package.json index 3631aa5..902b16d 100644 --- a/package.json +++ b/package.json @@ -32,7 +32,7 @@ "build": "yarn clean && yarn build:types && node ./scripts/build-rollup.js", "build:types": "tsc -p tsconfig.types.json", "clean": "rm -rf dist-types dist-cjs dist-es", - "test": "vitest run", + "test": "vitest run --reporter verbose", "test:watch": "vitest watch", "release": "yarn build && changeset publish" }, diff --git a/src/invoke-store.benchmark.ts b/src/invoke-store.benchmark.ts new file mode 100644 index 0000000..19508fc --- /dev/null +++ b/src/invoke-store.benchmark.ts @@ -0,0 +1,53 @@ +import { performance, PerformanceObserver } from "node:perf_hooks"; + +const obs = new PerformanceObserver((list) => { + const entries = list.getEntries(); + entries.forEach((entry) => { + console.log(`${entry.name}: ${entry.duration}ms`); + }); +}); +obs.observe({ entryTypes: ["measure"] }); + +async function runBenchmark() { + const iterations = 1000; + process.env["AWS_LAMBDA_BENCHMARK_MODE"] = "1"; + + performance.mark("direct-single-start"); + for (let i = 0; i < iterations; i++) { + const invokeStore = (await import("./invoke-store.js")).InvokeStore; + await invokeStore.getInstanceAsync(); + const testing = invokeStore._testing; + if (testing) { + testing.reset(); + } else { + throw "testing needs to be defined"; + } + } + performance.mark("direct-single-end"); + performance.measure( + "Direct SingleStore Creation (1000 iterations)", + "direct-single-start", + "direct-single-end", + ); + + performance.mark("direct-multi-start"); + process.env["AWS_LAMBDA_MAX_CONCURRENCY"] = "2"; + for (let i = 0; i < iterations; i++) { + const invokeStore = (await import("./invoke-store.js")).InvokeStore; + await invokeStore.getInstanceAsync(); + const testing = invokeStore._testing; + if (testing) { + testing.reset(); + } else { + throw "testing needs to be defined"; + } + } + performance.mark("direct-multi-end"); + performance.measure( + "Direct MultiStore Creation (1000 iterations)", + "direct-multi-start", + "direct-multi-end", + ); +} + +runBenchmark().catch(console.error); diff --git a/src/invoke-store.concurrency.spec.ts b/src/invoke-store.concurrency.spec.ts new file mode 100644 index 0000000..3cf771d --- /dev/null +++ b/src/invoke-store.concurrency.spec.ts @@ -0,0 +1,120 @@ +import { describe, it, expect, afterEach, beforeEach, vi } from "vitest"; +import { InvokeStore, InvokeStoreBase } from "./invoke-store.js"; + +describe("InvokeStore", async () => { + let invokeStore: InvokeStoreBase; + + beforeEach(async () => { + vi.stubEnv("AWS_LAMBDA_MAX_CONCURRENCY", "2"); + vi.useFakeTimers(); + invokeStore = await InvokeStore.getInstanceAsync(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + describe("run", () => { + it("should handle nested runs with different IDs", async () => { + // GIVEN + const traces: string[] = []; + + // WHEN + await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "outer", + }, + async () => { + traces.push(`outer-${invokeStore.getRequestId()}`); + await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "inner", + }, + async () => { + traces.push(`inner-${invokeStore.getRequestId()}`); + }, + ); + traces.push(`outer-again-${invokeStore.getRequestId()}`); + }, + ); + + // THEN + expect(traces).toEqual([ + "outer-outer", + "inner-inner", + "outer-again-outer", + ]); + }); + + it("should maintain isolation between concurrent executions", async () => { + // GIVEN + const traces: string[] = []; + + // WHEN - Simulate concurrent invocations + const isolateTasks = Promise.all([ + invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-1", + [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-1", + }, + async () => { + traces.push(`start-1-${invokeStore.getRequestId()}`); + await new Promise((resolve) => setTimeout(resolve, 10)); + traces.push(`end-1-${invokeStore.getRequestId()}`); + }, + ), + invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-2", + [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-2", + }, + async () => { + traces.push(`start-2-${invokeStore.getRequestId()}`); + await new Promise((resolve) => setTimeout(resolve, 5)); + traces.push(`end-2-${invokeStore.getRequestId()}`); + }, + ), + ]); + vi.runAllTimers(); + await isolateTasks; + + // THEN + expect(traces).toEqual([ + "start-1-request-1", + "start-2-request-2", + "end-2-request-2", + "end-1-request-1", + ]); + }); + + it("should maintain isolation across async operations", async () => { + // GIVEN + const traces: string[] = []; + + // WHEN + await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-1", + }, + async () => { + traces.push(`before-${invokeStore.getRequestId()}`); + const task = new Promise((resolve) => { + setTimeout(resolve, 1); + }).then(() => { + traces.push(`inside-${invokeStore.getRequestId()}`); + }); + vi.runAllTimers(); + await task; + traces.push(`after-${invokeStore.getRequestId()}`); + }, + ); + + // THEN + expect(traces).toEqual([ + "before-request-1", + "inside-request-1", + "after-request-1", + ]); + }); + }); +}); diff --git a/src/invoke-store.global.concurrency.spec.ts b/src/invoke-store.global.concurrency.spec.ts new file mode 100644 index 0000000..c4be9e2 --- /dev/null +++ b/src/invoke-store.global.concurrency.spec.ts @@ -0,0 +1,72 @@ +import { + describe, + it, + expect, + beforeAll, + afterAll, + beforeEach, + vi, +} from "vitest"; +import { + InvokeStoreBase, + InvokeStore as OriginalImport, +} from "./invoke-store.js"; + +describe("InvokeStore Global Singleton", () => { + const originalGlobalAwsLambda = globalThis.awslambda; + const originalEnv = process.env; + let invokeStore: InvokeStoreBase; + + beforeAll(() => { + globalThis.awslambda = originalGlobalAwsLambda; + }); + + afterAll(() => { + delete (globalThis as any).awslambda; + process.env = originalEnv; + }); + + beforeEach(async () => { + vi.stubEnv("AWS_LAMBDA_MAX_CONCURRENCY", "2"); + process.env = { ...originalEnv }; + invokeStore = await OriginalImport.getInstanceAsync(); + }); + + it("should maintain singleton behavior with dynamic imports", async () => { + // GIVEN + const testRequestId = "dynamic-import-test"; + const testTenantId = "dynamic-import-tenant-id-test"; + const testKey = "dynamic-key"; + const testValue = "dynamic-value"; + + // WHEN - Set up context with original import + await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.TENANT_ID]: testTenantId, + }, + async () => { + invokeStore.set(testKey, testValue); + + // Dynamically import the module again + const dynamicModule = await import("./invoke-store.js"); + const dynamicImport = + await dynamicModule.InvokeStore.getInstanceAsync(); + + // THEN - Dynamically imported instance should see the same context + expect(dynamicImport).toBe(invokeStore); // Same instance + expect(dynamicImport.getRequestId()).toBe(testRequestId); + expect(dynamicImport.getTenantId()).toBe(testTenantId); + expect(dynamicImport.get(testKey)).toBe(testValue); + + // WHEN - Set a new value using dynamic import + const newKey = "new-dynamic-key"; + const newValue = "new-dynamic-value"; + dynamicImport.set(newKey, newValue); + + // THEN - Original import should see the new value + expect(invokeStore.get(newKey)).toBe(newValue); + }, + ); + }); +}); diff --git a/src/invoke-store.global.spec.ts b/src/invoke-store.global.spec.ts index 251351f..cf165f5 100644 --- a/src/invoke-store.global.spec.ts +++ b/src/invoke-store.global.spec.ts @@ -7,204 +7,192 @@ import { beforeEach, vi, } from "vitest"; -import { InvokeStore as OriginalImport } from "./invoke-store.js"; - -describe("InvokeStore Global Singleton", () => { - const originalGlobalAwsLambda = globalThis.awslambda; - const originalEnv = process.env; - - beforeAll(() => { - globalThis.awslambda = originalGlobalAwsLambda; - }); - - afterAll(() => { - delete (globalThis as any).awslambda; - process.env = originalEnv; - }); - - beforeEach(() => { - process.env = { ...originalEnv }; - }); - - it("should store the instance in globalThis.awslambda", () => { - // THEN - expect(globalThis.awslambda.InvokeStore).toBeDefined(); - expect(globalThis.awslambda.InvokeStore).toBe(OriginalImport); - }); - - it("should share context between original import and global reference", async () => { - // GIVEN - const testRequestId = "shared-context-test"; - const testKey = "test-key"; - const testValue = "test-value"; - - // WHEN - Use the original import to set up context - await OriginalImport.run( - { [OriginalImport.PROTECTED_KEYS.REQUEST_ID]: testRequestId }, - () => { - OriginalImport.set(testKey, testValue); - - // THEN - Global reference should see the same context - const globalInstance = globalThis.awslambda.InvokeStore!; - expect(globalInstance.getRequestId()).toBe(testRequestId); - expect(globalInstance.get(testKey)).toBe(testValue); - } - ); - }); - - it("should maintain the same storage across different references", async () => { - // GIVEN - const globalInstance = globalThis.awslambda.InvokeStore!; - const testRequestId = "global-test"; - const testKey = "global-key"; - const testValue = "global-value"; - - // WHEN - Set context using global reference - await globalInstance.run( - { [globalInstance.PROTECTED_KEYS.REQUEST_ID]: testRequestId }, - () => { - globalInstance.set(testKey, testValue); - - // THEN - Original import should see the same context - expect(OriginalImport.getRequestId()).toBe(testRequestId); - expect(OriginalImport.get(testKey)).toBe(testValue); - } - ); - }); - - it("should maintain singleton behavior with dynamic imports", async () => { - // GIVEN - const testRequestId = "dynamic-import-test"; - const testTenantId = "dynamic-import-tenant-id-test"; - const testKey = "dynamic-key"; - const testValue = "dynamic-value"; - - // WHEN - Set up context with original import - await OriginalImport.run( - { - [OriginalImport.PROTECTED_KEYS.REQUEST_ID]: testRequestId, - [OriginalImport.PROTECTED_KEYS.TENANT_ID]: testTenantId, - }, - async () => { - OriginalImport.set(testKey, testValue); - - // Dynamically import the module again - const dynamicModule = await import("./invoke-store.js"); - const DynamicImport = dynamicModule.InvokeStore; - - // THEN - Dynamically imported instance should see the same context - expect(DynamicImport).toBe(OriginalImport); // Same instance - expect(DynamicImport.getRequestId()).toBe(testRequestId); - expect(DynamicImport.getTenantId()).toBe(testTenantId); - expect(DynamicImport.get(testKey)).toBe(testValue); - - // WHEN - Set a new value using dynamic import - const newKey = "new-dynamic-key"; - const newValue = "new-dynamic-value"; - DynamicImport.set(newKey, newValue); - - // THEN - Original import should see the new value - expect(OriginalImport.get(newKey)).toBe(newValue); +import { + InvokeStoreBase, + InvokeStore, + InvokeStore as OriginalImport, +} from "./invoke-store.js"; + +describe.each([ + { label: "multi-concurrency", isMultiConcurrent: true }, + { label: "single-concurrency", isMultiConcurrent: false }, +])("InvokeStore with %s", async ({ isMultiConcurrent }) => { + let invokeStore: InvokeStoreBase; + + describe("InvokeStore Global Singleton", () => { + const originalGlobalAwsLambda = globalThis.awslambda; + const originalEnv = process.env; + + beforeAll(() => { + globalThis.awslambda = originalGlobalAwsLambda; + }); + + afterAll(() => { + delete (globalThis as any).awslambda; + process.env = originalEnv; + }); + + beforeEach(async () => { + if (isMultiConcurrent) { + vi.stubEnv("AWS_LAMBDA_MAX_CONCURRENCY", "2"); } - ); + process.env = { ...originalEnv }; + invokeStore = await InvokeStore.getInstanceAsync(); + }); + + it("should store the instance in globalThis.awslambda", async () => { + // THEN + expect(globalThis.awslambda.InvokeStore).toBeDefined(); + expect(await globalThis.awslambda.InvokeStore).toBe( + await OriginalImport.getInstanceAsync(), + ); + }); + + it("should share context between original import and global reference", async () => { + // GIVEN + const testRequestId = "shared-context-test"; + const testKey = "test-key"; + const testValue = "test-value"; + + const originalImport = await OriginalImport.getInstanceAsync(); + + // WHEN - Use the original import to set up context + await originalImport.run( + { [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId }, + async () => { + originalImport.set(testKey, testValue); + + // THEN - Global reference should see the same context + const globalInstance = globalThis.awslambda.InvokeStore!; + expect(globalInstance.getRequestId()).toBe(testRequestId); + expect(globalInstance.get(testKey)).toBe(testValue); + }, + ); + }); + + it("should maintain the same storage across different references", async () => { + // GIVEN + const globalInstance = globalThis.awslambda.InvokeStore!; + const originalImport = await OriginalImport.getInstanceAsync(); + const testRequestId = "global-test"; + const testKey = "global-key"; + const testValue = "global-value"; + + // WHEN - Set context using global reference + await globalInstance.run( + { [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId }, + async () => { + globalInstance.set(testKey, testValue); + + // THEN - Original import should see the same context + expect(originalImport.getRequestId()).toBe(testRequestId); + expect(originalImport.get(testKey)).toBe(testValue); + }, + ); + }); }); -}); - -describe("InvokeStore Existing Instance", () => { - const originalGlobalAwsLambda = globalThis.awslambda; - beforeEach(() => { - delete (globalThis as any).awslambda; - globalThis.awslambda = {}; - - vi.resetModules(); + describe("InvokeStore Existing Instance", () => { + const originalGlobalAwsLambda = globalThis.awslambda; + + beforeEach(() => { + delete (globalThis as any).awslambda; + globalThis.awslambda = {}; + + vi.resetModules(); + }); + + afterAll(() => { + globalThis.awslambda = originalGlobalAwsLambda; + }); + + it("should use existing instance from globalThis.awslambda.InvokeStore", async () => { + // GIVEN + const mockInstance = { + PROTECTED_KEYS: { + REQUEST_ID: "_AWS_LAMBDA_REQUEST_ID", + X_RAY_TRACE_ID: "_AWS_LAMBDA_TRACE_ID", + }, + run: vi.fn(), + getContext: vi.fn(), + get: vi.fn(), + set: vi.fn(), + getRequestId: vi.fn().mockReturnValue("mock-request-id"), + getXRayTraceId: vi.fn(), + getTenantId: vi.fn().mockReturnValue("my-test-tenant-id"), + hasContext: vi.fn(), + }; + + // @ts-expect-error - mockInstance can be loosely related to original type + globalThis.awslambda.InvokeStore = mockInstance; + + // WHEN + const { InvokeStore: ReimportedStore } = await import( + "./invoke-store.js" + ); + const awaitedReimportedStore = await ReimportedStore.getInstanceAsync(); + + // THEN + expect(awaitedReimportedStore).toBe(mockInstance); + expect(awaitedReimportedStore.getRequestId()).toBe("mock-request-id"); + expect(awaitedReimportedStore.getTenantId()).toBe("my-test-tenant-id"); + }); }); - afterAll(() => { - globalThis.awslambda = originalGlobalAwsLambda; - }); - - it("should use existing instance from globalThis.awslambda.InvokeStore", async () => { - // GIVEN - const mockInstance = { - PROTECTED_KEYS: { - REQUEST_ID: "_AWS_LAMBDA_REQUEST_ID", - X_RAY_TRACE_ID: "_AWS_LAMBDA_TRACE_ID", - }, - run: vi.fn(), - getContext: vi.fn(), - get: vi.fn(), - set: vi.fn(), - getRequestId: vi.fn().mockReturnValue("mock-request-id"), - getXRayTraceId: vi.fn(), - getTenantId: vi.fn().mockReturnValue("my-test-tenant-id"), - hasContext: vi.fn(), - }; - - // @ts-expect-error - mockInstance can be loosely related to original type - globalThis.awslambda.InvokeStore = mockInstance; - - // WHEN - const { InvokeStore: ReimportedStore } = await import("./invoke-store.js"); - - // THEN - expect(ReimportedStore).toBe(mockInstance); - expect(ReimportedStore.getRequestId()).toBe("mock-request-id"); - expect(ReimportedStore.getTenantId()).toBe("my-test-tenant-id"); - }); -}); - -describe("InvokeStore Environment Variable Opt-Out", () => { - const originalEnv = process.env; - - beforeEach(() => { - process.env = { ...originalEnv }; - delete (globalThis as any).awslambda; - - vi.resetModules(); - }); - - afterAll(() => { - process.env = originalEnv; - }); - - it("should respect AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA=1", async () => { - // GIVEN - process.env.AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA = "1"; - - // WHEN - Import the module with the environment variable set - const { InvokeStore } = await import("./invoke-store.js"); - - // THEN - The global namespace should not be modified - expect(globalThis.awslambda?.InvokeStore).toBeUndefined(); - - let requestId: string | undefined; - await InvokeStore.run( - { [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id" }, - () => { - requestId = InvokeStore.getRequestId(); - } - ); - expect(requestId).toBe("test-id"); - }); - - it("should respect AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA=true", async () => { - // GIVEN - process.env.AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA = "true"; - - // WHEN - Import the module with the environment variable set - const { InvokeStore } = await import("./invoke-store.js"); - - // THEN - The global namespace should not be modified - expect(globalThis.awslambda?.InvokeStore).toBeUndefined(); - - let requestId: string | undefined; - await InvokeStore.run( - { [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id" }, - () => { - requestId = InvokeStore.getRequestId(); - } - ); - expect(requestId).toBe("test-id"); + describe("InvokeStore Environment Variable Opt-Out", () => { + const originalEnv = process.env; + + beforeEach(() => { + process.env = { ...originalEnv }; + delete (globalThis as any).awslambda; + + vi.resetModules(); + }); + + afterAll(() => { + process.env = originalEnv; + }); + + it("should respect AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA=1", async () => { + // GIVEN + process.env.AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA = "1"; + + // WHEN - Import the module with the environment variable set + const { InvokeStore } = await import("./invoke-store.js"); + const invokeStore = await InvokeStore.getInstanceAsync(); + + // THEN - The global namespace should not be modified + expect(globalThis.awslambda?.InvokeStore).toBeUndefined(); + + let requestId: string | undefined; + await invokeStore.run( + { [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id" }, + () => { + requestId = invokeStore.getRequestId(); + }, + ); + expect(requestId).toBe("test-id"); + }); + + it("should respect AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA=true", async () => { + // GIVEN + process.env.AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA = "true"; + + // WHEN - Import the module with the environment variable set + const { InvokeStore } = await import("./invoke-store.js"); + const invokeStore = await InvokeStore.getInstanceAsync(); + + // THEN - The global namespace should not be modified + expect(globalThis.awslambda?.InvokeStore).toBeUndefined(); + + let requestId: string | undefined; + await invokeStore.run( + { [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id" }, + () => { + requestId = invokeStore.getRequestId(); + }, + ); + expect(requestId).toBe("test-id"); + }); }); }); diff --git a/src/invoke-store.module.loading.concurrency.spec.ts b/src/invoke-store.module.loading.concurrency.spec.ts new file mode 100644 index 0000000..a79bcf5 --- /dev/null +++ b/src/invoke-store.module.loading.concurrency.spec.ts @@ -0,0 +1,10 @@ +import { describe, vi, it, expect } from "vitest"; +import { InvokeStore } from "./invoke-store.js"; + +describe("InvokeStore implementations", () => { + it("should load the correct class", async () => { + vi.stubEnv("AWS_LAMBDA_MAX_CONCURRENCY", "2"); + const singleStore = await InvokeStore.getInstanceAsync(); + expect(singleStore.constructor.name).toBe("InvokeStoreMulti"); + }); +}); diff --git a/src/invoke-store.module.loading.spec.ts b/src/invoke-store.module.loading.spec.ts new file mode 100644 index 0000000..144a38e --- /dev/null +++ b/src/invoke-store.module.loading.spec.ts @@ -0,0 +1,9 @@ +import { describe, it, expect } from "vitest"; +import { InvokeStore } from "./invoke-store.js"; + +describe("InvokeStore implementations", () => { + it("should load the correct class", async () => { + const singleStore = await InvokeStore.getInstanceAsync(); + expect(singleStore.constructor.name).toBe("InvokeStoreSingle"); + }); +}); diff --git a/src/invoke-store.spec.ts b/src/invoke-store.spec.ts index bff52e5..eff2e6d 100644 --- a/src/invoke-store.spec.ts +++ b/src/invoke-store.spec.ts @@ -1,338 +1,249 @@ import { describe, it, expect, afterEach, beforeEach, vi } from "vitest"; -import { InvokeStore } from "./invoke-store.js"; +import { InvokeStoreBase, InvokeStore } from "./invoke-store.js"; + +describe.each([ + { label: "multi-concurrency", isMultiConcurrent: true }, + { label: "single-concurrency", isMultiConcurrent: false }, +])("InvokeStore with %s", async ({ isMultiConcurrent }) => { + describe("InvokeStore", async () => { + let invokeStore: InvokeStoreBase; + beforeEach(() => { + if (isMultiConcurrent) { + vi.stubEnv("AWS_LAMBDA_MAX_CONCURRENCY", "2"); + } + vi.useFakeTimers(); + }); -describe("InvokeStore", () => { - beforeEach(() => { - vi.useFakeTimers(); - }); + afterEach(() => { + vi.useRealTimers(); + }); - afterEach(() => { - vi.useRealTimers(); - }); + invokeStore = await InvokeStore.getInstanceAsync(); - describe("run", () => { - it("should maintain isolation between concurrent executions", async () => { - // GIVEN - const traces: string[] = []; + describe("getRequestId and getXRayTraceId", () => { + it("should return placeholder when called outside run context", () => { + // WHEN + const requestId = invokeStore.getRequestId(); + const traceId = invokeStore.getXRayTraceId(); - // WHEN - Simulate concurrent invocations - const isolateTasks = Promise.all([ - InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-1", - [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-1", - }, - async () => { - traces.push(`start-1-${InvokeStore.getRequestId()}`); - await new Promise((resolve) => setTimeout(resolve, 10)); - traces.push(`end-1-${InvokeStore.getRequestId()}`); - }, - ), - InvokeStore.run( + // THEN + expect(requestId).toBe("-"); + expect(traceId).toBeUndefined(); + }); + + it("should return current invoke IDs when called within run context", async () => { + // WHEN + const result = await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-2", - [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-2", + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-id", }, - async () => { - traces.push(`start-2-${InvokeStore.getRequestId()}`); - await new Promise((resolve) => setTimeout(resolve, 5)); - traces.push(`end-2-${InvokeStore.getRequestId()}`); + () => { + return { + requestId: invokeStore.getRequestId(), + traceId: invokeStore.getXRayTraceId(), + }; }, - ), - ]); - vi.runAllTimers(); - await isolateTasks; + ); - // THEN - expect(traces).toEqual([ - "start-1-request-1", - "start-2-request-2", - "end-2-request-2", - "end-1-request-1", - ]); - }); - - it("should maintain isolation across async operations", async () => { - // GIVEN - const traces: string[] = []; - - // WHEN - await InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-1", - }, - async () => { - traces.push(`before-${InvokeStore.getRequestId()}`); - const task = new Promise((resolve) => { - setTimeout(resolve, 1); - }).then(() => { - traces.push(`inside-${InvokeStore.getRequestId()}`); - }); - vi.runAllTimers(); - await task; - traces.push(`after-${InvokeStore.getRequestId()}`); - }, - ); - - // THEN - expect(traces).toEqual([ - "before-request-1", - "inside-request-1", - "after-request-1", - ]); + // THEN + expect(result.requestId).toBe("test-id"); + expect(result.traceId).toBe("trace-id"); + }); }); - it("should handle nested runs with different IDs", async () => { - // GIVEN - const traces: string[] = []; - - // WHEN - await InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "outer", - }, - async () => { - traces.push(`outer-${InvokeStore.getRequestId()}`); - await InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "inner", - }, - async () => { - traces.push(`inner-${InvokeStore.getRequestId()}`); - }, - ); - traces.push(`outer-again-${InvokeStore.getRequestId()}`); - }, - ); - - // THEN - expect(traces).toEqual([ - "outer-outer", - "inner-inner", - "outer-again-outer", - ]); - }); - }); + describe("custom properties", () => { + it("should allow setting and getting custom properties", async () => { + // WHEN + const result = await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + customProp: "initial-value", + }, + () => { + invokeStore.set("dynamicProp", "dynamic-value"); + return { + initial: invokeStore.get("customProp"), + dynamic: invokeStore.get("dynamicProp"), + }; + }, + ); - describe("getRequestId and getXRayTraceId", () => { - it("should return placeholder when called outside run context", () => { - // WHEN - const requestId = InvokeStore.getRequestId(); - const traceId = InvokeStore.getXRayTraceId(); + // THEN + expect(result.initial).toBe("initial-value"); + expect(result.dynamic).toBe("dynamic-value"); + }); - // THEN - expect(requestId).toBe("-"); - expect(traceId).toBeUndefined(); + it("should prevent modifying protected Lambda fields", async () => { + // WHEN & THEN + await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + }, + () => { + expect(() => { + invokeStore.set( + InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID, + "new-id", + ); + }).toThrow(/Cannot modify protected Lambda context field/); + + expect(() => { + invokeStore.set( + InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID, + "new-trace", + ); + }).toThrow(/Cannot modify protected Lambda context field/); + }, + ); + }); }); - it("should return current invoke IDs when called within run context", async () => { - // WHEN - const result = await InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-id", - }, - () => { - return { - requestId: InvokeStore.getRequestId(), - traceId: InvokeStore.getXRayTraceId(), - }; - }, - ); + describe("getContext", () => { + it("should return undefined when outside run context", () => { + // WHEN + const context = invokeStore.getContext(); - // THEN - expect(result.requestId).toBe("test-id"); - expect(result.traceId).toBe("trace-id"); - }); - }); + // THEN + expect(context).toBeUndefined(); + }); - describe("custom properties", () => { - it("should allow setting and getting custom properties", async () => { - // WHEN - const result = await InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - customProp: "initial-value", - }, - () => { - InvokeStore.set("dynamicProp", "dynamic-value"); - return { - initial: InvokeStore.get("customProp"), - dynamic: InvokeStore.get("dynamicProp"), - }; - }, - ); + it("should return complete context with Lambda and custom fields", async () => { + // WHEN + const context = await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-id", + customField: "custom-value", + }, + () => { + invokeStore.set("dynamicField", "dynamic-value"); + return invokeStore.getContext(); + }, + ); - // THEN - expect(result.initial).toBe("initial-value"); - expect(result.dynamic).toBe("dynamic-value"); + // THEN + expect(context).toEqual({ + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-id", + customField: "custom-value", + dynamicField: "dynamic-value", + }); + }); }); - it("should prevent modifying protected Lambda fields", async () => { - // WHEN & THEN - await InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - }, - () => { - expect(() => { - InvokeStore.set(InvokeStore.PROTECTED_KEYS.REQUEST_ID, "new-id"); - }).toThrow(/Cannot modify protected Lambda context field/); - - expect(() => { - InvokeStore.set( - InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID, - "new-trace", - ); - }).toThrow(/Cannot modify protected Lambda context field/); - }, - ); - }); - }); + describe("hasContext", () => { + it("should return false when outside run context", () => { + // WHEN + const hasContext = invokeStore.hasContext(); - describe("getContext", () => { - it("should return undefined when outside run context", () => { - // WHEN - const context = InvokeStore.getContext(); - - // THEN - expect(context).toBeUndefined(); - }); + // THEN + expect(hasContext).toBe(false); + }); - it("should return complete context with Lambda and custom fields", async () => { - // WHEN - const context = await InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-id", - customField: "custom-value", - }, - () => { - InvokeStore.set("dynamicField", "dynamic-value"); - return InvokeStore.getContext(); - }, - ); + it("should return true when inside run context", async () => { + // WHEN + const result = await invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + }, + () => { + return invokeStore.hasContext(); + }, + ); - // THEN - expect(context).toEqual({ - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-id", - customField: "custom-value", - dynamicField: "dynamic-value", + // THEN + expect(result).toBe(true); }); }); - }); - describe("hasContext", () => { - it("should return false when outside run context", () => { - // WHEN - const hasContext = InvokeStore.hasContext(); + describe("error handling", () => { + it("should propagate errors while maintaining isolation", async () => { + // GIVEN + const error = new Error("test error"); - // THEN - expect(hasContext).toBe(false); - }); - - it("should return true when inside run context", async () => { - // WHEN - const result = await InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - }, - () => { - return InvokeStore.hasContext(); - }, - ); + // WHEN + const promise = invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + }, + async () => { + throw error; + }, + ); - // THEN - expect(result).toBe(true); - }); - }); + // THEN + await expect(promise).rejects.toThrow(error); + expect(invokeStore.getRequestId()).toBe("-"); + }); - describe("error handling", () => { - it("should propagate errors while maintaining isolation", async () => { - // GIVEN - const error = new Error("test error"); + it("should handle errors in concurrent executions independently", async () => { + // GIVEN + const traces: string[] = []; - // WHEN - const promise = InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - }, - async () => { - throw error; - }, - ); + // WHEN + await Promise.allSettled([ + invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "success-id", + }, + async () => { + traces.push(`success-${invokeStore.getRequestId()}`); + }, + ), + invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "error-id", + }, + async () => { + traces.push(`before-error-${invokeStore.getRequestId()}`); + throw new Error("test error"); + }, + ), + ]); - // THEN - await expect(promise).rejects.toThrow(error); - expect(InvokeStore.getRequestId()).toBe("-"); + // THEN + expect(traces).toContain("success-success-id"); + expect(traces).toContain("before-error-error-id"); + expect(invokeStore.getRequestId()).toBe("-"); + }); }); - it("should handle errors in concurrent executions independently", async () => { - // GIVEN - const traces: string[] = []; - - // WHEN - await Promise.allSettled([ - InvokeStore.run( + describe("edge cases", () => { + it("should handle synchronous functions", () => { + // WHEN + console.log(InvokeStore); + const result = invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "success-id", + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", }, - async () => { - traces.push(`success-${InvokeStore.getRequestId()}`); - }, - ), - InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "error-id", - }, - async () => { - traces.push(`before-error-${InvokeStore.getRequestId()}`); - throw new Error("test error"); + () => { + return invokeStore.getRequestId(); }, - ), - ]); + ); - // THEN - expect(traces).toContain("success-success-id"); - expect(traces).toContain("before-error-error-id"); - expect(InvokeStore.getRequestId()).toBe("-"); - }); - }); - - describe("edge cases", () => { - it("should handle synchronous functions", () => { - // WHEN - const result = InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - }, - () => { - return InvokeStore.getRequestId(); - }, - ); - - // THEN - expect(result).toBe("test-id"); - }); + // THEN + expect(result).toBe("test-id"); + }); - it("should handle promises that reject synchronously", async () => { - // GIVEN - const error = new Error("immediate rejection"); + it("should handle promises that reject synchronously", async () => { + // GIVEN + const error = new Error("immediate rejection"); - // WHEN - const promise = InvokeStore.run( - { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "test-id", - }, - () => { - return Promise.reject(error); - }, - ); + // WHEN + const promise = invokeStore.run( + { + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "test-id", + }, + () => { + return Promise.reject(error); + }, + ); - // THEN - await expect(promise).rejects.toThrow(error); - expect(InvokeStore.getRequestId()).toBe("-"); + // THEN + await expect(promise).rejects.toThrow(error); + expect(invokeStore.getRequestId()).toBe("-"); + }); }); }); }); diff --git a/src/invoke-store.timers.spec.ts b/src/invoke-store.timers.concurrency.spec.ts similarity index 67% rename from src/invoke-store.timers.spec.ts rename to src/invoke-store.timers.concurrency.spec.ts index fb37ad1..b0a4bf1 100644 --- a/src/invoke-store.timers.spec.ts +++ b/src/invoke-store.timers.concurrency.spec.ts @@ -1,11 +1,18 @@ -import { describe, it, expect } from "vitest"; -import { InvokeStore } from "./invoke-store.js"; +import { describe, beforeEach, vi, it, expect } from "vitest"; +import { InvokeStore, InvokeStoreBase } from "./invoke-store.js"; /** * These tests specifically verify context preservation across various * timer and async APIs without using fake timers. */ -describe("InvokeStore timer functions context preservation", () => { +describe("InvokeStore timer functions context preservation", async () => { + let invokeStore: InvokeStoreBase; + + beforeEach(async () => { + vi.stubEnv("AWS_LAMBDA_MAX_CONCURRENCY", "2"); + invokeStore = await InvokeStore.getInstanceAsync(); + }); + const delay = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); @@ -16,21 +23,21 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`before-${InvokeStore.getRequestId()}`); + traces.push(`before-${invokeStore.getRequestId()}`); await new Promise((resolve) => { setTimeout(() => { - traces.push(`inside-timeout-${InvokeStore.getRequestId()}`); + traces.push(`inside-timeout-${invokeStore.getRequestId()}`); resolve(); }, 10); }); - traces.push(`after-${InvokeStore.getRequestId()}`); + traces.push(`after-${invokeStore.getRequestId()}`); }, ); @@ -48,25 +55,25 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`level-0-${InvokeStore.getRequestId()}`); + traces.push(`level-0-${invokeStore.getRequestId()}`); await new Promise((resolve) => { setTimeout(() => { - traces.push(`level-1-${InvokeStore.getRequestId()}`); + traces.push(`level-1-${invokeStore.getRequestId()}`); setTimeout(() => { - traces.push(`level-2-${InvokeStore.getRequestId()}`); + traces.push(`level-2-${invokeStore.getRequestId()}`); resolve(); }, 10); }, 10); }); - traces.push(`done-${InvokeStore.getRequestId()}`); + traces.push(`done-${invokeStore.getRequestId()}`); }, ); @@ -87,21 +94,21 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`before-${InvokeStore.getRequestId()}`); + traces.push(`before-${invokeStore.getRequestId()}`); await new Promise((resolve) => { setImmediate(() => { - traces.push(`inside-immediate-${InvokeStore.getRequestId()}`); + traces.push(`inside-immediate-${invokeStore.getRequestId()}`); resolve(); }); }); - traces.push(`after-${InvokeStore.getRequestId()}`); + traces.push(`after-${invokeStore.getRequestId()}`); }, ); @@ -119,25 +126,25 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`level-0-${InvokeStore.getRequestId()}`); + traces.push(`level-0-${invokeStore.getRequestId()}`); await new Promise((resolve) => { setImmediate(() => { - traces.push(`level-1-${InvokeStore.getRequestId()}`); + traces.push(`level-1-${invokeStore.getRequestId()}`); setImmediate(() => { - traces.push(`level-2-${InvokeStore.getRequestId()}`); + traces.push(`level-2-${invokeStore.getRequestId()}`); resolve(); }); }); }); - traces.push(`done-${InvokeStore.getRequestId()}`); + traces.push(`done-${invokeStore.getRequestId()}`); }, ); @@ -158,21 +165,21 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`before-${InvokeStore.getRequestId()}`); + traces.push(`before-${invokeStore.getRequestId()}`); await new Promise((resolve) => { process.nextTick(() => { - traces.push(`inside-nexttick-${InvokeStore.getRequestId()}`); + traces.push(`inside-nexttick-${invokeStore.getRequestId()}`); resolve(); }); }); - traces.push(`after-${InvokeStore.getRequestId()}`); + traces.push(`after-${invokeStore.getRequestId()}`); }, ); @@ -190,25 +197,25 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`level-0-${InvokeStore.getRequestId()}`); + traces.push(`level-0-${invokeStore.getRequestId()}`); await new Promise((resolve) => { process.nextTick(() => { - traces.push(`level-1-${InvokeStore.getRequestId()}`); + traces.push(`level-1-${invokeStore.getRequestId()}`); process.nextTick(() => { - traces.push(`level-2-${InvokeStore.getRequestId()}`); + traces.push(`level-2-${invokeStore.getRequestId()}`); resolve(); }); }); }); - traces.push(`done-${InvokeStore.getRequestId()}`); + traces.push(`done-${invokeStore.getRequestId()}`); }, ); @@ -229,18 +236,18 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`before-${InvokeStore.getRequestId()}`); + traces.push(`before-${invokeStore.getRequestId()}`); await Promise.resolve().then(() => { - traces.push(`inside-promise-${InvokeStore.getRequestId()}`); + traces.push(`inside-promise-${invokeStore.getRequestId()}`); }); - traces.push(`after-${InvokeStore.getRequestId()}`); + traces.push(`after-${invokeStore.getRequestId()}`); }, ); @@ -258,27 +265,27 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`start-${InvokeStore.getRequestId()}`); + traces.push(`start-${invokeStore.getRequestId()}`); await Promise.resolve() .then(() => { - traces.push(`then-1-${InvokeStore.getRequestId()}`); + traces.push(`then-1-${invokeStore.getRequestId()}`); return delay(10); }) .then(() => { - traces.push(`then-2-${InvokeStore.getRequestId()}`); + traces.push(`then-2-${invokeStore.getRequestId()}`); return Promise.resolve(); }) .then(() => { - traces.push(`then-3-${InvokeStore.getRequestId()}`); + traces.push(`then-3-${invokeStore.getRequestId()}`); }); - traces.push(`end-${InvokeStore.getRequestId()}`); + traces.push(`end-${invokeStore.getRequestId()}`); }, ); @@ -300,36 +307,36 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`start-${InvokeStore.getRequestId()}`); + traces.push(`start-${invokeStore.getRequestId()}`); const immediatePromise = new Promise((resolve) => { setImmediate(() => { - traces.push(`immediate-${InvokeStore.getRequestId()}`); + traces.push(`immediate-${invokeStore.getRequestId()}`); resolve(); }); }); const timeoutPromise = new Promise((resolve) => { setTimeout(() => { - traces.push(`timeout-${InvokeStore.getRequestId()}`); + traces.push(`timeout-${invokeStore.getRequestId()}`); resolve(); }, 0); }); const nextTickPromise = new Promise((resolve) => { process.nextTick(() => { - traces.push(`nextTick-${InvokeStore.getRequestId()}`); + traces.push(`nextTick-${invokeStore.getRequestId()}`); resolve(); }); }); const promisePromise = Promise.resolve().then(() => { - traces.push(`promise-${InvokeStore.getRequestId()}`); + traces.push(`promise-${invokeStore.getRequestId()}`); }); await Promise.all([ @@ -339,7 +346,7 @@ describe("InvokeStore timer functions context preservation", () => { promisePromise, ]); - traces.push(`end-${InvokeStore.getRequestId()}`); + traces.push(`end-${invokeStore.getRequestId()}`); }, ); @@ -361,26 +368,26 @@ describe("InvokeStore timer functions context preservation", () => { // WHEN - Simulate concurrent invocations await Promise.all([ - InvokeStore.run( + invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-1", - [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-1", + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-1", + [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-1", }, async () => { - traces.push(`start-1-${InvokeStore.getRequestId()}`); + traces.push(`start-1-${invokeStore.getRequestId()}`); await delay(20); // Longer delay - traces.push(`end-1-${InvokeStore.getRequestId()}`); + traces.push(`end-1-${invokeStore.getRequestId()}`); }, ), - InvokeStore.run( + invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-2", - [InvokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-2", + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-2", + [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-2", }, async () => { - traces.push(`start-2-${InvokeStore.getRequestId()}`); + traces.push(`start-2-${invokeStore.getRequestId()}`); await delay(10); // Shorter delay - traces.push(`end-2-${InvokeStore.getRequestId()}`); + traces.push(`end-2-${invokeStore.getRequestId()}`); }, ), ]); @@ -400,58 +407,58 @@ describe("InvokeStore timer functions context preservation", () => { // WHEN - Simulate concurrent invocations with different async operations await Promise.all([ - InvokeStore.run( + invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-1", + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-1", }, async () => { - traces.push(`start-1-${InvokeStore.getRequestId()}`); + traces.push(`start-1-${invokeStore.getRequestId()}`); // Use setTimeout await new Promise((resolve) => { setTimeout(() => { - traces.push(`timeout-1-${InvokeStore.getRequestId()}`); + traces.push(`timeout-1-${invokeStore.getRequestId()}`); resolve(); }, 15); }); - traces.push(`end-1-${InvokeStore.getRequestId()}`); + traces.push(`end-1-${invokeStore.getRequestId()}`); }, ), - InvokeStore.run( + invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-2", + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-2", }, async () => { - traces.push(`start-2-${InvokeStore.getRequestId()}`); + traces.push(`start-2-${invokeStore.getRequestId()}`); // Use setImmediate await new Promise((resolve) => { setImmediate(() => { - traces.push(`immediate-2-${InvokeStore.getRequestId()}`); + traces.push(`immediate-2-${invokeStore.getRequestId()}`); resolve(); }); }); - traces.push(`end-2-${InvokeStore.getRequestId()}`); + traces.push(`end-2-${invokeStore.getRequestId()}`); }, ), - InvokeStore.run( + invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-3", + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: "request-3", }, async () => { - traces.push(`start-3-${InvokeStore.getRequestId()}`); + traces.push(`start-3-${invokeStore.getRequestId()}`); // Use process.nextTick await new Promise((resolve) => { process.nextTick(() => { - traces.push(`nextTick-3-${InvokeStore.getRequestId()}`); + traces.push(`nextTick-3-${invokeStore.getRequestId()}`); resolve(); }); }); - traces.push(`end-3-${InvokeStore.getRequestId()}`); + traces.push(`end-3-${invokeStore.getRequestId()}`); }, ), ]); @@ -504,17 +511,17 @@ describe("InvokeStore timer functions context preservation", () => { const iterations = 3; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`start-${InvokeStore.getRequestId()}`); + traces.push(`start-${invokeStore.getRequestId()}`); let count = 0; await new Promise((resolve) => { function recursive() { - traces.push(`iteration-${count}-${InvokeStore.getRequestId()}`); + traces.push(`iteration-${count}-${invokeStore.getRequestId()}`); count++; if (count < iterations) { @@ -527,7 +534,7 @@ describe("InvokeStore timer functions context preservation", () => { recursive(); }); - traces.push(`end-${InvokeStore.getRequestId()}`); + traces.push(`end-${invokeStore.getRequestId()}`); }, ); @@ -547,30 +554,30 @@ describe("InvokeStore timer functions context preservation", () => { const traces: string[] = []; // WHEN - await InvokeStore.run( + await invokeStore.run( { - [InvokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId, + [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: testRequestId, }, async () => { - traces.push(`start-${InvokeStore.getRequestId()}`); + traces.push(`start-${invokeStore.getRequestId()}`); // Queue a setTimeout that triggers setImmediate await new Promise((resolve) => { setTimeout(() => { - traces.push(`timeout-${InvokeStore.getRequestId()}`); + traces.push(`timeout-${invokeStore.getRequestId()}`); setImmediate(() => { - traces.push(`immediate-${InvokeStore.getRequestId()}`); + traces.push(`immediate-${invokeStore.getRequestId()}`); process.nextTick(() => { - traces.push(`nextTick-${InvokeStore.getRequestId()}`); + traces.push(`nextTick-${invokeStore.getRequestId()}`); resolve(); }); }); }, 10); }); - traces.push(`end-${InvokeStore.getRequestId()}`); + traces.push(`end-${invokeStore.getRequestId()}`); }, ); diff --git a/src/invoke-store.ts b/src/invoke-store.ts index 9cda776..cc3d3a5 100644 --- a/src/invoke-store.ts +++ b/src/invoke-store.ts @@ -1,126 +1,195 @@ -import { AsyncLocalStorage } from "async_hooks"; - -// AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA provides an escape hatch since we're modifying the global object which may not be expected to a customer's handler. -const noGlobalAwsLambda = - process.env["AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA"] === "1" || - process.env["AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA"] === "true"; - -if (!noGlobalAwsLambda) { - globalThis.awslambda = globalThis.awslambda || {}; +import type { AsyncLocalStorage } from "node:async_hooks"; +interface Context { + [key: string]: unknown; + [key: symbol]: unknown; } const PROTECTED_KEYS = { - REQUEST_ID: Symbol("_AWS_LAMBDA_REQUEST_ID"), - X_RAY_TRACE_ID: Symbol("_AWS_LAMBDA_X_RAY_TRACE_ID"), - TENANT_ID: Symbol("_AWS_LAMBDA_TENANT_ID"), + 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"), } as const; +const NO_GLOBAL_AWS_LAMBDA = ["true", "1"].includes( + process.env?.AWS_LAMBDA_NODEJS_NO_GLOBAL_AWSLAMBDA ?? "", +); + +declare global { + var awslambda: { + InvokeStore?: InvokeStoreBase; + [key: string]: unknown; + }; +} + +if (!NO_GLOBAL_AWS_LAMBDA) { + globalThis.awslambda = globalThis.awslambda || {}; +} + /** - * Generic store context that uses protected keys for Lambda fields - * and allows custom user properties + * Base class for AWS Lambda context storage implementations. + * Provides core functionality for managing Lambda execution context. + * + * Implementations handle either single-context (InvokeStoreSingle) or + * multi-context (InvokeStoreMulti) scenarios based on Lambda's execution environment. + * + * @public */ -export interface InvokeStoreContext { - [key: string | symbol]: unknown; +export abstract class InvokeStoreBase { + public static readonly PROTECTED_KEYS = PROTECTED_KEYS; + + abstract getContext(): Context | undefined; + abstract hasContext(): boolean; + abstract get(key: string | symbol): T | undefined; + abstract set(key: string | symbol, value: T): void; + abstract run(context: Context, fn: () => T): T; + + protected isProtectedKey(key: string | symbol): boolean { + return Object.values(PROTECTED_KEYS).includes(key as symbol); + } + + getRequestId(): string { + return this.get(PROTECTED_KEYS.REQUEST_ID) ?? "-"; + } + + getXRayTraceId(): string | undefined { + return this.get(PROTECTED_KEYS.X_RAY_TRACE_ID); + } + + getTenantId(): string | undefined { + return this.get(PROTECTED_KEYS.TENANT_ID); + } } /** - * InvokeStore implementation class + * Single Context Implementation + * @internal */ -class InvokeStoreImpl { - private static storage = new AsyncLocalStorage(); +class InvokeStoreSingle extends InvokeStoreBase { + private currentContext?: Context; - // Protected keys for Lambda context fields - public static readonly PROTECTED_KEYS = PROTECTED_KEYS; - - /** - * Initialize and run code within an invoke context - */ - public static run( - context: InvokeStoreContext, - fn: () => T | Promise, - ): T | Promise { - return this.storage.run({ ...context }, fn); + getContext(): Context | undefined { + return this.currentContext; } - /** - * Get the complete current context - */ - public static getContext(): InvokeStoreContext | undefined { - return this.storage.getStore(); + hasContext(): boolean { + return this.currentContext !== undefined; } - /** - * Get a specific value from the context by key - */ - public static get(key: string | symbol): T | undefined { - const context = this.storage.getStore(); - return context?.[key] as T | undefined; + get(key: string | symbol): T | undefined { + return this.currentContext?.[key] as T | undefined; } - /** - * Set a custom value in the current context - * Protected Lambda context fields cannot be overwritten - */ - public static set(key: string | symbol, value: unknown): void { + set(key: string | symbol, value: T): void { if (this.isProtectedKey(key)) { - throw new Error(`Cannot modify protected Lambda context field`); + throw new Error( + `Cannot modify protected Lambda context field: ${String(key)}`, + ); } - const context = this.storage.getStore(); - if (context) { - context[key] = value; - } + this.currentContext = this.currentContext || {}; + this.currentContext[key] = value; } - /** - * Get the current request ID - */ - public static getRequestId(): string { - return this.get(this.PROTECTED_KEYS.REQUEST_ID) ?? "-"; + run(context: Context, fn: () => T): T { + this.currentContext = context; + try { + return fn(); + } finally { + this.currentContext = undefined; + } } +} - /** - * Get the current X-ray trace ID - */ - public static getXRayTraceId(): string | undefined { - return this.get(this.PROTECTED_KEYS.X_RAY_TRACE_ID); +/** + * Multi Context Implementation + * @internal + */ +class InvokeStoreMulti extends InvokeStoreBase { + private als!: AsyncLocalStorage; + + static async create(): Promise { + const instance = new InvokeStoreMulti(); + const asyncHooks = await import("node:async_hooks"); + instance.als = new asyncHooks.AsyncLocalStorage(); + return instance; } - /** - * Get the current tenant ID - */ - public static getTenantId(): string | undefined { - return this.get(this.PROTECTED_KEYS.TENANT_ID); + getContext(): Context | undefined { + return this.als.getStore(); } - /** - * Check if we're currently within an invoke context - */ - public static hasContext(): boolean { - return this.storage.getStore() !== undefined; + hasContext(): boolean { + return this.als.getStore() !== undefined; } - /** - * Check if a key is protected (readonly Lambda context field) - */ - private static isProtectedKey(key: string | symbol): boolean { - return ( - key === this.PROTECTED_KEYS.REQUEST_ID || - key === this.PROTECTED_KEYS.X_RAY_TRACE_ID - ); + get(key: string | symbol): T | undefined { + return this.als.getStore()?.[key] as T | undefined; } -} -let instance: typeof InvokeStoreImpl; + set(key: string | symbol, value: T): void { + if (this.isProtectedKey(key)) { + throw new Error( + `Cannot modify protected Lambda context field: ${String(key)}`, + ); + } -if (!noGlobalAwsLambda && globalThis.awslambda?.InvokeStore) { - instance = globalThis.awslambda.InvokeStore; -} else { - instance = InvokeStoreImpl; + const store = this.als.getStore(); + if (!store) { + throw new Error("No context available"); + } - if (!noGlobalAwsLambda && globalThis.awslambda) { - globalThis.awslambda.InvokeStore = instance; + store[key] = value; + } + + run(context: Context, fn: () => T): T { + return this.als.run(context, fn); } } -export const InvokeStore = instance; +/** + * Provides access to AWS Lambda execution context storage. + * Supports both single-context and multi-context environments through different implementations. + * + * The store manages protected Lambda context fields and allows storing/retrieving custom values + * within the execution context. + * @public + */ +export namespace InvokeStore { + let instance: Promise | null = null; + + export async function getInstanceAsync(): Promise { + if (!instance) { + // Lock synchronously on first invoke by immediately assigning the promise + instance = (async () => { + const isMulti = "AWS_LAMBDA_MAX_CONCURRENCY" in process.env; + const newInstance = isMulti + ? await InvokeStoreMulti.create() + : new InvokeStoreSingle(); + + if (!NO_GLOBAL_AWS_LAMBDA && globalThis.awslambda?.InvokeStore) { + return globalThis.awslambda.InvokeStore; + } else if (!NO_GLOBAL_AWS_LAMBDA && globalThis.awslambda) { + globalThis.awslambda.InvokeStore = newInstance; + return newInstance; + } else { + return newInstance; + } + })(); + } + + return instance; + } + + export const _testing = + process.env.AWS_LAMBDA_BENCHMARK_MODE === "1" + ? { + reset: () => { + instance = null; + if (globalThis.awslambda?.InvokeStore) { + delete globalThis.awslambda.InvokeStore; + } + globalThis.awslambda = {}; + }, + } + : undefined; +}