1/**
2 * KeyedCoalescingWorker - A keyed worker that keeps only the latest value per key.
3 *
4 * Enqueues for an active or already-queued key are merged atomically instead of
5 * creating duplicate queued items. Each batch yields to other queued keys before
6 * processing newer updates. `drainKey()` resolves only when that key has no
7 * queued, pending, or active work left.
8 *
9 * @module KeyedCoalescingWorker
10 */
11import * as Scope from "effect/Scope";
12import * as Effect from "effect/Effect";
13import * as TxQueue from "effect/TxQueue";
14import * as TxRef from "effect/TxRef";
15
16export interface KeyedCoalescingWorker<K, V> {
17 readonly enqueue: (key: K, value: V) => Effect.Effect<void>;
18 readonly drainKey: (key: K) => Effect.Effect<void>;
19}
20
21interface KeyedCoalescingWorkerState<K, V> {
22 readonly latestByKey: Map<K, V>;
23 readonly queuedKeys: Set<K>;
24 readonly activeKeys: Set<K>;
25}
26
27export const makeKeyedCoalescingWorker = <K, V, E, R>(options: {
28 readonly merge: (current: V, next: V) => V;
29 readonly process: (key: K, value: V) => Effect.Effect<void, E, R>;
30}): Effect.Effect<KeyedCoalescingWorker<K, V>, never, Scope.Scope | R> =>
31 Effect.gen(function* () {
32 const queue = yield* Effect.acquireRelease(TxQueue.unbounded<K>(), TxQueue.shutdown);
33 const stateRef = yield* TxRef.make<KeyedCoalescingWorkerState<K, V>>({
34 latestByKey: new Map(),
35 queuedKeys: new Set(),
36 activeKeys: new Set(),
37 });
38
39 // Return dirty keys to the tail on success or failure. Keep pending values
40 // available for merging until the next batch starts.
41 const completeKey = (key: K): Effect.Effect<void> =>
42 TxRef.modify(stateRef, (state) => {
43 const activeKeys = new Set(state.activeKeys);
44 activeKeys.delete(key);
45
46 if (state.latestByKey.has(key) && !state.queuedKeys.has(key)) {
47 const queuedKeys = new Set(state.queuedKeys);
48 queuedKeys.add(key);
49 return [true, { ...state, activeKeys, queuedKeys }] as const;
50 }
51
52 return [false, { ...state, activeKeys }] as const;
53 }).pipe(
54 Effect.flatMap((shouldRequeue) =>
55 shouldRequeue ? TxQueue.offer(queue, key) : Effect.void,
56 ),
57 Effect.tx,
58 Effect.asVoid,
59 );
60
61 yield* TxQueue.take(queue).pipe(
62 Effect.flatMap((key) =>
63 TxRef.modify(stateRef, (state) => {
64 const queuedKeys = new Set(state.queuedKeys);
65 queuedKeys.delete(key);
66
67 const value = state.latestByKey.get(key);
68 if (value === undefined) {
69 return [null, { ...state, queuedKeys }] as const;
70 }
71
72 const latestByKey = new Map(state.latestByKey);
73 latestByKey.delete(key);
74 const activeKeys = new Set(state.activeKeys);
75 activeKeys.add(key);
76
77 return [
78 { key, value } as const,
79 { ...state, latestByKey, queuedKeys, activeKeys },
80 ] as const;
81 }).pipe(Effect.tx),
82 ),
83 Effect.flatMap((item) =>
84 item === null
85 ? Effect.void
86 : options
87 .process(item.key, item.value)
88 .pipe(Effect.ignoreCause, Effect.ensuring(completeKey(item.key))),
89 ),
90 Effect.forever,
91 Effect.forkScoped,
92 );
93
94 const enqueue: KeyedCoalescingWorker<K, V>["enqueue"] = (key, value) =>
95 TxRef.modify(stateRef, (state) => {
96 const latestByKey = new Map(state.latestByKey);
97 const existing = latestByKey.get(key);
98 latestByKey.set(key, existing === undefined ? value : options.merge(existing, value));
99
100 if (state.queuedKeys.has(key) || state.activeKeys.has(key)) {
101 return [false, { ...state, latestByKey }] as const;
102 }
103
104 const queuedKeys = new Set(state.queuedKeys);
105 queuedKeys.add(key);
106 return [true, { ...state, latestByKey, queuedKeys }] as const;
107 }).pipe(
108 Effect.flatMap((shouldOffer) => (shouldOffer ? TxQueue.offer(queue, key) : Effect.void)),
109 Effect.tx,
110 Effect.asVoid,
111 );
112
113 const drainKey: KeyedCoalescingWorker<K, V>["drainKey"] = (key) =>
114 TxRef.get(stateRef).pipe(
115 Effect.tap((state) =>
116 state.latestByKey.has(key) || state.queuedKeys.has(key) || state.activeKeys.has(key)
117 ? Effect.txRetry
118 : Effect.void,
119 ),
120 Effect.asVoid,
121 Effect.tx,
122 );
123
124 return { enqueue, drainKey } satisfies KeyedCoalescingWorker<K, V>;
125 });