feat: multi-pending /next polling for multi-concurrent mode

Dockerfile.js

@@ -22,8 +22,8 @@ RUN mkdir -p /build && \
     tar xzf /tmp/deps.tar.gz -C /build && \
     ls -R /build/deps
 
-# Copy bare config
-COPY package.json package-lock.json tsconfig.json eslint.config.js vitest.config.js vitest.setup.ts  /app/
+# Copy dependency manifests and bare config
+COPY package.json package-lock.json tsconfig.json eslint.config.js vitest.config.js vitest.setup.ts /app/
 
 WORKDIR /app
 

src/client/rapid-client.ts

@@ -17,6 +17,7 @@ import {
   RetryOptions,
   DEFAULT_RETRY_OPTIONS,
   calculateBackoffDelay,
+  getPollingCount,
 } from "../utils/index.js";
 
 import { StreamOptions } from "../context/types.js";
@@ -71,7 +72,7 @@ export class RAPIDClient {
     this.nativeClient = nativeClient;
     this.agent = new this.httpClient.Agent({
       keepAlive: true,
-      maxSockets: 1,
+      maxSockets: isMultiConcurrent ? getPollingCount() : 1,
     });
     this.useAlternativeClient = shouldUseAlternativeClient();
     this.isMultiConcurrent = isMultiConcurrent;

src/context/constants.ts

@@ -39,3 +39,7 @@ export const CALLBACK_ERROR_NODEJS22 =
 
 export const CALLBACK_ERROR_NODEJS24_ABOVE =
   "ERROR: AWS Lambda has removed support for callback-based function handlers starting with Node.js 24. You need to modify this function to use a supported handler signature to use Node.js 24 or later. For more information see https://docs.aws.amazon.com/lambda/latest/dg/nodejs-handler.html.";
+
+// Number of concurrent /next calls to keep pending in multi-concurrent mode.
+// Override via AWS_LAMBDA_NODEJS_POLLING_COUNT env var.
+export const DEFAULT_POLLING_COUNT = 4;

src/runtime/runtime.test.ts

@@ -418,12 +418,15 @@ describe("Runtime", () => {
     beforeEach(async () => {
       // Set AWS_LAMBDA_MAX_CONCURRENCY to ensure InvokeStoreMulti is used
       process.env.AWS_LAMBDA_MAX_CONCURRENCY = "1";
+      // Use polling count of 1 for existing tests (they test handler concurrency, not multi-pending)
+      process.env.AWS_LAMBDA_NODEJS_POLLING_COUNT = "1";
       InvokeStore._testing?.reset?.();
       invokeStore = await InvokeStore.getInstanceAsync();
     });
 
     afterEach(() => {
       delete process.env.AWS_LAMBDA_MAX_CONCURRENCY;
+      delete process.env.AWS_LAMBDA_NODEJS_POLLING_COUNT;
       InvokeStore._testing?.reset();
     });
     it("should process invocations concurrently", async () => {
@@ -875,6 +878,202 @@ describe("Runtime", () => {
       });
     });
   });
+
+  describe("[multi-pending]", () => {
+    beforeEach(async () => {
+      process.env.AWS_LAMBDA_MAX_CONCURRENCY = "32";
+      InvokeStore._testing?.reset?.();
+      await InvokeStore.getInstanceAsync();
+    });
+
+    afterEach(() => {
+      delete process.env.AWS_LAMBDA_MAX_CONCURRENCY;
+      delete process.env.AWS_LAMBDA_NODEJS_POLLING_COUNT;
+      InvokeStore._testing?.reset();
+    });
+
+    it("should fire pollingCount /next calls at startup", async () => {
+      // GIVEN
+      process.env.AWS_LAMBDA_NODEJS_POLLING_COUNT = "4";
+      // All calls hang (simulating RAPID holding /next connections)
+      const mockRapidClient = {
+        nextInvocation: vi.fn().mockImplementation(() => {
+          return new Promise(() => {});
+        }),
+        postInvocationResponse: vi.fn(async () => {}),
+        postInvocationError: vi.fn(async () => {}),
+        postInitError: vi.fn(async () => {}),
+      } as unknown as RAPIDClient;
+
+      const handler: BufferedHandler = async () => "done";
+
+      const runtime = Runtime.create({
+        rapidClient: mockRapidClient,
+        handler,
+        isMultiConcurrent: true,
+      });
+
+      // WHEN
+      runtime.start();
+      await vi.advanceTimersByTimeAsync(10);
+
+      // THEN - exactly 4 calls issued (all pending, no resolutions)
+      expect(mockRapidClient.nextInvocation).toHaveBeenCalledTimes(4);
+    });
+
+    it("should respect AWS_LAMBDA_NODEJS_POLLING_COUNT env var", async () => {
+      // GIVEN
+      process.env.AWS_LAMBDA_NODEJS_POLLING_COUNT = "2";
+      const mockRapidClient = {
+        nextInvocation: vi.fn().mockImplementation(() => {
+          return new Promise(() => {}); // all hang
+        }),
+        postInvocationResponse: vi.fn(async () => {}),
+        postInvocationError: vi.fn(async () => {}),
+        postInitError: vi.fn(async () => {}),
+      } as unknown as RAPIDClient;
+
+      const handler: BufferedHandler = async () => "done";
+
+      const runtime = Runtime.create({
+        rapidClient: mockRapidClient,
+        handler,
+        isMultiConcurrent: true,
+      });
+
+      // WHEN
+      runtime.start();
+      await vi.advanceTimersByTimeAsync(10);
+
+      // THEN - only 2 calls issued
+      expect(mockRapidClient.nextInvocation).toHaveBeenCalledTimes(2);
+    });
+
+    it("should replenish /next after invocation resolves", async () => {
+      // GIVEN
+      vi.useRealTimers();
+      process.env.AWS_LAMBDA_NODEJS_POLLING_COUNT = "2";
+      const resolvers: Array<(value: unknown) => void> = [];
+
+      const mockRapidClient = {
+        nextInvocation: vi.fn().mockImplementation(() => {
+          return new Promise((resolve) => {
+            resolvers.push(resolve);
+          });
+        }),
+        postInvocationResponse: vi.fn(async () => {}),
+        postInvocationError: vi.fn(async () => {}),
+        postInitError: vi.fn(async () => {}),
+      } as unknown as RAPIDClient;
+
+      const handler: BufferedHandler = async () => "done";
+
+      const runtime = Runtime.create({
+        rapidClient: mockRapidClient,
+        handler,
+        isMultiConcurrent: true,
+      });
+
+      // WHEN - start: 2 pending
+      runtime.start();
+      await new Promise((r) => setTimeout(r, 10));
+      expect(mockRapidClient.nextInvocation).toHaveBeenCalledTimes(2);
+      expect(resolvers.length).toBe(2);
+
+      // Resolve one - should trigger replenishment
+      resolvers[0]({
+        bodyJson: JSON.stringify({ key: "value" }),
+        headers: {
+          "lambda-runtime-aws-request-id": "test-id-1",
+          "lambda-runtime-deadline-ms": "1234567890",
+          "lambda-runtime-invoked-function-arn": "test:arn",
+        },
+      });
+      await new Promise((r) => setTimeout(r, 10));
+
+      // THEN - 2 initial + 1 replenishment = 3
+      expect(mockRapidClient.nextInvocation).toHaveBeenCalledTimes(3);
+    });
+
+    it("should not exceed pollingCount pending calls", async () => {
+      // GIVEN
+      process.env.AWS_LAMBDA_NODEJS_POLLING_COUNT = "3";
+      let concurrentPending = 0;
+      let maxConcurrentPending = 0;
+      let invocationCount = 0;
+
+      const mockRapidClient = {
+        nextInvocation: vi.fn().mockImplementation(() => {
+          invocationCount++;
+          concurrentPending++;
+          maxConcurrentPending = Math.max(
+            maxConcurrentPending,
+            concurrentPending,
+          );
+
+          if (invocationCount >= 10) {
+            return new Promise(() => {}); // throttle
+          }
+
+          return new Promise((resolve) => {
+            setTimeout(() => {
+              concurrentPending--;
+              resolve({
+                ...baseMockInvocation,
+                headers: {
+                  ...baseMockInvocation.headers,
+                  "lambda-runtime-aws-request-id": `test-id-${invocationCount}`,
+                },
+              });
+            }, 10);
+          });
+        }),
+        postInvocationResponse: vi.fn(async () => {}),
+        postInvocationError: vi.fn(async () => {}),
+        postInitError: vi.fn(async () => {}),
+      } as unknown as RAPIDClient;
+
+      const handler: BufferedHandler = async () => "done";
+
+      const runtime = Runtime.create({
+        rapidClient: mockRapidClient,
+        handler,
+        isMultiConcurrent: true,
+      });
+
+      // WHEN
+      await startAndTickRuntime(runtime, 100);
+
+      // THEN - pending never exceeded 3
+      expect(maxConcurrentPending).toBeLessThanOrEqual(3);
+    });
+
+    it("should dispatch all handlers without capping in-flight", async () => {
+      // GIVEN
+      process.env.AWS_LAMBDA_NODEJS_POLLING_COUNT = "2";
+      let handlersDispatched = 0;
+
+      const mockRapidClient = createMockRapidClient(8); // allow 7 invocations
+
+      const handler: BufferedHandler = async () => {
+        handlersDispatched++;
+        await new Promise((resolve) => setTimeout(resolve, 100));
+        return "done";
+      };
+
+      const runtime = Runtime.create({
+        rapidClient: mockRapidClient,
+        handler,
+        isMultiConcurrent: true,
+      });
+
+      // WHEN
+      await startAndTickRuntime(runtime, 50);
+
+      // THEN - all resolved invocations dispatched (no cap)
+      expect(handlersDispatched).toBeGreaterThanOrEqual(2);
+    });
+  });
 });
 
 const baseMockInvocation: InvocationRequest = {
@@ -902,7 +1101,7 @@ function createMockRapidClient(
   return {
     nextInvocation: vi.fn().mockImplementation(() => {
       invocationCount++;
-      if (invocationCount == throttleAt) {
+      if (throttleAt > 0 && invocationCount >= throttleAt) {
         return new Promise(() => {});
       }
       return Promise.resolve({

src/runtime/runtime.ts

@@ -6,10 +6,15 @@ import {
   StreamingHandler,
   UserHandler,
 } from "./types.js";
+import { EventAndContext } from "../context/types.js";
 import { InvokeStore, InvokeStoreBase } from "@aws/lambda-invoke-store";
 import { LifecycleManager } from "./lifecycle-manager.js";
 import { BufferedInvokeProcessor } from "./buffered-invoke-processor.js";
 import { StreamingInvokeProcessor } from "./streaming-invoke-processor.js";
+import { logger } from "../logging/verbose-log.js";
+import { getPollingCount } from "../utils/env.js";
+
+const log = logger("RUNTIME");
 
 export class Runtime {
   public static create({
@@ -75,9 +80,32 @@ export class Runtime {
   private async processMultiConcurrent(
     processor: InvokeProcessor,
   ): Promise<void> {
+    const pollingCount = getPollingCount();
+    log.vvverbose(`Multi-concurrent polling count: ${pollingCount}`);
+
+    const pendingNextCalls = new Set<Promise<EventAndContext>>();
+
+    // Fill initial pending /next calls
+    while (pendingNextCalls.size < pollingCount) {
+      pendingNextCalls.add(this.lifecycle.next());
+    }
+
     while (true) {
-      const { context, event } = await this.lifecycle.next();
+      // Wait for any pending /next to return with an invocation
+      const resolved = await Promise.race(
+        [...pendingNextCalls].map((p) =>
+          p.then((result) => ({ promise: p, result })),
+        ),
+      );
+
+      // Remove the resolved promise from tracking
+      pendingNextCalls.delete(resolved.promise);
+
+      // Replenish: fire replacement /next BEFORE dispatching handler
+      pendingNextCalls.add(this.lifecycle.next());
 
+      // Dispatch handler via setImmediate (fire-and-forget)
+      const { context, event } = resolved.result;
       setImmediate(async () => {
         await this.runWithInvokeContext(
           context.awsRequestId,

src/utils/env.ts

@@ -1,5 +1,6 @@
 import { HEADERS } from "../context/index.js";
 import { InvokeHeaders } from "../context/types.js";
+import { DEFAULT_POLLING_COUNT } from "../context/constants.js";
 import {
   LogFormat,
   LOG_FORMAT,
@@ -35,6 +36,60 @@ export function isMultiConcurrentMode(): boolean {
   return process.env["AWS_LAMBDA_MAX_CONCURRENCY"] !== undefined;
 }
 
+// Default polling count is 4. Override via AWS_LAMBDA_NODEJS_POLLING_COUNT.
+// Capped at AWS_LAMBDA_MAX_CONCURRENCY to prevent excessive connections to RAPID.
+export function getPollingCount(): number {
+  const maxConcurrency = parseInt(
+    process.env["AWS_LAMBDA_MAX_CONCURRENCY"] || "0",
+    10,
+  );
+  const envValue = process.env["AWS_LAMBDA_NODEJS_POLLING_COUNT"];
+  if (envValue) {
+    const parsed = parseInt(envValue, 10);
+    if (!parsed || parsed <= 0) {
+      console.warn(
+        `Invalid AWS_LAMBDA_NODEJS_POLLING_COUNT="${envValue}", using default ${DEFAULT_POLLING_COUNT}`,
+      );
+      return DEFAULT_POLLING_COUNT;
+    }
+    if (maxConcurrency > 0 && parsed > maxConcurrency) {
+      console.warn(
+        `AWS_LAMBDA_NODEJS_POLLING_COUNT=${parsed} exceeds MC=${maxConcurrency}, capping`,
+      );
+      return maxConcurrency;
+    }
+    return parsed;
+  }
+  return DEFAULT_POLLING_COUNT;
+}
+
+// Override via AWS_LAMBDA_NODEJS_WORKER_COUNT.
+// Capped at AWS_LAMBDA_MAX_CONCURRENCY to prevent spawning more workers than MC.
+export function getWorkerCountFromEnv(): number | undefined {
+  const envValue = process.env["AWS_LAMBDA_NODEJS_WORKER_COUNT"];
+  if (!envValue) return undefined;
+
+  const parsed = parseInt(envValue, 10);
+  if (!parsed || parsed <= 0) {
+    console.warn(
+      `Invalid AWS_LAMBDA_NODEJS_WORKER_COUNT="${envValue}", ignoring`,
+    );
+    return undefined;
+  }
+
+  const maxConcurrency = parseInt(
+    process.env["AWS_LAMBDA_MAX_CONCURRENCY"] || "0",
+    10,
+  );
+  if (maxConcurrency > 0 && parsed > maxConcurrency) {
+    console.warn(
+      `AWS_LAMBDA_NODEJS_WORKER_COUNT=${parsed} exceeds MC=${maxConcurrency}, capping`,
+    );
+    return maxConcurrency;
+  }
+  return parsed;
+}
+
 export function moveXRayHeaderToEnv(headers: InvokeHeaders) {
   if (!isMultiConcurrentMode()) {
     if (headers[HEADERS.X_RAY_TRACE_ID]) {