diff --git a/packages/pgconductor-js/src/lib/async-queue.ts b/packages/pgconductor-js/src/lib/async-queue.ts index bae15d8..83357b4 100644 --- a/packages/pgconductor-js/src/lib/async-queue.ts +++ b/packages/pgconductor-js/src/lib/async-queue.ts @@ -1,11 +1,13 @@ export interface PollableAsyncIterable extends AsyncIterable { tryNext(): T | undefined; + onNextItemAvailable?(listener: () => void): () => void; } export class AsyncQueue implements PollableAsyncIterable { private queue: T[] = []; private resolvers: ((value: IteratorResult) => void)[] = []; private pushResolvers: (() => void)[] = []; + private itemAvailableListeners = new Set<() => void>(); private closed = false; constructor(private readonly capacity: number) {} @@ -30,6 +32,7 @@ export class AsyncQueue implements PollableAsyncIterable { // No consumer waiting, add to queue for later this.queue.push(item); } + this.notifyItemAvailable(); } tryNext(): T | undefined { @@ -51,6 +54,21 @@ export class AsyncQueue implements PollableAsyncIterable { 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) { diff --git a/packages/pgconductor-js/src/lib/batching-async-queue.ts b/packages/pgconductor-js/src/lib/batching-async-queue.ts index 350a9e4..b07893b 100644 --- a/packages/pgconductor-js/src/lib/batching-async-queue.ts +++ b/packages/pgconductor-js/src/lib/batching-async-queue.ts @@ -20,6 +20,7 @@ export class BatchingAsyncQueue implements Polla > { private queue: BatchGroup[] = []; private resolvers: ((value: IteratorResult>) => void)[] = []; + private itemAvailableListeners = new Set<() => void>(); private closed = false; // Track batches per task @@ -117,6 +118,7 @@ export class BatchingAsyncQueue implements Polla // Backpressure is managed at push() level, not here this.queue.push(group); } + this.notifyItemAvailable(); } tryNext(): BatchGroup | undefined { @@ -136,6 +138,21 @@ export class BatchingAsyncQueue 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; diff --git a/packages/pgconductor-js/src/lib/map-concurrent.ts b/packages/pgconductor-js/src/lib/map-concurrent.ts index 4eb0808..0fc59dc 100644 --- a/packages/pgconductor-js/src/lib/map-concurrent.ts +++ b/packages/pgconductor-js/src/lib/map-concurrent.ts @@ -8,14 +8,17 @@ export async function* mapConcurrent( limit: number, mapper: (item: T) => Promise, ): AsyncGenerator { + type Task = { id: number; promise: Promise }; + type RaceResult = { id: number; result: R } | { id: null }; + type ItemWait = { promise: Promise; cancel(): void }; + const it = source[Symbol.asyncIterator](); + const onNextItemAvailable = source.onNextItemAvailable?.bind(source); + const active = new Map(); let sourceDone = false; + let itemWait: ItemWait | null = null; let nextId = 0; - // Track promises with unique IDs - type Task = { id: number; promise: Promise }; - const active = new Map(); - const nextItem = async (): Promise => { if (sourceDone) return null; const { value, done } = await it.next(); @@ -52,30 +55,39 @@ export async function* mapConcurrent( 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((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[] = 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(); } } diff --git a/packages/pgconductor-js/tests/unit/lib/map-concurrent.test.ts b/packages/pgconductor-js/tests/unit/lib/map-concurrent.test.ts index 4e8db55..c96e677 100644 --- a/packages/pgconductor-js/tests/unit/lib/map-concurrent.test.ts +++ b/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 implements PollableAsyncIterable { 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(2); + const firstStarted = new Deferred(); + const secondStarted = new Deferred(); + const releaseFirst = new Deferred(); + + 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(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;