@remit/search-index-worker 0.0.22 → 0.0.24

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@remit/search-index-worker",
-  "version": "0.0.22",
+  "version": "0.0.24",
   "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/services.ts

@@ -8,6 +8,8 @@ import {
 import {
 	buildEmbeddingServiceFromEnv,
 	buildVectorStoreFromEnv,
+	EMBEDDING_PROVIDER_OFF,
+	readEmbeddingProviderFromEnv,
 } from "@remit/search-service/from-env";
 import { createHeartbeat } from "@remit/sqs-client/heartbeat";
 import type { StorageService } from "@remit/storage-service";
@@ -74,7 +76,7 @@ let governorResolved = false;
 export const getMemoryGovernor = (): MemoryGovernor | undefined => {
 	if (governorResolved) return governor;
 	governorResolved = true;
-	if (process.env.SEARCH_EMBEDDING_PROVIDER !== "local") return undefined;
+	if (readEmbeddingProviderFromEnv() !== "local") return undefined;
 
 	const config = readAdaptiveEmbeddingConfigFromEnv();
 	const metrics = registerAdaptiveEmbedding({
@@ -118,6 +120,19 @@ const governed = (embedder: EmbeddingService): EmbeddingService => {
 export const getServices = async (): Promise<Services> => {
 	if (cached) return cached;
 
+	// This process exists to embed. `off` is the self-host default and holds the
+	// container down behind the `semantic` compose profile
+	// (deploy/vps/docker-compose.sqlite.yml), so reaching here with it set means
+	// the worker was started against a deployment that asked for no embedding:
+	// every message it took off the queue would fail one at a time, forever. Say
+	// so once, at startup, and name the command that settles it.
+	if (readEmbeddingProviderFromEnv() === EMBEDDING_PROVIDER_OFF) {
+		throw new Error(
+			"SEARCH_EMBEDDING_PROVIDER is off, so there is nothing for this worker to embed. " +
+				"Turn semantic search on with 'remit semantic on', or leave this service down.",
+		);
+	}
+
 	const dataPorts = await buildDataPortsFromEnv();
 
 	const storageService = createStorageService();

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

@@ -113,3 +113,38 @@ describe("SqliteOutboxStore", () => {
 		db.close();
 	});
 });
+
+// `remit semantic on` against an existing mailbox makes the first drain pass the
+// whole back catalogue, and the pass re-runs every 2 s. The read is bounded so
+// one tick relays a batch rather than tens of thousands of rows; what it does not
+// take stays unprocessed and is what the next tick selects.
+describe("SqliteOutboxStore, bounded", () => {
+	test("reads at most one batch, and the rest on the next pass", async () => {
+		const db = makeOutboxDb();
+		const total = 620;
+		for (let i = 0; i < total; i++) {
+			insertRow(db, `r${i}`, `m${i}`, "message.body_synced");
+		}
+		const store = new SqliteOutboxStore(db as unknown as never);
+
+		const first = await store.listUnprocessedEvents();
+		assert.equal(first.length, 500);
+
+		const sent: string[] = [];
+		const relay = new OutboxRelay({
+			store,
+			sqs: fakeSqs(sent),
+			queueUrl: "q",
+		});
+
+		assert.equal(await relay.drainPending(), 500);
+		assert.equal(await relay.drainPending(), total - 500);
+		assert.equal(await relay.drainPending(), 0);
+
+		const stillPending = db
+			.prepare("SELECT count(*) AS n FROM outbox WHERE processed_at IS NULL")
+			.get() as { n: number };
+		assert.equal(stillPending.n, 0, "every row was relayed, none skipped");
+		db.close();
+	});
+});

package/src/sqlite-outbox-drain.ts

@@ -24,6 +24,13 @@ import {
 
 const DRAIN_INTERVAL_MS = 2_000;
 
+// How many distinct messages one drain pass relays. `remit semantic on` on an
+// existing mailbox makes the first pass the whole back catalogue — tens of
+// thousands of rows read, enqueued and marked in one tick — so the pass is
+// bounded and the next tick, 2 s later, takes the following batch. Nothing is
+// dropped: unmarked rows stay unprocessed and are what the next pass selects.
+const DRAIN_BATCH_SIZE = 500;
+
 // A minimal view of the better-sqlite3 surface used here, so the module carries
 // no static type dependency on the native package (imported dynamically to stay
 // out of the Lambda bundle).
@@ -46,7 +53,9 @@ export class SqliteOutboxStore implements OutboxStore {
 		const rows = this.db
 			.prepare(
 				`SELECT DISTINCT message_id, event FROM outbox
-				 WHERE event IN (${placeholders}) AND processed_at IS NULL`,
+				 WHERE event IN (${placeholders}) AND processed_at IS NULL
+				 ORDER BY message_id
+				 LIMIT ${DRAIN_BATCH_SIZE}`,
 			)
 			.all(...DRAIN_EVENTS) as Array<{ message_id: string; event: string }>;
 		return rows.map((row) => ({