1import * as Schema from "effect/Schema";
2import * as Arr from "effect/Array";
3import type * as Cause from "effect/Cause";
4import * as Deferred from "effect/Deferred";
5import * as Effect from "effect/Effect";
6import * as Exit from "effect/Exit";
7import * as Queue from "effect/Queue";
8import * as PubSub from "effect/PubSub";
9import * as Scope from "effect/Scope";
10import * as Stream from "effect/Stream";
11
12export class LiveStreamBufferError extends Schema.TaggedError<LiveStreamBufferError>()(
13 "LiveStreamBufferError",
14 { message: Schema.String },
15) {}
16
17export const LIVE_STREAM_MAX_ITEMS = 1_000;
18const LIVE_STREAM_MAX_SERIALIZED_BYTES = 8 * 1024 * 1024;
19
20export interface RetainedLiveItem<A> {
21 readonly value: A;
22 readonly serializedBytes: number;
23}
24
25// Published events are immutable and shared across subscriptions. Measure each
26// object once without keeping the event or its serialized copy alive.
27const serializedSizes = new WeakMap<object, number>();
28
29function serializedSize(value: object): number {
30 const cached = serializedSizes.get(value);
31 if (cached !== undefined) {
32 return cached;
33 }
34 const bytes = Buffer.byteLength(JSON.stringify(value));
35 serializedSizes.set(value, bytes);
36 return bytes;
37}
38
39/** One budget covers one subscription and its delivery stream, including the batch waiting for an RPC ACK. */
40export const makeLiveStreamBudget = Effect.fn("makeLiveStreamBudget")(function* (limits?: {
41 readonly maxItems?: number;
42 readonly maxSerializedBytes?: number;
43}) {
44 const maxItems = limits?.maxItems ?? LIVE_STREAM_MAX_ITEMS;
45 const maxSerializedBytes = limits?.maxSerializedBytes ?? LIVE_STREAM_MAX_SERIALIZED_BYTES;
46 const failed = yield* Deferred.make<never, LiveStreamBufferError>();
47 const cleanupComplete = yield* Deferred.make<void>();
48 let failure: LiveStreamBufferError | undefined;
49 const retained = new Set<RetainedLiveItem<unknown>>();
50 let retainedSerializedBytes = 0;
51
52 const release = (items: Iterable<RetainedLiveItem<unknown>>) => {
53 for (const item of items) {
54 if (!retained.delete(item)) {
55 continue;
56 }
57 retainedSerializedBytes -= item.serializedBytes;
58 }
59 };
60 yield* Effect.addFinalizer(() =>
61 Effect.sync(() => release(retained)).pipe(
62 Effect.andThen(Deferred.succeed(cleanupComplete, undefined)),
63 ),
64 );
65
66 const check = Effect.suspend(() => (failure ? Effect.fail(failure) : Effect.void));
67
68 const overflow = Effect.fn("LiveStreamBudget.overflow")(function* (
69 nextItems: number,
70 nextSerializedBytes: number,
71 ) {
72 failure ??= new LiveStreamBufferError({
73 message: "The live event buffer is full. Resume from the last received sequence.",
74 });
75 yield* Deferred.fail(failed, failure);
76 yield* Effect.logWarning("orchestration live event buffer is full", {
77 retainedItems: retained.size,
78 retainedSerializedBytes,
79 nextItems,
80 nextSerializedBytes,
81 maxItems,
82 maxSerializedBytes,
83 });
84 return yield* failure;
85 });
86
87 const retain = <A extends object>(value: A, payload: object = value) =>
88 Effect.suspend(() => {
89 if (failure) {
90 return Effect.fail(failure);
91 }
92 const serializedBytes = serializedSize(payload);
93 const nextItems = retained.size + 1;
94 const nextSerializedBytes = retainedSerializedBytes + serializedBytes;
95 if (nextItems > maxItems || nextSerializedBytes > maxSerializedBytes) {
96 return overflow(nextItems, nextSerializedBytes);
97 }
98 const item = { value, serializedBytes };
99 retained.add(item);
100 retainedSerializedBytes = nextSerializedBytes;
101 return Effect.succeed(item);
102 });
103
104 // Replace one coalescing batch atomically. Queued and coalesced payloads
105 // count against the same budget, and discarded updates release their charge.
106 const replace = <A extends object>(
107 previous: ReadonlyArray<RetainedLiveItem<unknown>>,
108 values: ReadonlyArray<A>,
109 payload: (value: A) => object = (value) => value,
110 ) =>
111 Effect.suspend(() => {
112 if (failure) {
113 return Effect.fail(failure);
114 }
115 const next = values.map((value) => ({
116 value,
117 serializedBytes: serializedSize(payload(value)),
118 }));
119 let nextItems = retained.size + next.length;
120 let nextSerializedBytes =
121 retainedSerializedBytes + next.reduce((sum, item) => sum + item.serializedBytes, 0);
122 for (const item of previous) {
123 if (retained.has(item)) {
124 nextItems -= 1;
125 nextSerializedBytes -= item.serializedBytes;
126 }
127 }
128 if (nextItems > maxItems || nextSerializedBytes > maxSerializedBytes) {
129 return overflow(nextItems, nextSerializedBytes);
130 }
131 release(previous);
132 for (const item of next) {
133 retained.add(item);
134 }
135 retainedSerializedBytes = nextSerializedBytes;
136 return Effect.succeed(next);
137 });
138
139 const deliver = <A, E, R>(stream: Stream.Stream<RetainedLiveItem<A>, E, R>) =>
140 Stream.fromPull(
141 Effect.gen(function* () {
142 yield* check;
143 const sourceScope = yield* Scope.fork(yield* Effect.scope);
144 const source = {
145 pull: yield* Stream.toPull(stream).pipe(Scope.provide(sourceScope)),
146 };
147 let inFlight: ReadonlyArray<RetainedLiveItem<A>> = [];
148 yield* Effect.addFinalizer(() =>
149 Effect.sync(() => {
150 release(inFlight);
151 inFlight = [];
152 source.pull = Effect.interrupt;
153 }),
154 );
155 yield* Deferred.await(failed).pipe(
156 Effect.catchTags({
157 LiveStreamBufferError: (error) =>
158 Effect.sync(() => {
159 // A grouped source can retain its own pending chunk. Close it
160 // and release the pull closure without waiting for an RPC ACK.
161 source.pull = Effect.interrupt;
162 }).pipe(
163 Effect.andThen(Scope.close(sourceScope, Exit.fail(error))),
164 Effect.andThen(
165 Effect.sync(() => {
166 const delivered = new Set<RetainedLiveItem<unknown>>(inFlight);
167 for (const item of retained) {
168 if (!delivered.has(item)) {
169 release([item]);
170 }
171 }
172 }),
173 ),
174 Effect.andThen(Deferred.succeed(cleanupComplete, undefined)),
175 ),
176 }),
177 Effect.forkScoped,
178 );
179 // @effect-diagnostics-next-line returnEffectInGen:off - Stream.fromPull needs the pull effect as its result.
180 return Effect.gen(function* () {
181 // RpcServer requests the next batch only after the client ACKs this
182 // one. Removing items from a queue alone does not mean delivery ended.
183 release(inFlight);
184 inFlight = [];
185 yield* check;
186 const items = yield* Effect.raceFirst(source.pull, Deferred.await(failed));
187 inFlight = items;
188 yield* check;
189 return Arr.map(items, (item) => item.value);
190 });
191 }),
192 ).pipe(Stream.scoped);
193
194 return {
195 retain,
196 replace,
197 release,
198 deliver,
199 check,
200 failed: Deferred.await(failed),
201 closed: Deferred.await(cleanupComplete),
202 usage: Effect.sync(() => ({ retainedItems: retained.size, retainedSerializedBytes })),
203 };
204});
205
206type LiveStreamBudget = Effect.Success<ReturnType<typeof makeLiveStreamBudget>>;
207
208type LiveStreamLimits = {
209 readonly maxItems?: number;
210 readonly maxSerializedBytes?: number;
211};
212
213const retainLiveStream = <A extends object, E, R>(
214 source: Stream.Stream<A, E, R>,
215 budget: LiveStreamBudget,
216) =>
217 Effect.gen(function* () {
218 const queue = yield* Queue.unbounded<
219 RetainedLiveItem<A>,
220 E | LiveStreamBufferError | Cause.Done
221 >();
222 let closed = false;
223 const close = (error?: LiveStreamBufferError) =>
224 Effect.gen(function* () {
225 if (closed) return;
226 closed = true;
227 budget.release(yield* Queue.clear(queue).pipe(Effect.orElseSucceed(() => [])));
228 if (error) yield* Queue.fail(queue, error);
229 yield* Queue.shutdown(queue);
230 });
231 yield* Effect.addFinalizer(() => close());
232 yield* budget.failed.pipe(
233 Effect.catchTags({ LiveStreamBufferError: close }),
234 Effect.forkScoped,
235 );
236 yield* source.pipe(
237 Stream.runForEach((value) =>
238 budget.retain(value).pipe(
239 Effect.flatMap((item) => Queue.offer(queue, item)),
240 Effect.uninterruptible,
241 ),
242 ),
243 Effect.raceFirst(budget.failed),
244 Effect.exit,
245 Effect.flatMap((exit) =>
246 Exit.isFailure(exit) ? Queue.failCause(queue, exit.cause) : Queue.end(queue),
247 ),
248 Effect.forkScoped({ startImmediately: true }),
249 );
250 return Stream.fromQueue(queue);
251 });
252
253/** Drain live events while RPC delivery waits for an ACK, with a bounded retained tail. */
254export const bufferLiveStream = <A extends object, E, R>(
255 source: Stream.Stream<A, E, R>,
256 limits?: LiveStreamLimits,
257) =>
258 Stream.unwrap(
259 Effect.gen(function* () {
260 const budget = yield* makeLiveStreamBudget(limits);
261 return budget.deliver(yield* retainLiveStream(source, budget));
262 }),
263 );
264
265/**
266 * Subscribe and start draining before reading the high-water mark. A blocked
267 * catch-up query or an unacknowledged replay batch must not strand an unbounded
268 * PubSub subscription. Raw replay/live events share one budget here; projected
269 * RPC delivery has its own budget, so the two stages retain at most twice the
270 * configured item and byte limits.
271 */
272type ReplayLiveInput<A extends { readonly sequence: number }, E, R> = {
273 readonly subscribe: Effect.Effect<PubSub.Subscription<A>, never, Scope.Scope>;
274 readonly latestSequence: Effect.Effect<number, E, R>;
275 readonly replay: (throughSequence: number) => Stream.Stream<A, E, R>;
276 readonly afterSequence?: number;
277 readonly filter?: (event: A) => boolean;
278};
279
280export const replayAndBufferLiveEvents = <A extends { readonly sequence: number }, E, R>(
281 input: ReplayLiveInput<A, E, R>,
282 limits?: LiveStreamLimits,
283) => replayAndBufferProjectedLiveEvents({ ...input, project: (event: A) => event }, limits);
284
285/** Reduce transport payloads before either replay or live retention charges the byte budget. */
286export const replayAndBufferProjectedLiveEvents = <
287 A extends { readonly sequence: number },
288 B extends { readonly sequence: number },
289 E,
290 R,
291>(
292 input: ReplayLiveInput<A, E, R> & { readonly project: (event: A) => B },
293 limits?: LiveStreamLimits,
294) =>
295 Stream.unwrap(
296 Effect.gen(function* () {
297 const budget = yield* makeLiveStreamBudget(limits);
298 const subscriptionScope = yield* Scope.fork(yield* Effect.scope);
299 const subscription = yield* input.subscribe.pipe(Scope.provide(subscriptionScope));
300 let throughSequence = input.afterSequence ?? 0;
301 const live = yield* retainLiveStream(
302 // Taking one event at a time avoids retaining an uncharged PubSub batch
303 // while the drain charges events to the budget.
304 Stream.fromEffectRepeat(PubSub.take(subscription)).pipe(
305 Stream.filter(
306 (event) => event.sequence > throughSequence && (input.filter?.(event) ?? true),
307 ),
308 Stream.map(input.project),
309 Stream.ensuring(Scope.close(subscriptionScope, Exit.void)),
310 ),
311 budget,
312 );
313 const replay = Stream.unwrap(
314 input.latestSequence.pipe(
315 Effect.map((highWater) => {
316 throughSequence = Math.max(highWater, input.afterSequence ?? 0);
317 return input.replay(highWater).pipe(
318 // The finite database page is independent of the live tail.
319 // Charge only the item being delivered so a fast reader can
320 // consume a page larger than the budget without overflowing.
321 Stream.rechunk(1),
322 Stream.mapEffect((event) => budget.retain(input.project(event))),
323 );
324 }),
325 ),
326 );
327 return budget.deliver(
328 Stream.concat(
329 replay,
330 live.pipe(
331 Stream.filter((item) => {
332 if (item.value.sequence > throughSequence) return true;
333 // An event published between subscribe and the high-water read
334 // is already present in replay.
335 budget.release([item]);
336 return false;
337 }),
338 ),
339 ),
340 );
341 }),
342 );