diff --git a/packages/shared/src/KeyedCoalescingWorker.test.ts b/packages/shared/src/KeyedCoalescingWorker.test.ts index 8bfc1a340a60..8ce44069d2d5 100644 --- a/packages/shared/src/KeyedCoalescingWorker.test.ts +++ b/packages/shared/src/KeyedCoalescingWorker.test.ts @@ -6,6 +6,98 @@ import * as Effect from "effect/Effect"; import { makeKeyedCoalescingWorker } from "./KeyedCoalescingWorker.ts"; describe("makeKeyedCoalescingWorker", () => { + it.live("lets a queued key drain before a continuously updated key finishes", () => + Effect.scoped( + Effect.gen(function* () { + const processed: string[] = []; + const busyStarted = yield* Deferred.make(); + const releaseBusy = yield* Deferred.make(); + + const worker = yield* makeKeyedCoalescingWorker({ + merge: (_current, next) => next, + process: (key, value): Effect.Effect => + Effect.gen(function* () { + processed.push(`${key}:${value}`); + if (key !== "busy") return; + + if (value === 1) { + yield* Deferred.succeed(busyStarted, undefined); + yield* Deferred.await(releaseBusy); + } + if (value < 4) yield* worker.enqueue(key, value + 1); + }), + }); + + yield* worker.enqueue("busy", 1); + yield* Deferred.await(busyStarted); + yield* worker.enqueue("quiet", 1); + yield* Deferred.succeed(releaseBusy, undefined); + yield* worker.drainKey("quiet"); + yield* worker.drainKey("busy"); + + expect(processed).toEqual(["busy:1", "quiet:1", "busy:2", "busy:3", "busy:4"]); + }), + ), + ); + + it.live("merges updates while waiting behind another key and drains the merged batch", () => + Effect.scoped( + Effect.gen(function* () { + const processed: string[] = []; + const busyStarted = yield* Deferred.make(); + const releaseBusy = yield* Deferred.make(); + const quietStarted = yield* Deferred.make(); + const releaseQuiet = yield* Deferred.make(); + const mergedStarted = yield* Deferred.make(); + const releaseMerged = yield* Deferred.make(); + + const worker = yield* makeKeyedCoalescingWorker({ + merge: (current, next) => `${current}+${next}`, + process: (key, value) => + Effect.gen(function* () { + processed.push(`${key}:${value}`); + if (value === "first") { + yield* Deferred.succeed(busyStarted, undefined); + yield* Deferred.await(releaseBusy); + } else if (key === "quiet") { + yield* Deferred.succeed(quietStarted, undefined); + yield* Deferred.await(releaseQuiet); + } else if (value === "a+b+c") { + yield* Deferred.succeed(mergedStarted, undefined); + yield* Deferred.await(releaseMerged); + } + }), + }); + + yield* worker.enqueue("busy", "first"); + yield* Deferred.await(busyStarted); + yield* worker.enqueue("quiet", "quiet"); + yield* worker.enqueue("busy", "a"); + yield* Deferred.succeed(releaseBusy, undefined); + yield* Deferred.await(quietStarted); + expect(processed).toEqual(["busy:first", "quiet:quiet"]); + yield* worker.enqueue("busy", "b"); + yield* worker.enqueue("busy", "c"); + + const drained = yield* Deferred.make(); + yield* Effect.forkChild( + worker.drainKey("busy").pipe(Effect.andThen(Deferred.succeed(drained, undefined))), + { startImmediately: true }, + ); + expect(yield* Deferred.isDone(drained)).toBe(false); + + yield* Deferred.succeed(releaseQuiet, undefined); + yield* Deferred.await(mergedStarted); + expect(yield* Deferred.isDone(drained)).toBe(false); + yield* Deferred.succeed(releaseMerged, undefined); + yield* Deferred.await(drained); + yield* worker.drainKey("quiet"); + + expect(processed).toEqual(["busy:first", "quiet:quiet", "busy:a+b+c"]); + }), + ), + ); + it.live("waits for latest work enqueued during active processing before draining the key", () => Effect.scoped( Effect.gen(function* () { @@ -85,12 +177,13 @@ describe("makeKeyedCoalescingWorker", () => { yield* worker.enqueue("terminal-1", "first"); yield* Deferred.await(firstStarted); + yield* worker.enqueue("terminal-2", "quiet"); yield* worker.enqueue("terminal-1", "second"); yield* Deferred.succeed(releaseFailure, undefined); yield* Deferred.await(secondProcessed); yield* worker.drainKey("terminal-1"); - expect(processed).toEqual(["terminal-1:first", "terminal-1:second"]); + expect(processed).toEqual(["terminal-1:first", "terminal-2:quiet", "terminal-1:second"]); }), ), ); diff --git a/packages/shared/src/KeyedCoalescingWorker.ts b/packages/shared/src/KeyedCoalescingWorker.ts index f4edebfafb30..ce1d89a41169 100644 --- a/packages/shared/src/KeyedCoalescingWorker.ts +++ b/packages/shared/src/KeyedCoalescingWorker.ts @@ -2,8 +2,9 @@ * KeyedCoalescingWorker - A keyed worker that keeps only the latest value per key. * * Enqueues for an active or already-queued key are merged atomically instead of - * creating duplicate queued items. `drainKey()` resolves only when that key has - * no queued, pending, or active work left. + * creating duplicate queued items. Each batch yields to other queued keys before + * processing newer updates. `drainKey()` resolves only when that key has no + * queued, pending, or active work left. * * @module KeyedCoalescingWorker */ @@ -35,28 +36,9 @@ export const makeKeyedCoalescingWorker = (options: { activeKeys: new Set(), }); - const processKey = (key: K, value: V): Effect.Effect => - options.process(key, value).pipe( - Effect.flatMap(() => - TxRef.modify(stateRef, (state) => { - const nextValue = state.latestByKey.get(key); - if (nextValue === undefined) { - const activeKeys = new Set(state.activeKeys); - activeKeys.delete(key); - return [null, { ...state, activeKeys }] as const; - } - - const latestByKey = new Map(state.latestByKey); - latestByKey.delete(key); - return [nextValue, { ...state, latestByKey }] as const; - }).pipe(Effect.tx), - ), - Effect.flatMap((nextValue) => - nextValue === null ? Effect.void : processKey(key, nextValue), - ), - ); - - const cleanupFailedKey = (key: K): Effect.Effect => + // Return dirty keys to the tail on success or failure. Keep pending values + // available for merging until the next batch starts. + const completeKey = (key: K): Effect.Effect => TxRef.modify(stateRef, (state) => { const activeKeys = new Set(state.activeKeys); activeKeys.delete(key); @@ -69,10 +51,11 @@ export const makeKeyedCoalescingWorker = (options: { return [false, { ...state, activeKeys }] as const; }).pipe( - Effect.tx, Effect.flatMap((shouldRequeue) => shouldRequeue ? TxQueue.offer(queue, key) : Effect.void, ), + Effect.tx, + Effect.asVoid, ); yield* TxQueue.take(queue).pipe( @@ -100,9 +83,9 @@ export const makeKeyedCoalescingWorker = (options: { Effect.flatMap((item) => item === null ? Effect.void - : processKey(item.key, item.value).pipe( - Effect.catchCause(() => cleanupFailedKey(item.key)), - ), + : options + .process(item.key, item.value) + .pipe(Effect.ignoreCause, Effect.ensuring(completeKey(item.key))), ), Effect.forever, Effect.forkScoped,