1import type { NonEmptyReadonlyArray } from "effect/Array";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Ref from "effect/Ref";
5import * as Schedule from "effect/Schedule";
6import * as Sink from "effect/Sink";
7
8/**
9 * Raised when a batch API keeps reporting entries as unprocessed after the
10 * bounded retry schedule is exhausted. Carries the stranded entries (in input
11 * order) so callers can decide what to do with them (DLQ, die, drop).
12 */
13export class BatchRetryExhaustedError<In> extends Data.TaggedError(
14 "BatchRetryExhaustedError",
15)<{
16 /** The entries that were still unprocessed when retries ran out, in input order. */
17 readonly entries: readonly In[];
18}> {}
19
20export interface BatchedSinkOptions<In, Out, Err> {
21 /**
22 * API max records per call (10 for SQS `SendMessageBatch` / SNS
23 * `PublishBatch`, 25 for `BatchWriteItem`, 500 for Kinesis `PutRecords`,
24 * 500 for Firehose, 10 for EventBridge, ...).
25 */
26 readonly maxRecords: number;
27 /** API max payload bytes per call; entries are greedily packed. */
28 readonly maxBytes?: number;
29 /** Approximate wire size of one record; required for `maxBytes` packing. */
30 readonly sizeOf?: (record: In) => number;
31 /** One API call per packed batch. Receives entries in input order. */
32 readonly send: (batch: readonly In[]) => Effect.Effect<Out, Err>;
33 /**
34 * Extract entries the API reports as *transiently* unprocessed (returned in
35 * input order). These are re-submitted on the bounded `retrySchedule`.
36 */
37 readonly unprocessed?: (out: Out, batch: readonly In[]) => readonly In[];
38 /**
39 * Extract entries the API reports as *permanently* rejected (returned in
40 * input order). Distinct from `unprocessed`: rejected entries are never
41 * retried — they are dropped and surfaced via `onRejected`.
42 */
43 readonly rejected?: (out: Out, batch: readonly In[]) => readonly In[];
44 /**
45 * Invoked with permanently-rejected entries before they are dropped.
46 * @default logs a warning with the rejected entry count
47 */
48 readonly onRejected?: (entries: readonly In[]) => Effect.Effect<void>;
49 /**
50 * Bounded retry for the unprocessed subset.
51 * @default Schedule.recurs(5) ∩ Schedule.exponential("200 millis")
52 */
53 readonly retrySchedule?: Schedule.Schedule<unknown>;
54}
55
56/**
57 * The shared engine behind every AWS batch-API sink (`SQS.QueueSink`,
58 * `SNS.TopicSink`, `Kinesis.StreamSink`, ...).
59 *
60 * Semantics:
61 * - each upstream chunk is split into `<= maxRecords` / `<= maxBytes` batches
62 * **preserving input order**;
63 * - batches are sent **sequentially** (order-preservation is the default);
64 * - after each `send`, the `rejected` subset is surfaced via `onRejected` and
65 * dropped, and the `unprocessed` subset is re-submitted (in order) on the
66 * bounded `retrySchedule`;
67 * - exhausting retries fails the sink with a typed
68 * {@link BatchRetryExhaustedError} carrying the stranded entries.
69 *
70 * Internal — NOT exported from the `AWS` barrel.
71 */
72export const makeBatchedSink = <In, Out, Err>(
73 options: BatchedSinkOptions<In, Out, Err>,
74): Sink.Sink<void, In, never, Err | BatchRetryExhaustedError<In>> => {
75 const schedule =
76 options.retrySchedule ??
77 Schedule.max([Schedule.recurs(5), Schedule.exponential("200 millis")]);
78
79 const onRejected =
80 options.onRejected ??
81 ((entries: readonly In[]) =>
82 Effect.logWarning(
83 `BatchedSink dropped ${entries.length} permanently rejected entr${
84 entries.length === 1 ? "y" : "ies"
85 }`,
86 ));
87
88 // One attempt over the currently-pending subset. Returns the entries that
89 // are still unprocessed after this attempt (empty = converged).
90 const attempt = (
91 pendingRef: Ref.Ref<readonly In[]>,
92 ): Effect.Effect<readonly In[], Err> =>
93 Effect.gen(function* () {
94 const pending = yield* Ref.get(pendingRef);
95 if (pending.length === 0) {
96 return pending;
97 }
98 const out = yield* options.send(pending);
99 const rejected = options.rejected?.(out, pending) ?? [];
100 if (rejected.length > 0) {
101 yield* onRejected(rejected);
102 }
103 const remaining = options.unprocessed?.(out, pending) ?? [];
104 yield* Ref.set(pendingRef, remaining);
105 return remaining;
106 });
107
108 const flushBatch = (
109 batch: readonly In[],
110 ): Effect.Effect<void, Err | BatchRetryExhaustedError<In>> =>
111 Effect.gen(function* () {
112 const pendingRef = yield* Ref.make(batch);
113 yield* repeatWhilePending(attempt(pendingRef), schedule);
114 const stranded = yield* Ref.get(pendingRef);
115 if (stranded.length > 0) {
116 return yield* new BatchRetryExhaustedError({ entries: stranded });
117 }
118 });
119
120 return Sink.forEachArray((chunk: NonEmptyReadonlyArray<In>) =>
121 Effect.gen(function* () {
122 const batches = yield* Effect.sync(() => packBatches(chunk, options));
123 yield* Effect.forEach(batches, flushBatch, { discard: true });
124 }),
125 );
126};
127
128/**
129 * Re-runs `step` (which re-submits the pending subset) on the bounded
130 * schedule until it reports nothing left to send.
131 *
132 * Extracted with an explicit return type so `Repeat.Return`'s conditional
133 * type never leaks into declaration emit (see PATTERNS §7 on inlined
134 * retry/repeat poisoning `.d.ts` inference).
135 */
136const repeatWhilePending = <In, Err>(
137 step: Effect.Effect<readonly In[], Err>,
138 schedule: Schedule.Schedule<unknown>,
139): Effect.Effect<readonly In[], Err> =>
140 Effect.repeat(step, {
141 schedule,
142 until: (remaining: readonly In[]) => remaining.length === 0,
143 });
144
145/** Greedily pack a chunk into `<= maxRecords` / `<= maxBytes` batches, preserving order. */
146const packBatches = <In>(
147 chunk: NonEmptyReadonlyArray<In>,
148 options: {
149 readonly maxRecords: number;
150 readonly maxBytes?: number;
151 readonly sizeOf?: (record: In) => number;
152 },
153): readonly (readonly In[])[] => {
154 const { maxBytes, maxRecords, sizeOf } = options;
155 const batches: In[][] = [];
156 let current: In[] = [];
157 let currentBytes = 0;
158 for (const record of chunk) {
159 const size = sizeOf?.(record) ?? 0;
160 const overRecords = current.length >= maxRecords;
161 const overBytes =
162 maxBytes !== undefined &&
163 current.length > 0 &&
164 currentBytes + size > maxBytes;
165 if (overRecords || overBytes) {
166 batches.push(current);
167 current = [];
168 currentBytes = 0;
169 }
170 // A single record larger than maxBytes still ships alone; the API
171 // rejects it and the error surfaces through the sink's error channel.
172 current.push(record);
173 currentBytes += size;
174 }
175 if (current.length > 0) {
176 batches.push(current);
177 }
178 return batches;
179};