1import * as Cause from "effect/Cause";
2import * as Scope from "effect/Scope";
3import * as Effect from "effect/Effect";
4import * as TxQueue from "effect/TxQueue";
5import * as TxRef from "effect/TxRef";
6
7export interface DrainableWorker<A> {
8 readonly enqueue: (item: A) => Effect.Effect<void>;
9
10 readonly drain: Effect.Effect<void>;
11}
12
13/**
14 * A finalizer shuts down the queue and drops queued items, so `drain` resolves after the scope
15 * closes instead of waiting on them. An item that fails or dies is logged and skipped; the worker
16 * keeps processing later items and `drain` still resolves.
17 */
18export const makeDrainableWorker = Effect.fn("DrainableWorker.makeDrainableWorker")(function* <
19 A,
20 E,
21 R,
22>(
23 process: (item: A) => Effect.Effect<void, E, R>,
24): Effect.fn.Return<DrainableWorker<A>, never, Scope.Scope | R> {
25 const outstanding = yield* TxRef.make(0);
26 const queue = yield* Effect.acquireRelease(TxQueue.unbounded<A>(), (queue) =>
27 // Uncount only the dropped items: an item still running uncounts itself,
28 // even when a parallel scope closes it after this finalizer.
29 TxQueue.clear(queue).pipe(
30 Effect.flatMap((dropped) => TxRef.update(outstanding, (n) => n - dropped.length)),
31 Effect.andThen(TxQueue.shutdown(queue)),
32 Effect.tx,
33 ),
34 );
35
36 yield* TxQueue.take(queue).pipe(
37 Effect.flatMap((a) =>
38 // `suspend` turns a `process` that throws while building its effect
39 // into this item's defect instead of the loop's.
40 Effect.suspend(() => process(a)).pipe(
41 // Only the item's own failure, defect, or interruption lands here and
42 // the loop continues; interrupting the worker fiber still stops it.
43 Effect.catchCause((cause) =>
44 Cause.hasInterruptsOnly(cause)
45 ? Effect.void
46 : Effect.logError("DrainableWorker item failed", cause),
47 ),
48 Effect.ensuring(TxRef.update(outstanding, (n) => n - 1)),
49 ),
50 ),
51 Effect.forever,
52 Effect.forkScoped,
53 );
54
55 const drain: DrainableWorker<A>["drain"] = TxRef.get(outstanding).pipe(
56 Effect.flatMap((n) => (n > 0 ? Effect.txRetry : Effect.void)),
57 Effect.tx,
58 );
59
60 const enqueue = (element: A): Effect.Effect<boolean, never, never> =>
61 TxQueue.offer(queue, element).pipe(
62 Effect.tap((offered) => (offered ? TxRef.update(outstanding, (n) => n + 1) : Effect.void)),
63 Effect.tx,
64 );
65
66 return { enqueue, drain } satisfies DrainableWorker<A>;
67});