Open History to see the last commit
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 Semaphore from "effect/Semaphore";
11import * as Stream from "effect/Stream";
12
13export class LiveStreamBufferError extends Schema.TaggedError<LiveStreamBufferError>()(
14 "LiveStreamBufferError",
15 { message: Schema.String },
16) {}
17
18export const LIVE_STREAM_MAX_ITEMS = 1_000;
19const LIVE_STREAM_MAX_SERIALIZED_BYTES = 8 * 1024 * 1024;
20
21export interface RetainedLiveItem<A> {
22 readonly value: A;
23 readonly serializedBytes: number;
24}
25
26// Immutable events are shared across subscribers; cache sizes without retaining events or serialized copies.
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
39export const makeLiveStreamBudget = Effect.fn("makeLiveStreamBudget")(function* (limits?: {
40 readonly maxItems?: number;
41 readonly maxSerializedBytes?: number;
42}) {
43 const maxItems = limits?.maxItems ?? LIVE_STREAM_MAX_ITEMS;
44 const maxSerializedBytes = limits?.maxSerializedBytes ?? LIVE_STREAM_MAX_SERIALIZED_BYTES;
45 const failed = yield* Deferred.make<never, LiveStreamBufferError>();
46 const cleanupComplete = yield* Deferred.make<void>();
47 let failure: LiveStreamBufferError | undefined;
48 const retained = new Set<RetainedLiveItem<unknown>>();
49 let retainedSerializedBytes = 0;
50
51 const release = (items: Iterable<RetainedLiveItem<unknown>>) => {
52 for (const item of items) {
53 if (!retained.delete(item)) {
54 continue;
55 }
56 retainedSerializedBytes -= item.serializedBytes;
57 }
58 };
59 yield* Effect.addFinalizer(() =>
60 Effect.sync(() => release(retained)).pipe(
61 Effect.andThen(Deferred.succeed(cleanupComplete, undefined)),
62 ),
63 );
64
65 const check = Effect.suspend(() => (failure ? Effect.fail(failure) : Effect.void));
66
67 const overflow = Effect.fn("LiveStreamBudget.overflow")(function* (
68 nextItems: number,
69 nextSerializedBytes: number,
70 ) {
71 failure ??= new LiveStreamBufferError({
72 message: "The live event buffer is full. Resume from the last received sequence.",
73 });
74 yield* Deferred.fail(failed, failure);
75 yield* Effect.logWarning("orchestration live event buffer is full", {
76 retainedItems: retained.size,
77 retainedSerializedBytes,
78 nextItems,
79 nextSerializedBytes,
80 maxItems,
81 maxSerializedBytes,
82 });
83 return yield* failure;
84 });
85
86 const retain = <A extends object>(value: A, payload: object = value) =>
87 Effect.suspend(() => {
88 if (failure) {
89 return Effect.fail(failure);
90 }
91 const serializedBytes = serializedSize(payload);
92 const nextItems = retained.size + 1;
93 const nextSerializedBytes = retainedSerializedBytes + serializedBytes;
94 if (nextItems > maxItems || nextSerializedBytes > maxSerializedBytes) {
95 return overflow(nextItems, nextSerializedBytes);
96 }
97 const item = { value, serializedBytes };
98 retained.add(item);
99 retainedSerializedBytes = nextSerializedBytes;
100 return Effect.succeed(item);
101 });
102
103 const replace = <A extends object>(
104 previous: ReadonlyArray<RetainedLiveItem<unknown>>,
105 values: ReadonlyArray<A>,
106 payload: (value: A) => object = (value) => value,
107 ) =>
108 Effect.suspend(() => {
109 if (failure) {
110 return Effect.fail(failure);
111 }
112 const next = values.map((value) => ({
113 value,
114 serializedBytes: serializedSize(payload(value)),
115 }));
116 let nextItems = retained.size + next.length;
117 let nextSerializedBytes =
118 retainedSerializedBytes + next.reduce((sum, item) => sum + item.serializedBytes, 0);
119 for (const item of previous) {
120 if (retained.has(item)) {
121 nextItems -= 1;
122 nextSerializedBytes -= item.serializedBytes;
123 }
124 }
125 if (nextItems > maxItems || nextSerializedBytes > maxSerializedBytes) {
126 return overflow(nextItems, nextSerializedBytes);
127 }
128 release(previous);
129 for (const item of next) {
130 retained.add(item);
131 }
132 retainedSerializedBytes = nextSerializedBytes;
133 return Effect.succeed(next);
134 });
135
136 const deliver = <A, E, R>(stream: Stream.Stream<RetainedLiveItem<A>, E, R>) =>
137 Stream.fromPull(
138 Effect.gen(function* () {
139 yield* check;
140 const sourceScope = yield* Scope.fork(yield* Effect.scope);
141 const source = {
142 pull: yield* Stream.toPull(stream).pipe(Scope.provide(sourceScope)),
143 };
144 let inFlight: ReadonlyArray<RetainedLiveItem<A>> = [];
145 yield* Effect.addFinalizer(() =>
146 Effect.sync(() => {
147 release(inFlight);
148 inFlight = [];
149 source.pull = Effect.interrupt;
150 }),
151 );
152 yield* Deferred.await(failed).pipe(
153 Effect.catchTags({
154 LiveStreamBufferError: (error) =>
155 Effect.sync(() => {
156 // Release retained source chunks without waiting for an RPC ACK.
157 source.pull = Effect.interrupt;
158 }).pipe(
159 Effect.andThen(Scope.close(sourceScope, Exit.fail(error))),
160 Effect.andThen(
161 Effect.sync(() => {
162 const delivered = new Set<RetainedLiveItem<unknown>>(inFlight);
163 for (const item of retained) {
164 if (!delivered.has(item)) {
165 release([item]);
166 }
167 }
168 }),
169 ),
170 Effect.andThen(Deferred.succeed(cleanupComplete, undefined)),
171 ),
172 }),
173 Effect.forkScoped,
174 );
175 return Effect.gen(function* () {
176 // RpcServer requests the next batch after ACK; dequeueing does not end delivery.
177 release(inFlight);
178 inFlight = [];
179 yield* check;
180 const items = yield* Effect.raceFirst(source.pull, Deferred.await(failed));
181 inFlight = items;
182 yield* check;
183 return Arr.map(items, (item) => item.value);
184 });
185 }),
186 ).pipe(Stream.scoped);
187
188 return {
189 retain,
190 replace,
191 release,
192 deliver,
193 check,
194 failed: Deferred.await(failed),
195 closed: Deferred.await(cleanupComplete),
196 usage: Effect.sync(() => ({ retainedItems: retained.size, retainedSerializedBytes })),
197 };
198});
199
200type LiveStreamBudget = Effect.Success<ReturnType<typeof makeLiveStreamBudget>>;
201
202type LiveStreamLimits = {
203 readonly maxItems?: number;
204 readonly maxSerializedBytes?: number;
205};
206
207const retainLiveStream = Effect.fn("LiveStreamBudget.retainLiveStream")(function* <
208 A extends object,
209 E,
210 R,
211>(source: Stream.Stream<A, E, R>, budget: LiveStreamBudget) {
212 const queue = yield* Queue.unbounded<
213 RetainedLiveItem<A>,
214 E | LiveStreamBufferError | Cause.Done
215 >();
216 let closed = false;
217 const close = Effect.fn("LiveStreamBudget.close")(function* (error?: LiveStreamBufferError) {
218 if (closed) return;
219 closed = true;
220 budget.release(yield* Queue.clear(queue).pipe(Effect.orElseSucceed(() => [])));
221 if (error) yield* Queue.fail(queue, error);
222 yield* Queue.shutdown(queue);
223 });
224 yield* Effect.addFinalizer(() => close());
225 yield* budget.failed.pipe(Effect.catchTags({ LiveStreamBufferError: close }), Effect.forkScoped);
226 yield* source.pipe(
227 Stream.runForEach((value) =>
228 budget.retain(value).pipe(
229 Effect.flatMap((item) => Queue.offer(queue, item)),
230 Effect.uninterruptible,
231 ),
232 ),
233 Effect.raceFirst(budget.failed),
234 Effect.exit,
235 Effect.flatMap((exit) =>
236 Exit.isFailure(exit) ? Queue.failCause(queue, exit.cause) : Queue.end(queue),
237 ),
238 Effect.forkScoped({ startImmediately: true }),
239 );
240 return Stream.fromQueue(queue);
241});
242
243/** Coalesce queued full states; delivered states stay charged until the next pull ACKs them. */
244export const bufferLatestLiveStream = <A extends { readonly sequence: number }, E, R>(
245 source: Stream.Stream<A, E, R>,
246 key: (value: A) => string,
247 limits?: LiveStreamLimits,
248) =>
249 Stream.unwrap(
250 Effect.gen(function* () {
251 const budget = yield* makeLiveStreamBudget(limits);
252 const ready = yield* Queue.bounded<void, E | LiveStreamBufferError | Cause.Done>(1);
253 const mutex = yield* Semaphore.make(1);
254 const pending = new Map<string, RetainedLiveItem<A>>();
255 let closed = false;
256 const close = (error?: LiveStreamBufferError) =>
257 mutex.withPermit(
258 Effect.gen(function* () {
259 if (closed) return;
260 closed = true;
261 budget.release(pending.values());
262 pending.clear();
263 if (error) yield* Queue.fail(ready, error);
264 yield* Queue.shutdown(ready);
265 }),
266 );
267 yield* Effect.addFinalizer(() => close());
268 yield* budget.failed.pipe(
269 Effect.catchTags({ LiveStreamBufferError: close }),
270 Effect.forkScoped,
271 );
272 yield* source.pipe(
273 Stream.runForEach((value) =>
274 mutex.withPermit(
275 Effect.gen(function* () {
276 yield* budget.check;
277 if (closed) return;
278 const identity = key(value);
279 const previous = pending.get(identity);
280 if (previous && previous.value.sequence >= value.sequence) return;
281 const wasEmpty = pending.size === 0;
282 const [item] = yield* budget.replace(previous ? [previous] : [], [value]);
283 pending.set(identity, item!);
284 // Only the empty-to-nonempty transition needs a wakeup; replacements
285 // must not leave another token after this batch has drained.
286 if (wasEmpty) yield* Queue.offer(ready, undefined);
287 }).pipe(Effect.uninterruptible),
288 ),
289 ),
290 Effect.raceFirst(budget.failed),
291 Effect.exit,
292 Effect.flatMap((exit) =>
293 Exit.isFailure(exit) ? Queue.failCause(ready, exit.cause) : Queue.end(ready),
294 ),
295 Effect.forkScoped({ startImmediately: true }),
296 );
297 return budget.deliver(
298 Stream.fromPull(
299 Effect.succeed(
300 Queue.take(ready).pipe(
301 Effect.andThen(
302 mutex.withPermit(
303 Effect.sync(() => {
304 const items = Array.from(pending.values()).sort(
305 (left, right) => left.value.sequence - right.value.sequence,
306 );
307 pending.clear();
308 // Moving to delivery preserves the budget charge until ACK.
309 return items as Arr.NonEmptyArray<RetainedLiveItem<A>>;
310 }),
311 ),
312 ),
313 ),
314 ),
315 ),
316 );
317 }),
318 );
319
320/**
321 * Subscribe and start draining before reading the high-water mark. A blocked
322 * catch-up query or an unacknowledged replay batch must not strand an unbounded
323 * PubSub subscription.
324 */
325type ReplayLiveInput<A extends { readonly sequence: number }, E, R> = {
326 readonly subscribe: Effect.Effect<PubSub.Subscription<A>, never, Scope.Scope>;
327 readonly latestSequence: Effect.Effect<number, E, R>;
328 readonly replay: (throughSequence: number) => Stream.Stream<A, E, R>;
329 readonly afterSequence?: number;
330 readonly filter?: (event: A) => boolean;
331};
332
333export const replayAndBufferProjectedLiveEvents = <
334 A extends { readonly sequence: number },
335 B extends { readonly sequence: number },
336 E,
337 R,
338>(
339 input: ReplayLiveInput<A, E, R> & { readonly project: (event: A) => B },
340 limits?: LiveStreamLimits,
341) =>
342 Stream.unwrap(
343 Effect.gen(function* () {
344 const budget = yield* makeLiveStreamBudget(limits);
345 const subscriptionScope = yield* Scope.fork(yield* Effect.scope);
346 const subscription = yield* input.subscribe.pipe(Scope.provide(subscriptionScope));
347 let throughSequence = input.afterSequence ?? 0;
348 const live = yield* retainLiveStream(
349 // Take one event at a time so no uncharged PubSub batch is retained.
350 Stream.fromEffectRepeat(PubSub.take(subscription)).pipe(
351 Stream.filter(
352 (event) => event.sequence > throughSequence && (input.filter?.(event) ?? true),
353 ),
354 Stream.map(input.project),
355 Stream.ensuring(Scope.close(subscriptionScope, Exit.void)),
356 ),
357 budget,
358 );
359 const replay = Stream.unwrap(
360 input.latestSequence.pipe(
361 Effect.map((highWater) => {
362 throughSequence = Math.max(highWater, input.afterSequence ?? 0);
363 return input.replay(highWater).pipe(
364 // The finite database page is independent of the live tail.
365 // Charge only the item being delivered so a fast reader can
366 // consume a page larger than the budget without overflowing.
367 Stream.filter((event) => input.filter?.(event) ?? true),
368 Stream.rechunk(1),
369 Stream.mapEffect((event) => budget.retain(input.project(event))),
370 );
371 }),
372 ),
373 );
374 return budget.deliver(
375 Stream.concat(
376 replay,
377 live.pipe(
378 Stream.filter((item) => {
379 if (item.value.sequence > throughSequence) return true;
380 budget.release([item]);
381 return false;
382 }),
383 ),
384 ),
385 );
386 }),
387 );