fix(worker): consume live work while tasks are active

packages/pgconductor-js/src/lib/async-queue.ts

@@ -1,11 +1,13 @@
 export interface PollableAsyncIterable<T> extends AsyncIterable<T> {
 	tryNext(): T | undefined;
+	onNextItemAvailable?(listener: () => void): () => void;
 }
 
 export class AsyncQueue<T> implements PollableAsyncIterable<T> {
 	private queue: T[] = [];
 	private resolvers: ((value: IteratorResult<T>) => void)[] = [];
 	private pushResolvers: (() => void)[] = [];
+	private itemAvailableListeners = new Set<() => void>();
 	private closed = false;
 
 	constructor(private readonly capacity: number) {}
@@ -30,6 +32,7 @@ export class AsyncQueue<T> implements PollableAsyncIterable<T> {
 			// No consumer waiting, add to queue for later
 			this.queue.push(item);
 		}
+		this.notifyItemAvailable();
 	}
 
 	tryNext(): T | undefined {
@@ -51,6 +54,21 @@ export class AsyncQueue<T> implements PollableAsyncIterable<T> {
 		return new Promise((resolve) => this.resolvers.push(resolve));
 	}
 
+	onNextItemAvailable(listener: () => void): () => void {
+		if (this.queue.length > 0) {
+			listener();
+			return () => {};
+		}
+		this.itemAvailableListeners.add(listener);
+		return () => this.itemAvailableListeners.delete(listener);
+	}
+
+	private notifyItemAvailable() {
+		const listeners = [...this.itemAvailableListeners];
+		this.itemAvailableListeners.clear();
+		for (const listener of listeners) listener();
+	}
+
 	private notifyPusher() {
 		// If there's space and someone waiting to push, wake them
 		if (this.queue.length < this.capacity && this.pushResolvers.length > 0) {

packages/pgconductor-js/src/lib/batching-async-queue.ts

@@ -20,6 +20,7 @@ export class BatchingAsyncQueue<T extends { task_key: string }> implements Polla
 > {
 	private queue: BatchGroup<T>[] = [];
 	private resolvers: ((value: IteratorResult<BatchGroup<T>>) => void)[] = [];
+	private itemAvailableListeners = new Set<() => void>();
 	private closed = false;
 
 	// Track batches per task
@@ -117,6 +118,7 @@ export class BatchingAsyncQueue<T extends { task_key: string }> implements Polla
 			// Backpressure is managed at push() level, not here
 			this.queue.push(group);
 		}
+		this.notifyItemAvailable();
 	}
 
 	tryNext(): BatchGroup<T> | undefined {
@@ -136,6 +138,21 @@ export class BatchingAsyncQueue<T extends { task_key: string }> implements Polla
 		return new Promise((resolve) => this.resolvers.push(resolve));
 	}
 
+	onNextItemAvailable(listener: () => void): () => void {
+		if (this.queue.length > 0) {
+			listener();
+			return () => {};
+		}
+		this.itemAvailableListeners.add(listener);
+		return () => this.itemAvailableListeners.delete(listener);
+	}
+
+	private notifyItemAvailable() {
+		const listeners = [...this.itemAvailableListeners];
+		this.itemAvailableListeners.clear();
+		for (const listener of listeners) listener();
+	}
+
 	close(): void {
 		if (this.closed) return;
 

packages/pgconductor-js/src/lib/map-concurrent.ts

@@ -8,14 +8,17 @@ export async function* mapConcurrent<T, R>(
 	limit: number,
 	mapper: (item: T) => Promise<R>,
 ): AsyncGenerator<R> {
+	type Task = { id: number; promise: Promise<R> };
+	type RaceResult = { id: number; result: R } | { id: null };
+	type ItemWait = { promise: Promise<RaceResult>; cancel(): void };
+
 	const it = source[Symbol.asyncIterator]();
+	const onNextItemAvailable = source.onNextItemAvailable?.bind(source);
+	const active = new Map<number, Task>();
 	let sourceDone = false;
+	let itemWait: ItemWait | null = null;
 	let nextId = 0;
 
-	// Track promises with unique IDs
-	type Task = { id: number; promise: Promise<R> };
-	const active = new Map<number, Task>();
-
 	const nextItem = async (): Promise<T | null> => {
 		if (sourceDone) return null;
 		const { value, done } = await it.next();
@@ -52,30 +55,39 @@ export async function* mapConcurrent<T, R>(
 
 	await fillSlots();
 
-	while (active.size > 0) {
-		// Wrap each promise to include its ID
-		const wrappedPromises = Array.from(active.values()).map(async (task) => ({
-			id: task.id,
-			result: await task.promise,
-		}));
+	try {
+		while (active.size > 0) {
+			if (!sourceDone && active.size < limit && !itemWait && onNextItemAvailable) {
+				let cancel = () => {};
+				const promise = new Promise<void>((resolve) => {
+					cancel = onNextItemAvailable(resolve);
+				}).then(() => ({ id: null }) as const);
+				itemWait = { promise, cancel: () => cancel() };
+			}
 
-		// Race to get first completed task
-		const { id, result } = await Promise.race(wrappedPromises);
+			const wrappedPromises: Promise<RaceResult>[] = Array.from(active.values()).map(
+				async (task) => ({ id: task.id, result: await task.promise }),
+			);
+			if (itemWait) wrappedPromises.push(itemWait.promise);
 
-		// Remove the completed task
-		active.delete(id);
+			// Race to get completed work or newly available input
+			const event = await Promise.race(wrappedPromises);
 
-		yield result;
+			if (event.id === null) {
+				itemWait = null;
+				await fillSlots();
+				continue;
+			}
 
-		// Refill slots
-		await fillSlots();
-	}
+			// Remove the completed task
+			active.delete(event.id);
+
+			yield event.result;
 
-	if (typeof it.return === "function") {
-		try {
-			await it.return();
-		} catch {
-			// ignore
+			// Refill slots
+			await fillSlots();
 		}
+	} finally {
+		itemWait?.cancel();
 	}
 }

packages/pgconductor-js/tests/unit/lib/map-concurrent.test.ts

@@ -1,6 +1,8 @@
 import { test, expect, describe } from "bun:test";
 import { mapConcurrent } from "../../../src/lib/map-concurrent";
-import type { PollableAsyncIterable } from "../../../src/lib/async-queue";
+import { AsyncQueue, type PollableAsyncIterable } from "../../../src/lib/async-queue";
+import { BatchingAsyncQueue } from "../../../src/lib/batching-async-queue";
+import { Deferred } from "../../../src/lib/deferred";
 
 class PollableGenerator<T> implements PollableAsyncIterable<T> {
 	private buffer: T[] = [];
@@ -57,6 +59,78 @@ describe("mapConcurrent", () => {
 		expect(results).toEqual([0, 2, 4, 6, 8]);
 	});
 
+	test("starts items that arrive while another mapper is running", async () => {
+		const source = new AsyncQueue<number>(2);
+		const firstStarted = new Deferred<void>();
+		const secondStarted = new Deferred<void>();
+		const releaseFirst = new Deferred<void>();
+
+		const consuming = (async () => {
+			for await (const _ of mapConcurrent(source, 2, async (value) => {
+				if (value === 1) {
+					firstStarted.resolve();
+					await releaseFirst.promise;
+				} else {
+					secondStarted.resolve();
+				}
+				return value;
+			})) {
+				// Consume all results.
+			}
+		})();
+
+		try {
+			await source.push(1);
+			await firstStarted.promise;
+			await source.push(2);
+
+			const startedBeforeFirstCompleted = await Promise.race([
+				secondStarted.promise.then(() => true),
+				Bun.sleep(100).then(() => false),
+			]);
+			expect(startedBeforeFirstCompleted).toBe(true);
+		} finally {
+			releaseFirst.resolve();
+			source.close();
+			await consuming;
+		}
+	});
+
+	test("leaves the source usable when the consumer exits early", async () => {
+		const source = new AsyncQueue<number>(2);
+		await source.push(1);
+
+		for await (const result of mapConcurrent(source, 2, async (value) => value)) {
+			expect(result).toBe(1);
+			break;
+		}
+
+		await source.push(2);
+		expect(await source.next()).toEqual({ value: 2, done: false });
+		source.close();
+	});
+
+	test("preserves pending batches when the consumer exits early", async () => {
+		const source = new BatchingAsyncQueue<{ task_key: string; id: number }>(
+			2,
+			new Map([["batched", { size: 2, timeoutMs: 60_000 }]]),
+		);
+		await source.push({ task_key: "immediate", id: 1 });
+		await source.push({ task_key: "batched", id: 2 });
+
+		for await (const result of mapConcurrent(source, 1, async (group) => group)) {
+			expect(result.items.map((item) => item.id)).toEqual([1]);
+			break;
+		}
+
+		await source.push({ task_key: "batched", id: 3 });
+		const pendingBatch = await source.next();
+		expect(pendingBatch.done).toBe(false);
+		if (pendingBatch.done) throw new Error("Expected a pending batch");
+		expect(pendingBatch.value.items.map((item) => item.id)).toEqual([2, 3]);
+		source.close();
+	});
+
 	test("respects concurrency limit", async () => {
 		const concurrent: number[] = [];
 		let maxConcurrent = 0;