feat: lazy loading

.changeset/goofy-planes-tease.md

@@ -0,0 +1,10 @@
+---
+"@aws/lambda-invoke-store": minor
+---
+
+- `InvokeStore` is now accessible via `InvokeStore.getInstance()` instead of direct instantiation
+
+- Lazy loading of `node:async_hooks` to improve startup performance
+- Dynamic implementation selection based on Lambda environment:
+  - Single-context implementation for standard Lambda executions
+  - Multi-context implementation (using AsyncLocalStorage)

README.md

@@ -33,109 +33,117 @@ 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.getInstance();
+  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,
   };
 };
 
 // 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.getInstance();
+  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.getInstance()
+First, get an instance of the InvokeStore:
+```typescript
+const invokeStore = await InvokeStore.getInstance();
+```
+
+### 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(InvokeStore.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(invokeStore.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
+    [invokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-123",
+    [invokeStore.PROTECTED_KEYS.X_RAY_TRACE_ID]: "trace-456", // Optional X-Ray trace ID
     customField: "value", // Optional custom fields
   },
   () => {

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"
   },

src/invoke-store.concurrency.spec.ts

@@ -0,0 +1,122 @@
+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.getInstance();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  describe("run", () => {
+
+    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",
+      ]);
+    });
+
+    it("should maintain isolation between concurrent executions", async () => {
+      // GIVEN
+      const traces: string[] = [];
+
+      // 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(
+          {
+            [invokeStore.PROTECTED_KEYS.REQUEST_ID]: "request-2",
+            [invokeStore.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(
+        {
+          [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",
+      ]);
+    });
+  });
+});

src/invoke-store.global.concurrency.spec.ts

@@ -0,0 +1,68 @@
+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.getInstance();
+  });
+
+  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(
+      {
+        [invokeStore.PROTECTED_KEYS.REQUEST_ID]: testRequestId,
+        [invokeStore.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.getInstance();
+
+        // 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);
+      }
+    );
+  });
+});