@remit/sqs-client 0.0.9 → 0.0.11

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@remit/sqs-client",
-  "version": "0.0.9",
+  "version": "0.0.11",
   "type": "module",
   "main": "src/index.ts",
   "types": "src/index.ts",

package/src/poller.test.ts

@@ -1,7 +1,7 @@
 import assert from "node:assert/strict";
 import { describe, it, mock } from "node:test";
 import { SQSClient } from "@aws-sdk/client-sqs";
-import type { SQSBatchResponse, SQSEvent } from "aws-lambda";
+import type { Context, SQSBatchResponse, SQSEvent } from "aws-lambda";
 import { runQueuePoller } from "./poller.js";
 
 const QUEUE_URL = "http://localhost:9324/000000000000/remit-test";
@@ -81,6 +81,60 @@ describe("runQueuePoller", () => {
 		assert.deepEqual(deletedReceiptHandles, ["receipt-m1"]);
 	});
 
+	it("aborts an in-flight idle long-poll on signal instead of waiting it out", async () => {
+		let receiveCalls = 0;
+		let sawAbortSignal = false;
+
+		const sendMock = mock.method(
+			SQSClient.prototype,
+			"send",
+			// biome-ignore lint/suspicious/noExplicitAny: minimal SDK command shape
+			async function (this: SQSClient, command: any, options?: any) {
+				if (command.constructor.name === "ReceiveMessageCommand") {
+					receiveCalls++;
+					if (receiveCalls > 1) {
+						throw new Error(
+							"a second ReceiveMessage was issued — the abort did not end the poll",
+						);
+					}
+					const signal: AbortSignal | undefined = options?.abortSignal;
+					sawAbortSignal = signal instanceof AbortSignal;
+					// Model an empty long poll that only ever ends because it is
+					// aborted: request shutdown, then resolve solely on abort. A
+					// poller that waited out WaitTimeSeconds would hang this test.
+					process.emit("SIGWINCH", "SIGWINCH");
+					return await new Promise((_resolve, reject) => {
+						const fail = () =>
+							reject(
+								Object.assign(new Error("aborted"), { name: "AbortError" }),
+							);
+						if (signal?.aborted) return fail();
+						signal?.addEventListener("abort", fail, { once: true });
+					});
+				}
+				throw new Error(`unexpected command: ${command.constructor.name}`);
+			},
+		);
+
+		try {
+			await runQueuePoller({
+				targets: [
+					{ queueUrl: QUEUE_URL, handler: mock.fn(), functionName: "test-fn" },
+				],
+				log: buildLog(),
+				signals: ["SIGWINCH"],
+				// A no-op heartbeat keeps this case about the abort alone, off the
+				// real heartbeat file the default writer would touch.
+				createHeartbeat: () => async () => {},
+			});
+		} finally {
+			sendMock.mock.restore();
+		}
+
+		assert.ok(sawAbortSignal, "the receive was issued without an abort signal");
+		assert.equal(receiveCalls, 1);
+	});
+
 	it("gives every target its own heartbeat, beaten at the top of each receive attempt", async () => {
 		const beatsByQueue = new Map<string, number>();
 		let receiveCalls = 0;
@@ -177,6 +231,67 @@ describe("runQueuePoller", () => {
 		assert.match(String(failures[0]?.arguments[0]?.error), /ENOSPC/);
 	});
 
+	// The `requestId` every line of one invocation is grouped by
+	// (deploy/vps/README.md, "Logs") is minted here, and it has to reach both the
+	// handler and the poller's own bracketing lines — otherwise grouping on it
+	// stops right before the line saying what happened to the message.
+	it("gives each batch its own requestId, on the handler context and on its own lines", async () => {
+		let receiveCalls = 0;
+		const sendMock = mock.method(
+			SQSClient.prototype,
+			"send",
+			// biome-ignore lint/suspicious/noExplicitAny: minimal SDK command shape, not worth typing per-command here
+			async function (this: SQSClient, command: any) {
+				if (command.constructor.name === "ReceiveMessageCommand") {
+					receiveCalls++;
+					if (receiveCalls >= 2) process.emit("SIGWINCH", "SIGWINCH");
+					return { Messages: [buildMessage(`m${receiveCalls}`)] };
+				}
+				if (command.constructor.name === "DeleteMessageCommand") return {};
+				throw new Error(`unexpected command: ${command.constructor.name}`);
+			},
+		);
+
+		const handler = mock.fn(
+			async (
+				_event: SQSEvent,
+				_context: Context,
+			): Promise<SQSBatchResponse> => ({ batchItemFailures: [] }),
+		);
+		const log = buildLog();
+
+		try {
+			await runQueuePoller({
+				targets: [{ queueUrl: QUEUE_URL, handler, functionName: "test-fn" }],
+				log,
+				signals: ["SIGWINCH"],
+				createHeartbeat: () => async () => {},
+			});
+		} finally {
+			sendMock.mock.restore();
+		}
+
+		const requestIds = handler.mock.calls.map(
+			(call) => call.arguments[1].awsRequestId,
+		);
+		assert.equal(requestIds.length, 2);
+		for (const requestId of requestIds)
+			assert.match(requestId, /^[0-9a-f-]{36}$/);
+		assert.notEqual(requestIds[0], requestIds[1]);
+
+		const bracketing = ["poller: invoking handler", "poller: batch processed"];
+		for (const message of bracketing) {
+			const lines = log.info.mock.calls.filter(
+				(call) => call.arguments[1] === message,
+			);
+			assert.equal(lines.length, 2, `expected two ${message} lines`);
+			assert.deepEqual(
+				lines.map((call) => call.arguments[0].requestId),
+				requestIds,
+			);
+		}
+	});
+
 	it("throws when constructed with no targets", async () => {
 		await assert.rejects(
 			() => runQueuePoller({ targets: [], log: buildLog() }),

package/src/poller.ts

@@ -1,7 +1,9 @@
+import { randomUUID } from "node:crypto";
 import { inspect } from "node:util";
 import {
 	DeleteMessageCommand,
 	ReceiveMessageCommand,
+	type ReceiveMessageCommandOutput,
 } from "@aws-sdk/client-sqs";
 import type {
 	Context,
@@ -65,6 +67,7 @@ const pollTarget = async (
 	target: QueuePollerTarget,
 	log: QueuePollerLog,
 	isShuttingDown: () => boolean,
+	shutdownSignal: AbortSignal,
 	buildHeartbeat: (name: string) => Heartbeat,
 ): Promise<void> => {
 	const queueName = new URL(target.queueUrl).pathname.split("/").pop();
@@ -73,8 +76,6 @@ const pollTarget = async (
 	const maxMessages = target.maxMessages ?? DEFAULT_MAX_MESSAGES;
 	const visibilityTimeout =
 		target.visibilityTimeoutSeconds ?? DEFAULT_VISIBILITY_TIMEOUT_SECONDS;
-	const lambdaContext = { functionName: target.functionName } as Context;
-
 	log.info({ queue: queueName, maxMessages }, "poller: started");
 
 	// Always long-poll (SQS max WaitTimeSeconds). An empty long poll is free
@@ -101,15 +102,30 @@ const pollTarget = async (
 			);
 		});
 
-		const response = await sqs.send(
-			new ReceiveMessageCommand({
-				QueueUrl: target.queueUrl,
-				MaxNumberOfMessages: maxMessages,
-				WaitTimeSeconds: LONG_POLL_WAIT_SECONDS,
-				VisibilityTimeout: visibilityTimeout,
-				MessageSystemAttributeNames: ["ApproximateReceiveCount"],
-			}),
-		);
+		let response: ReceiveMessageCommandOutput;
+		try {
+			// Abort the receive on shutdown. Without this, a container stopped
+			// while idle waits out the current 20s long poll before the loop
+			// checks the flag again, and a self-update then burns the whole
+			// stop grace on every idle worker — the bulk of the visible restart
+			// downtime. The abort ends the in-flight receive at once; an
+			// in-flight handler below is never interrupted, so a message being
+			// processed still finishes cleanly. The heartbeat above already wrote
+			// this iteration, so the abort exits without a further beat.
+			response = await sqs.send(
+				new ReceiveMessageCommand({
+					QueueUrl: target.queueUrl,
+					MaxNumberOfMessages: maxMessages,
+					WaitTimeSeconds: LONG_POLL_WAIT_SECONDS,
+					VisibilityTimeout: visibilityTimeout,
+					MessageSystemAttributeNames: ["ApproximateReceiveCount"],
+				}),
+				{ abortSignal: shutdownSignal },
+			);
+		} catch (error) {
+			if (shutdownSignal.aborted) break;
+			throw error;
+		}
 
 		if (!response.Messages || response.Messages.length === 0) {
 			continue;
@@ -129,8 +145,26 @@ const pollTarget = async (
 		);
 		if (messages.length === 0) continue;
 
+		// One batch is one handler invocation, so the id is minted per batch: it is
+		// the `requestId` every line of that invocation is grouped by (see
+		// deploy/vps/README.md, "Logs"). A context built once per target would put
+		// the whole queue's lifetime under one id, which correlates nothing.
+		//
+		// Minted before the poller's own bracketing lines rather than next to the
+		// handler call: grouping on the id has to reach the line saying whether the
+		// message was deleted or left for redelivery, which is what someone looks
+		// for after a failure.
+		const lambdaContext = {
+			functionName: target.functionName,
+			awsRequestId: randomUUID(),
+		} as Context;
+
 		log.info(
-			{ queue: queueName, count: messages.length },
+			{
+				queue: queueName,
+				count: messages.length,
+				requestId: lambdaContext.awsRequestId,
+			},
 			"poller: invoking handler",
 		);
 
@@ -181,6 +215,7 @@ const pollTarget = async (
 				queue: queueName,
 				deleted: succeeded.length,
 				leftForRedelivery: failedIds.size,
+				requestId: lambdaContext.awsRequestId,
 			},
 			"poller: batch processed",
 		);
@@ -190,10 +225,11 @@ const pollTarget = async (
 };
 
 /**
- * Runs every target's poll loop concurrently until a shutdown signal is
- * received (default: SIGINT, SIGTERM), then lets each loop finish its
- * current iteration and returns. Rejects (crashes the process) if any loop
- * throws — a stuck poller should exit loudly, not degrade silently.
+ * received (default: SIGINT, SIGTERM). The signal aborts each loop's in-flight
+ * long-poll receive so an idle worker returns at once instead of waiting out
+ * the current 20s poll; a loop mid-handler finishes that handler first. Rejects
+ * (crashes the process) if any loop throws — a stuck poller should exit loudly,
+ * not degrade silently.
  *
  * Every loop writes its own heartbeat file (see heartbeat.ts) at the top of its
  * receive attempt, which is how a loop that hangs instead of throwing — the
@@ -219,9 +255,13 @@ export const runQueuePoller = async (
 
 	let shuttingDown = false;
 	const isShuttingDown = () => shuttingDown;
+	// Aborted alongside the flag so a loop currently blocked in its long-poll
+	// receive returns at once, rather than after up to WaitTimeSeconds.
+	const shutdown = new AbortController();
 	const onSignal = (signal: NodeJS.Signals) => {
 		log.info({ signal }, "poller: shutdown signal received");
 		shuttingDown = true;
+		shutdown.abort();
 	};
 	for (const signal of signals) {
 		process.on(signal, onSignal);
@@ -231,7 +271,7 @@ export const runQueuePoller = async (
 
 	await Promise.all(
 		targets.map((target) =>
-			pollTarget(target, log, isShuttingDown, buildHeartbeat),
+			pollTarget(target, log, isShuttingDown, shutdown.signal, buildHeartbeat),
 		),
 	);
 };