@remit/search-index-worker 0.0.23 → 0.0.25

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@remit/search-index-worker",
-  "version": "0.0.23",
+  "version": "0.0.25",
   "type": "module",
   "main": "dist/index.js",
   "exports": {

package/src/adaptive-embedder.test.ts

@@ -1,6 +1,7 @@
 import assert from "node:assert/strict";
 import { describe, it } from "node:test";
 import type { EmbeddingService } from "@remit/search-service";
+import { DEFAULT_VISIBILITY_TIMEOUT_SECONDS } from "@remit/sqs-client/poller";
 import {
 	type AdaptiveEmbeddingConfig,
 	createAdaptiveEmbeddingService,
@@ -53,10 +54,28 @@ const reads = (script: readonly Reading[]): MemoryReader => {
 const available = (...availableMb: number[]): MemoryReader =>
 	reads(availableMb.map((mb) => ({ availableMb: mb })));
 
+/**
+ * A box that sits below the critical floor, frees memory for exactly one
+ * reading, and dips again — the intermittent shape, as against a sustained
+ * stall. `lowReads` sets how long each dip lasts.
+ */
+const recoveringBox = (lowReads: number): MemoryReader => {
+	let lowLeft = lowReads;
+	return () => {
+		const low = lowLeft > 0;
+		lowLeft = low ? lowLeft - 1 : lowReads;
+		return {
+			availableBytes: (low ? 300 : ROOMY) * MB,
+			rssBytes: 512 * MB,
+		};
+	};
+};
+
 class Harness {
 	readonly plans: EmbeddingPlan[] = [];
 	readonly sleeps: number[] = [];
 	readonly lines: string[] = [];
+	readonly logged: { message: string; fields?: Record<string, unknown> }[] = [];
 	stalls = 0;
 	beats = 0;
 	clock = 0;
@@ -64,8 +83,9 @@ class Harness {
 	readonly deps: GovernorDeps;
 
 	constructor(readMemory: MemoryReader) {
-		const record = (message: string) => {
+		const record = (message: string, fields?: Record<string, unknown>) => {
 			this.lines.push(message);
+			this.logged.push({ message, fields });
 		};
 		this.deps = {
 			readMemory,
@@ -123,6 +143,26 @@ const embedderRecording = (): {
 	};
 };
 
+/** An embedder whose work advances the fake clock, so waves spend the budget. */
+const embedderCosting = (
+	h: Harness,
+	costMs: number,
+): { service: EmbeddingService; batches: number[] } => {
+	const batches: number[] = [];
+	return {
+		batches,
+		service: {
+			dimensions: 3,
+			embeddingId: "fake@3",
+			embed: async (texts: string[]) => {
+				h.clock += costMs;
+				batches.push(texts.length);
+				return texts.map(() => [0, 0, 0]);
+			},
+		},
+	};
+};
+
 describe("the memory governor", () => {
 	it("starts at the floor: the smallest batch, one inference", () => {
 		const governor = new MemoryGovernor(CONFIG, harness(available(ROOMY)).deps);
@@ -223,10 +263,16 @@ describe("the memory governor", () => {
 		const governor = new MemoryGovernor(CONFIG, h.deps);
 		settleTimes(governor, CONFIG.rampAfterReadings + 1);
 
-		assert.equal(await governor.admit(), "admitted");
+		assert.equal(
+			(await governor.admit(governor.stallDeadline())).status,
+			"admitted",
+		);
 		assert.deepEqual(h.sleeps, [CONFIG.pauseMs]);
 		// The pause is per shed, not sticky: an admit that follows no shed runs on.
-		assert.equal(await governor.admit(), "admitted");
+		assert.equal(
+			(await governor.admit(governor.stallDeadline())).status,
+			"admitted",
+		);
 		assert.deepEqual(h.sleeps, [CONFIG.pauseMs]);
 	});
 
@@ -236,7 +282,10 @@ describe("the memory governor", () => {
 		governor.settle();
 		assert.deepEqual(governor.plan, { batchSize: 2, concurrency: 1 });
 		assert.deepEqual(h.lines, []);
-		assert.equal(await governor.admit(), "admitted");
+		assert.equal(
+			(await governor.admit(governor.stallDeadline())).status,
+			"admitted",
+		);
 		assert.deepEqual(h.sleeps, [CONFIG.pauseMs]);
 	});
 
@@ -246,7 +295,10 @@ describe("the memory governor", () => {
 		settleTimes(governor, CONFIG.rampAfterReadings);
 		assert.deepEqual(governor.plan, { batchSize: 4, concurrency: 1 });
 
-		assert.equal(await governor.admit(), "admitted");
+		assert.equal(
+			(await governor.admit(governor.stallDeadline())).status,
+			"admitted",
+		);
 		assert.equal(h.stalls, 1);
 		assert.equal(h.sleeps.length, 2);
 		// A stop is the loudest signal the worker has, and it comes back at the
@@ -262,7 +314,10 @@ describe("the memory governor", () => {
 		const h = harness(available(300));
 		const governor = new MemoryGovernor(CONFIG, h.deps);
 
-		assert.equal(await governor.admit(), "expired");
+		assert.equal(
+			(await governor.admit(governor.stallDeadline())).status,
+			"expired",
+		);
 		assert.equal(h.clock, CONFIG.stallMaxMs);
 		assert.match(h.lines.at(-1) ?? "", /gave up/);
 	});
@@ -271,7 +326,7 @@ describe("the memory governor", () => {
 		const h = harness(available(300));
 		const governor = new MemoryGovernor(CONFIG, h.deps);
 
-		await governor.admit();
+		await governor.admit(governor.stallDeadline());
 		assert.equal(h.beats, CONFIG.stallMaxMs / CONFIG.pauseMs);
 	});
 
@@ -280,14 +335,20 @@ describe("the memory governor", () => {
 		h.beatFails = true;
 		const governor = new MemoryGovernor(CONFIG, h.deps);
 
-		assert.equal(await governor.admit(), "expired");
+		assert.equal(
+			(await governor.admit(governor.stallDeadline())).status,
+			"expired",
+		);
 		assert.ok(h.lines.some((line) => /heartbeat/.test(line)));
 	});
 
 	it("does not stall while the box is merely tight", async () => {
 		const h = harness(available(500));
 		const governor = new MemoryGovernor(CONFIG, h.deps);
-		assert.equal(await governor.admit(), "admitted");
+		assert.equal(
+			(await governor.admit(governor.stallDeadline())).status,
+			"admitted",
+		);
 		assert.equal(h.stalls, 0);
 		assert.deepEqual(h.sleeps, []);
 	});
@@ -384,6 +445,55 @@ describe("the governed embedder", () => {
 		assert.deepEqual(inner.batches, []);
 	});
 
+	// The budget bounds the handler, not one wave of it. A box that recovers just
+	// inside it on every wave held the record for waves × budget, well past the
+	// queue's visibility timeout, and reported no failure.
+	it("holds the budget across waves when memory recovers between them", async () => {
+		const inner = embedderRecording();
+		const h = harness(recoveringBox(CONFIG.stallMaxMs / CONFIG.pauseMs - 1));
+		const service = createAdaptiveEmbeddingService(
+			inner.service,
+			new MemoryGovernor(CONFIG, h.deps),
+			CONFIG.stallMaxMs,
+		);
+
+		await assert.rejects(
+			() => service.embed(Array.from({ length: 20 }, (_, i) => `chunk ${i}`)),
+			(error: unknown) => error instanceof MemoryStallTimeoutError,
+		);
+		// Memory did come back and work did get through, so this is the
+		// intermittent stall rather than the sustained one.
+		assert.ok(h.stalls > 1);
+		assert.ok(inner.batches.length > 0);
+		assert.ok(h.clock <= CONFIG.stallMaxMs + CONFIG.pauseMs);
+	});
+
+	// Waves spend the same budget the stall does, so a message can fail on a dip
+	// that arrives late and is short. What it reports has to be the dip it
+	// measured, not the budget: the two are no longer the same number.
+	it("spends the budget on the waves themselves, and says what it waited", async () => {
+		const h = harness(available(ROOMY, ROOMY, ROOMY, ROOMY, ROOMY, ROOMY, 300));
+		const inner = embedderCosting(h, 30);
+		const service = createAdaptiveEmbeddingService(
+			inner.service,
+			new MemoryGovernor(CONFIG, h.deps),
+			CONFIG.stallMaxMs,
+		);
+
+		await assert.rejects(
+			() => service.embed(Array.from({ length: 20 }, (_, i) => `chunk ${i}`)),
+			(error: unknown) =>
+				error instanceof MemoryStallTimeoutError &&
+				error.waitedMs === CONFIG.pauseMs &&
+				/ran out of its/.test(error.message),
+		);
+		assert.equal(inner.batches.length, 3);
+		assert.equal(h.clock, CONFIG.stallMaxMs);
+		assert.equal(h.stalls, 1);
+		const gaveUp = h.logged.find((line) => /gave up/.test(line.message));
+		assert.equal(gaveUp?.fields?.waitedMs, CONFIG.pauseMs);
+	});
+
 	// The throttle paces work; it never turns a fault into a quiet retry.
 	it("lets a model failure propagate", async () => {
 		const h = harness(available(ROOMY));
@@ -459,7 +569,10 @@ describe("the configured thresholds", () => {
 	// The stall has to end before the queue redelivers the record underneath it.
 	it("gives up well inside the poller's 300 s visibility timeout", () => {
 		withEnv(UNSET, () => {
-			assert.ok(readAdaptiveEmbeddingConfigFromEnv().stallMaxMs < 300_000);
+			assert.ok(
+				readAdaptiveEmbeddingConfigFromEnv().stallMaxMs <
+					DEFAULT_VISIBILITY_TIMEOUT_SECONDS * 1000,
+			);
 		});
 	});
 
@@ -512,6 +625,25 @@ describe("the configured thresholds", () => {
 		);
 	});
 
+	// A budget that reaches the visibility timeout has the record redelivered
+	// underneath the handler still holding it: the redelivery it exists to avoid.
+	it("refuses a stall budget at or above the visibility timeout", () => {
+		withEnv(
+			{
+				...UNSET,
+				SEARCH_INDEX_MEMORY_STALL_MAX_MS: String(
+					DEFAULT_VISIBILITY_TIMEOUT_SECONDS * 1000,
+				),
+			},
+			() => {
+				assert.throws(
+					readAdaptiveEmbeddingConfigFromEnv,
+					/SEARCH_INDEX_MEMORY_STALL_MAX_MS must be below/,
+				);
+			},
+		);
+	});
+
 	it("refuses a value that is not a positive integer", () => {
 		withEnv({ ...UNSET, SEARCH_INDEX_EMBED_BATCH_MAX: "0" }, () => {
 			assert.throws(

package/src/adaptive-embedder.ts

@@ -1,4 +1,5 @@
 import type { EmbeddingService } from "@remit/search-service";
+import { DEFAULT_VISIBILITY_TIMEOUT_SECONDS } from "@remit/sqs-client/poller";
 import type { MemoryReader, MemoryReading } from "./memory.js";
 
 /**
@@ -76,10 +77,14 @@ export const DEFAULT_ADAPTIVE_EMBEDDING_CONFIG: AdaptiveEmbeddingConfig = {
  */
 export class MemoryStallTimeoutError extends Error {
 	readonly code = "ERR_SEARCH_INDEX_MEMORY_STALL";
-	constructor(waitedMs: number) {
+	constructor(
+		stallMaxMs: number,
+		readonly waitedMs: number,
+	) {
 		super(
-			`Search index waited ${Math.round(waitedMs / 1000)}s for the box to ` +
-				"free memory and gave up; the message goes back on the queue",
+			`Search index ran out of its ${Math.round(stallMaxMs / 1000)}s budget ` +
+				`for this message with the box below the memory floor for the last ` +
+				`${Math.round(waitedMs / 1000)}s; the message goes back on the queue`,
 		);
 		this.name = "MemoryStallTimeoutError";
 	}
@@ -100,11 +105,21 @@ const fromEnv = (name: string, fallback: number, scale = 1): number => {
 	return positiveInt(name, raw) * scale;
 };
 
+/**
+ * A stall budget that reaches the queue's visibility timeout has the record
+ * redelivered underneath the handler still holding it, which is the redelivery
+ * the budget exists to prevent. The poller's default is the authority: the
+ * search index queue passes no override, so what `deploy/vps/queues.json` sets
+ * on the queue never reaches this code and has to match it by hand.
+ */
+const VISIBILITY_TIMEOUT_MS = DEFAULT_VISIBILITY_TIMEOUT_SECONDS * 1000;
+
 /**
  * Every threshold is an env var so the same image bounds itself against a 4 GB
  * VPS and a 32 GB box without a rebuild. A configuration that cannot hold —
  * a critical floor at or above the ramp headroom, a floor batch above the
- * ceiling — is a startup error, not something to correct silently at runtime.
+ * ceiling, a stall budget at or above the visibility timeout — is a startup
+ * error, not something to correct silently at runtime.
  */
 export const readAdaptiveEmbeddingConfigFromEnv =
 	(): AdaptiveEmbeddingConfig => {
@@ -153,10 +168,22 @@ export const readAdaptiveEmbeddingConfigFromEnv =
 				"SEARCH_INDEX_MEMORY_CRITICAL_MB must be below SEARCH_INDEX_MEMORY_HEADROOM_MB",
 			);
 		}
+		if (config.stallMaxMs >= VISIBILITY_TIMEOUT_MS) {
+			throw new Error(
+				`SEARCH_INDEX_MEMORY_STALL_MAX_MS must be below the queue's ${VISIBILITY_TIMEOUT_MS} ms visibility timeout`,
+			);
+		}
 		return config;
 	};
 
-type Admission = "admitted" | "expired";
+/**
+ * `waitedMs` is the time this stop actually spent below the floor, which the
+ * budget no longer stands in for: embedding work spends the same budget, so a
+ * message can fail on a two-second dip that arrived late.
+ */
+type Admission =
+	| { readonly status: "admitted" }
+	| { readonly status: "expired"; readonly waitedMs: number };
 
 /**
  * Bounds the worker's resident memory against the box it shares, which
@@ -181,7 +208,9 @@ type Admission = "admitted" | "expired";
  * swap, where the kernel picks its OOM victim by size and takes the backend
  * rather than the indexer. That stop is bounded: past the budget the message
  * goes back on the queue, because a handler that waits longer than the queue's
- * visibility timeout has its record redelivered underneath it anyway.
+ * visibility timeout has its record redelivered underneath it anyway. The
+ * budget belongs to the message, so its deadline is handed in rather than
+ * minted here — one that restarted per stop would bound no handler at all.
  */
 export class MemoryGovernor {
 	private batchSize: number;
@@ -200,14 +229,17 @@ export class MemoryGovernor {
 		return { batchSize: this.batchSize, concurrency: this.concurrency };
 	}
 
+	/** The end of one message's stall budget, taken once per governed call. */
+	stallDeadline = (): number => this.deps.now() + this.config.stallMaxMs;
+
 	/** Blocks until the box can afford the next batch, or the budget runs out. */
-	admit = async (): Promise<Admission> => {
+	admit = async (stallDeadline: number): Promise<Admission> => {
 		let reading = this.deps.readMemory();
 		if (reading.availableBytes >= this.config.criticalBytes) {
-			if (!this.pauseBeforeNextBatch) return "admitted";
+			if (!this.pauseBeforeNextBatch) return { status: "admitted" };
 			this.pauseBeforeNextBatch = false;
 			await this.deps.sleep(this.config.pauseMs);
-			return "admitted";
+			return { status: "admitted" };
 		}
 
 		this.deps.onStall?.();
@@ -218,13 +250,15 @@ export class MemoryGovernor {
 		this.reset();
 		const startedAt = this.deps.now();
 		while (reading.availableBytes < this.config.criticalBytes) {
-			const waitedMs = this.deps.now() - startedAt;
-			if (waitedMs >= this.config.stallMaxMs) {
-				this.deps.log.warn("Search index gave up waiting for memory", {
+			const now = this.deps.now();
+			if (now >= stallDeadline) {
+				const waitedMs = now - startedAt;
+				this.deps.log.warn("Search index gave up: its memory budget ran out", {
 					waitedMs,
+					stallMaxMs: this.config.stallMaxMs,
 					...this.fields(reading),
 				});
-				return "expired";
+				return { status: "expired", waitedMs };
 			}
 			await this.keepAlive();
 			await this.deps.sleep(this.config.pauseMs);
@@ -235,7 +269,7 @@ export class MemoryGovernor {
 			"Search index resumed: memory recovered",
 			this.fields(reading),
 		);
-		return "admitted";
+		return { status: "admitted" };
 	};
 
 	/** Measures what the batch just cost and moves the plan at most one step. */
@@ -342,6 +376,11 @@ export class MemoryGovernor {
  * governed batches and holds only the current wave's inputs, so the model's
  * outputs from a finished batch are unreachable before the next one starts.
  *
+ * The stall budget is taken once here and spans every wave, because what has to
+ * stay inside the queue's visibility timeout is the handler, not one wave of it:
+ * a box that recovers just inside the budget on each wave would otherwise hold
+ * the record for as many budgets as the email has chunks.
+ *
  * `dimensions` and `embeddingId` pass straight through: `embeddingId` feeds the
  * content hash that decides what needs re-embedding, so wrapping the embedder
  * must not invalidate an existing index.
@@ -355,10 +394,12 @@ export const createAdaptiveEmbeddingService = (
 	embeddingId: inner.embeddingId,
 	embed: async (texts: string[]): Promise<number[][]> => {
 		const vectors: number[][] = [];
+		const deadline = governor.stallDeadline();
 		let next = 0;
 		while (next < texts.length) {
-			if ((await governor.admit()) === "expired") {
-				throw new MemoryStallTimeoutError(stallMaxMs);
+			const admission = await governor.admit(deadline);
+			if (admission.status === "expired") {
+				throw new MemoryStallTimeoutError(stallMaxMs, admission.waitedMs);
 			}
 			const { batchSize, concurrency } = governor.plan;
 			const wave: string[][] = [];

package/src/handler.test.ts

@@ -68,7 +68,7 @@ describe("a message the worker cannot index right now", () => {
 	it("goes back on the queue when indexing stalls out of memory", async () => {
 		const response = await processBatch(
 			[record()],
-			servicesThatFail(new MemoryStallTimeoutError(240_000)),
+			servicesThatFail(new MemoryStallTimeoutError(240_000, 2_000)),
 			silent,
 		);
 

package/src/outbox-drain-coverage.test.ts

@@ -0,0 +1,138 @@
+import assert from "node:assert/strict";
+import { randomUUID } from "node:crypto";
+import { describe, test } from "node:test";
+import type { SendMessageCommand } from "@aws-sdk/client-sqs";
+import {
+	DrizzleMessageRepository,
+	OUTBOX_EVENTS,
+} from "@remit/drizzle-service";
+import { createShippedSqliteDb } from "@remit/drizzle-service/test-sqlite";
+import { OutboxRelay } from "@remit/outbox-relay";
+import { SqliteOutboxStore } from "./sqlite-outbox-drain.js";
+
+type SqliteHandle = ReturnType<typeof createShippedSqliteDb>["sqlite"];
+
+// The outbox invariant (reader#1063): every event kind a producer writes has a
+// consumer that drains it. `message.created` was written once per message and
+// drained by nothing, so a live instance carried 30,906 rows that could never
+// clear. This runs the real producers (the message repo) and the real consumer
+// (the relay over the SQLite store) against the shipped migrations, so a kind
+// with no drain shows up as a row that stays undrained.
+//
+// What holds the invariant is the union `OUTBOX_EVENTS`, not this file: the
+// column's type is what stops a producer in another package inventing a kind,
+// and the compiler is what reports it. This covers the kinds the union declares
+// and the producers in this repo's message repo — a producer elsewhere reaches
+// it only by adding its kind to the union first.
+
+const fakeSqs = (sent: string[]) =>
+	({
+		send: async (cmd: SendMessageCommand) => {
+			sent.push(String(cmd.input.MessageBody));
+			return {};
+		},
+	}) as unknown as ConstructorParameters<typeof OutboxRelay>[0]["sqs"];
+
+const drainAll = async (sqlite: SqliteHandle): Promise<void> => {
+	const relay = new OutboxRelay({
+		store: new SqliteOutboxStore(sqlite as unknown as never),
+		sqs: fakeSqs([]),
+		queueUrl: "q",
+	});
+	while ((await relay.drainPending()) > 0) {
+		// Drain reads a bounded batch; repeat until the table is quiet.
+	}
+};
+
+const undrainedEvents = (sqlite: SqliteHandle): string[] =>
+	(
+		sqlite
+			.prepare(
+				"SELECT DISTINCT event FROM outbox WHERE processed_at IS NULL ORDER BY event",
+			)
+			.all() as Array<{ event: string }>
+	).map((row) => row.event);
+
+const writtenEvents = (sqlite: SqliteHandle): string[] =>
+	(
+		sqlite
+			.prepare("SELECT DISTINCT event FROM outbox ORDER BY event")
+			.all() as Array<{ event: string }>
+	).map((row) => row.event);
+
+const NOW = 1700000000000;
+const MAILBOX_ID = "00000000-0000-0000-4444-000000000001";
+const DEST_MAILBOX_ID = "00000000-0000-0000-4444-000000000002";
+
+const createInput = (messageId: string) => ({
+	messageId,
+	mailboxId: MAILBOX_ID,
+	uid: 1,
+	sequenceNumber: 1,
+	rfc822Size: 512,
+	internalDate: NOW,
+	envelopeId: randomUUID(),
+	rootBodyPartId: randomUUID(),
+});
+
+describe("outbox drain coverage", () => {
+	test("every event the message repo writes is drained", async () => {
+		const { db, sqlite, close } = createShippedSqliteDb();
+		const repo = new DrizzleMessageRepository(
+			db as unknown as ConstructorParameters<
+				typeof DrizzleMessageRepository
+			>[0],
+		);
+
+		const bodySyncedId = randomUUID();
+		await repo.create(createInput(bodySyncedId));
+		await repo.update(bodySyncedId, { bodyStorageKey: "body/1.json" });
+
+		const movedId = randomUUID();
+		await repo.create(createInput(movedId));
+		await repo.updateUid(movedId, 99, DEST_MAILBOX_ID);
+
+		const removedId = randomUUID();
+		await repo.create(createInput(removedId));
+		await repo.delete(removedId);
+
+		assert.deepEqual(
+			writtenEvents(sqlite),
+			[...OUTBOX_EVENTS].sort(),
+			"the producers write exactly the declared outbox vocabulary",
+		);
+
+		await drainAll(sqlite);
+
+		assert.deepEqual(
+			undrainedEvents(sqlite),
+			[],
+			"a kind left undrained here grows the outbox forever",
+		);
+		close();
+	});
+
+	test("every kind the union declares is drained", async () => {
+		const { sqlite, close } = createShippedSqliteDb();
+
+		for (const event of OUTBOX_EVENTS) {
+			sqlite
+				.prepare(
+					`INSERT INTO outbox (id, message_id, event, payload, created_at)
+					 VALUES (?, ?, ?, ?, ?)`,
+				)
+				.run(
+					randomUUID(),
+					`m-${event}`,
+					event,
+					JSON.stringify({ messageId: `m-${event}` }),
+					Date.now(),
+				);
+		}
+
+		await drainAll(sqlite);
+
+		assert.deepEqual(undrainedEvents(sqlite), []);
+		close();
+	});
+});