Open History to see the last commit
1import * as Scope from "effect/Scope";
2import * as Effect from "effect/Effect";
3import * as TxQueue from "effect/TxQueue";
4import * as TxRef from "effect/TxRef";
5
6export interface KeyedCoalescingWorker<K, V> {
7 readonly enqueue: (key: K, value: V) => Effect.Effect<void>;
8 readonly drainKey: (key: K) => Effect.Effect<void>;
9}
10
11interface KeyedCoalescingWorkerState<K, V> {
12 readonly latestByKey: Map<K, V>;
13 readonly queuedKeys: Set<K>;
14 readonly activeKeys: Set<K>;
15}
16
17export const makeKeyedCoalescingWorker = Effect.fn(
18 "KeyedCoalescingWorker.makeKeyedCoalescingWorker",
19)(function* <K, V, E, R>(options: {
20 readonly merge: (current: V, next: V) => V;
21 readonly process: (key: K, value: V) => Effect.Effect<void, E, R>;
22}): Effect.fn.Return<KeyedCoalescingWorker<K, V>, never, Scope.Scope | R> {
23 const queue = yield* Effect.acquireRelease(TxQueue.unbounded<K>(), TxQueue.shutdown);
24 const stateRef = yield* TxRef.make<KeyedCoalescingWorkerState<K, V>>({
25 latestByKey: new Map(),
26 queuedKeys: new Set(),
27 activeKeys: new Set(),
28 });
29
30 const completeKey = (key: K): Effect.Effect<void> =>
31 TxRef.modify(stateRef, (state) => {
32 const activeKeys = new Set(state.activeKeys);
33 activeKeys.delete(key);
34
35 if (state.latestByKey.has(key) && !state.queuedKeys.has(key)) {
36 const queuedKeys = new Set(state.queuedKeys);
37 queuedKeys.add(key);
38 return [true, { ...state, activeKeys, queuedKeys }] as const;
39 }
40
41 return [false, { ...state, activeKeys }] as const;
42 }).pipe(
43 Effect.flatMap((shouldRequeue) => (shouldRequeue ? TxQueue.offer(queue, key) : Effect.void)),
44 Effect.tx,
45 Effect.asVoid,
46 );
47
48 yield* TxQueue.take(queue).pipe(
49 Effect.flatMap((key) =>
50 TxRef.modify(stateRef, (state) => {
51 const queuedKeys = new Set(state.queuedKeys);
52 queuedKeys.delete(key);
53
54 const value = state.latestByKey.get(key);
55 if (value === undefined) {
56 return [null, { ...state, queuedKeys }] as const;
57 }
58
59 const latestByKey = new Map(state.latestByKey);
60 latestByKey.delete(key);
61 const activeKeys = new Set(state.activeKeys);
62 activeKeys.add(key);
63
64 return [
65 { key, value } as const,
66 { ...state, latestByKey, queuedKeys, activeKeys },
67 ] as const;
68 }).pipe(Effect.tx),
69 ),
70 Effect.flatMap((item) =>
71 item === null
72 ? Effect.void
73 : options
74 .process(item.key, item.value)
75 .pipe(Effect.ignoreCause, Effect.ensuring(completeKey(item.key))),
76 ),
77 Effect.forever,
78 Effect.forkScoped,
79 );
80
81 const enqueue: KeyedCoalescingWorker<K, V>["enqueue"] = (key, value) =>
82 TxRef.modify(stateRef, (state) => {
83 const latestByKey = new Map(state.latestByKey);
84 const existing = latestByKey.get(key);
85 latestByKey.set(key, existing === undefined ? value : options.merge(existing, value));
86
87 if (state.queuedKeys.has(key) || state.activeKeys.has(key)) {
88 return [false, { ...state, latestByKey }] as const;
89 }
90
91 const queuedKeys = new Set(state.queuedKeys);
92 queuedKeys.add(key);
93 return [true, { ...state, latestByKey, queuedKeys }] as const;
94 }).pipe(
95 Effect.flatMap((shouldOffer) => (shouldOffer ? TxQueue.offer(queue, key) : Effect.void)),
96 Effect.tx,
97 Effect.asVoid,
98 );
99
100 const drainKey: KeyedCoalescingWorker<K, V>["drainKey"] = (key) =>
101 TxRef.get(stateRef).pipe(
102 Effect.tap((state) =>
103 state.latestByKey.has(key) || state.queuedKeys.has(key) || state.activeKeys.has(key)
104 ? Effect.txRetry
105 : Effect.void,
106 ),
107 Effect.asVoid,
108 Effect.tx,
109 );
110
111 return { enqueue, drainKey } satisfies KeyedCoalescingWorker<K, V>;
112});