diff --git a/Dockerfile.js b/Dockerfile.js index 3d920cd..d9e04e7 100644 --- a/Dockerfile.js +++ b/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 diff --git a/src/client/rapid-client.ts b/src/client/rapid-client.ts index f8b402b..a27bd48 100644 --- a/src/client/rapid-client.ts +++ b/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; diff --git a/src/context/constants.ts b/src/context/constants.ts index 302f22b..8ebc055 100644 --- a/src/context/constants.ts +++ b/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; diff --git a/src/runtime/runtime.test.ts b/src/runtime/runtime.test.ts index 781cb22..1df2570 100644 --- a/src/runtime/runtime.test.ts +++ b/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({ diff --git a/src/runtime/runtime.ts b/src/runtime/runtime.ts index 379ee8e..6405070 100644 --- a/src/runtime/runtime.ts +++ b/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 { + const pollingCount = getPollingCount(); + log.vvverbose(`Multi-concurrent polling count: ${pollingCount}`); + + const pendingNextCalls = new Set>(); + + // 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, diff --git a/src/utils/env.ts b/src/utils/env.ts index 695940c..b23107e 100644 --- a/src/utils/env.ts +++ b/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]) { diff --git a/src/utils/worker.test.ts b/src/utils/worker.test.ts index 2aff31b..36365df 100644 --- a/src/utils/worker.test.ts +++ b/src/utils/worker.test.ts @@ -21,9 +21,9 @@ const CUSTOMER_ENV_RESPONSE = 5; const CPU_MAX_QUOTA = 200000; const CPU_MAX_PERIOD = 100000; const CPU_MAX_CALCULATED_CPUS = 2; -const CPU_MAX_EXPECTED_WORKERS = 8 * CPU_MAX_CALCULATED_CPUS; +const CPU_MAX_EXPECTED_WORKERS = 2 * CPU_MAX_CALCULATED_CPUS + 1; const DEFAULT_CPU_COUNT = 1; -const DEFAULT_WORKER_COUNT = 8 * DEFAULT_CPU_COUNT; +const DEFAULT_WORKER_COUNT = 2 * DEFAULT_CPU_COUNT + 1; const OS_DETECTED_CPUS = 4; describe("worker utils", () => { @@ -106,7 +106,7 @@ describe("worker utils", () => { const result = getWorkerCount(); // THEN - expect(result).toBe(8 * OS_DETECTED_CPUS); + expect(result).toBe(2 * OS_DETECTED_CPUS + 1); }); it("should handle invalid cpu.max values", async () => { @@ -128,8 +128,8 @@ describe("worker utils", () => { // WHEN const result = getWorkerCount(); - // THEN - expect(result).toBeNaN(); + // THEN - invalid value is ignored, falls back to CPU detection + expect(result).toBe(DEFAULT_WORKER_COUNT); }); it("should handle zero value", () => { @@ -139,8 +139,8 @@ describe("worker utils", () => { // WHEN const result = getWorkerCount(); - // THEN - expect(result).toBe(0); + // THEN - zero is invalid, falls back to CPU detection + expect(result).toBe(DEFAULT_WORKER_COUNT); }); }); }); diff --git a/src/utils/worker.ts b/src/utils/worker.ts index 4ef6be9..94b1b29 100644 --- a/src/utils/worker.ts +++ b/src/utils/worker.ts @@ -1,5 +1,6 @@ import { logger } from "../logging/verbose-log.js"; import { cjsRequire } from "./cjs-require.js"; +import { getWorkerCountFromEnv } from "./env.js"; const { readFileSync } = cjsRequire("node:fs"); const os = cjsRequire("node:os"); @@ -48,19 +49,18 @@ function getAvailableCpus(): number { } export function getWorkerCount(): number { - const envValue = process.env.AWS_LAMBDA_NODEJS_WORKER_COUNT; - if (envValue) { - const workerCount = parseInt(envValue, 10); + const envWorkerCount = getWorkerCountFromEnv(); + if (envWorkerCount !== undefined) { verboseLog.vvverbose( - `Using AWS_LAMBDA_NODEJS_WORKER_COUNT: ${envValue} = ${workerCount} workers`, + `Using AWS_LAMBDA_NODEJS_WORKER_COUNT: ${envWorkerCount} workers`, ); - return workerCount; + return envWorkerCount; } const detectedCpus = getAvailableCpus(); - const calculatedWorkers = 8 * detectedCpus; + const calculatedWorkers = 2 * detectedCpus + 1; verboseLog.vvverbose( - `No env var set, using 8 * ${detectedCpus} = ${calculatedWorkers} workers`, + `No env var set, using 2 * ${detectedCpus} + 1 = ${calculatedWorkers} workers`, ); return calculatedWorkers; }