feat: lazy loading (#13)

Co-authored-by: Trivikram Kamat <[email protected]>

Author: maxday

.changeset/changelog.mjs

@@ -14,17 +14,19 @@ 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");
 };
 
 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") : ""
   }`;
-};
+};

.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)

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.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,
   };
 };
 
 // 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
   },
   () => {

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.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);

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",
+      ]);
+    });
+  });
+});