1import { it } from "@effect/vitest";
2import * as Deferred from "effect/Deferred";
3import * as Effect from "effect/Effect";
4import * as Fiber from "effect/Fiber";
5import * as PubSub from "effect/PubSub";
6import * as Queue from "effect/Queue";
7import * as Stream from "effect/Stream";
8import { describe, expect } from "vite-plus/test";
9
10import {
11 bufferLiveStream,
12 makeLiveStreamBudget,
13 replayAndBufferLiveEvents,
14 type RetainedLiveItem,
15} from "./LiveStreamBudget.ts";
16
17describe("LiveStreamBudget", () => {
18 it.effect("closes the source without releasing a batch still waiting for an ACK", () =>
19 Effect.scoped(
20 Effect.gen(function* () {
21 const budget = yield* makeLiveStreamBudget({ maxItems: 3 });
22 const queue = yield* Queue.unbounded<RetainedLiveItem<{ text: string }>>();
23 const sourceClosed = yield* Deferred.make<void>();
24 const first = yield* budget.retain({ text: "first" });
25 const second = yield* budget.retain({ text: "second" });
26 const third = yield* budget.retain({ text: "third" });
27 yield* Queue.offerAll(queue, [first, second, third]);
28 yield* Effect.scoped(
29 Effect.gen(function* () {
30 const pull = yield* Stream.toPull(
31 budget.deliver(
32 Stream.fromQueue(queue).pipe(
33 Stream.rechunk(1),
34 Stream.ensuring(Deferred.succeed(sourceClosed, undefined)),
35 ),
36 ),
37 );
38 expect(yield* pull).toEqual([{ text: "first" }]);
39 // The other two items are in the source's pull state, not its queue.
40 expect(yield* Queue.size(queue)).toBe(0);
41 expect((yield* budget.usage).retainedItems).toBe(3);
42 const overflow = yield* budget.retain({ text: "fourth" }).pipe(Effect.result);
43 expect(overflow._tag).toBe("Failure");
44 // Do not resume the consumer. Its source scope must close now.
45 yield* Deferred.await(sourceClosed);
46 yield* budget.closed;
47 expect((yield* budget.usage).retainedItems).toBe(1);
48 }),
49 );
50 expect(yield* budget.usage).toEqual({ retainedItems: 0, retainedSerializedBytes: 0 });
51 }),
52 ),
53 );
54});
55
56it.effect("stops draining a slow subscriber when its unacknowledged tail fills", () =>
57 Effect.scoped(
58 Effect.gen(function* () {
59 const input = yield* Queue.unbounded<{ text: string }>();
60 const sourceClosed = yield* Deferred.make<void>();
61 const pull = yield* Stream.toPull(
62 bufferLiveStream(
63 Stream.fromQueue(input).pipe(Stream.ensuring(Deferred.succeed(sourceClosed, undefined))),
64 { maxItems: 2 },
65 ),
66 );
67 yield* Queue.offer(input, { text: "first" });
68 expect(yield* pull).toEqual([{ text: "first" }]);
69 yield* Queue.offerAll(input, [{ text: "second" }, { text: "third" }]);
70 yield* Deferred.await(sourceClosed);
71 const next = yield* pull.pipe(Effect.result);
72 expect(next._tag).toBe("Failure");
73 if (next._tag === "Failure") expect(next.failure._tag).toBe("LiveStreamBufferError");
74 }),
75 ),
76);
77
78type Event = { readonly sequence: number; readonly text: string; readonly threadId?: string };
79
80describe("replayAndBufferLiveEvents", () => {
81 it.effect.each(["high-water", "replay"] as const)(
82 "unsubscribes and cancels a blocked %s read on live overflow",
83 (phase) =>
84 Effect.scoped(
85 Effect.gen(function* () {
86 const pubsub = yield* PubSub.unbounded<Event>();
87 const subscribed = yield* Deferred.make<PubSub.Subscription<Event>>();
88 const readStarted = yield* Deferred.make<void>();
89 const readClosed = yield* Deferred.make<void>();
90 const blockedRead = Deferred.succeed(readStarted, undefined).pipe(
91 Effect.andThen(Effect.never),
92 Effect.ensuring(Deferred.succeed(readClosed, undefined)),
93 );
94 const pull = yield* Stream.toPull(
95 replayAndBufferLiveEvents(
96 {
97 subscribe: PubSub.subscribe(pubsub).pipe(
98 Effect.tap((subscription) => Deferred.succeed(subscribed, subscription)),
99 ),
100 latestSequence:
101 phase === "high-water" ? blockedRead.pipe(Effect.as(0)) : Effect.succeed(0),
102 replay: () => (phase === "replay" ? Stream.fromEffect(blockedRead) : Stream.empty),
103 },
104 { maxItems: 2 },
105 ),
106 );
107 const reader = yield* pull.pipe(Effect.result, Effect.forkScoped);
108 yield* Deferred.await(readStarted);
109 yield* PubSub.publishAll(pubsub, [
110 { sequence: 1, text: "one" },
111 { sequence: 2, text: "two" },
112 { sequence: 3, text: "overflow" },
113 ]);
114 yield* (yield* Deferred.await(subscribed)).shutdownHook.await;
115 yield* Deferred.await(readClosed);
116 const result = yield* Fiber.join(reader);
117 expect(result._tag).toBe("Failure");
118 if (result._tag === "Failure") expect(result.failure._tag).toBe("LiveStreamBufferError");
119 expect(yield* PubSub.size(pubsub)).toBe(0);
120 // Further publications do not accumulate for the abandoned subscriber.
121 yield* PubSub.publish(pubsub, { sequence: 4, text: "after overflow" });
122 expect(yield* PubSub.size(pubsub)).toBe(0);
123 }),
124 ),
125 );
126
127 it.effect.each(["items", "bytes"] as const)(
128 "counts an unacknowledged replay batch toward the live %s limit",
129 (limit) =>
130 Effect.scoped(
131 Effect.gen(function* () {
132 const pubsub = yield* PubSub.unbounded<Event>();
133 const subscribed = yield* Deferred.make<PubSub.Subscription<Event>>();
134 const replayed = { sequence: 1, text: "replay" };
135 const live = { sequence: 2, text: "live" };
136 const pull = yield* Stream.toPull(
137 replayAndBufferLiveEvents(
138 {
139 subscribe: PubSub.subscribe(pubsub).pipe(
140 Effect.tap((subscription) => Deferred.succeed(subscribed, subscription)),
141 ),
142 latestSequence: Effect.succeed(1),
143 replay: () => Stream.succeed(replayed),
144 },
145 limit === "items"
146 ? { maxItems: 1 }
147 : {
148 maxSerializedBytes: 40,
149 },
150 ),
151 );
152 expect(yield* pull).toEqual([replayed]);
153 yield* PubSub.publish(pubsub, live);
154 // The client has not acknowledged replay; cleanup must still finish.
155 yield* (yield* Deferred.await(subscribed)).shutdownHook.await;
156 expect(yield* PubSub.size(pubsub)).toBe(0);
157 const result = yield* pull.pipe(Effect.result);
158 expect(result._tag).toBe("Failure");
159 if (result._tag === "Failure") expect(result.failure._tag).toBe("LiveStreamBufferError");
160 }),
161 ),
162 );
163
164 it.effect("lets a fast reader replay a database page larger than both buffer limits", () =>
165 Effect.scoped(
166 Effect.gen(function* () {
167 const pubsub = yield* PubSub.unbounded<Event>();
168 const events = Array.from({ length: 20 }, (_, index) => ({
169 sequence: index + 1,
170 text: "historical event",
171 }));
172 const replayed = yield* replayAndBufferLiveEvents(
173 {
174 subscribe: PubSub.subscribe(pubsub),
175 latestSequence: Effect.succeed(20),
176 replay: () => Stream.fromIterable(events),
177 },
178 { maxItems: 2, maxSerializedBytes: 100 },
179 ).pipe(Stream.take(20), Stream.runCollect);
180 expect(replayed).toEqual(events);
181 expect(yield* PubSub.size(pubsub)).toBe(0);
182 }),
183 ),
184 );
185
186 it.effect(
187 "filters unrelated events before buffering and deduplicates the replay/live overlap",
188 () =>
189 Effect.scoped(
190 Effect.gen(function* () {
191 const pubsub = yield* PubSub.unbounded<Event>();
192 const highWaterStarted = yield* Deferred.make<void>();
193 const highWater = yield* Deferred.make<number>();
194 const first = { sequence: 1, text: "first", threadId: "selected" };
195 const overlap = { sequence: 2, text: "overlap", threadId: "selected" };
196 const last = { sequence: 10, text: "last", threadId: "selected" };
197 const reader = yield* replayAndBufferLiveEvents(
198 {
199 subscribe: PubSub.subscribe(pubsub),
200 latestSequence: Deferred.succeed(highWaterStarted, undefined).pipe(
201 Effect.andThen(Deferred.await(highWater)),
202 ),
203 replay: () => Stream.make(first, overlap),
204 filter: (event) => event.threadId === "selected",
205 },
206 { maxItems: 4 },
207 ).pipe(Stream.take(3), Stream.runCollect, Effect.forkScoped);
208 yield* Deferred.await(highWaterStarted);
209 yield* PubSub.publishAll(pubsub, [
210 overlap,
211 ...Array.from({ length: 6 }, (_, index) => ({
212 sequence: index + 3,
213 text: "unrelated",
214 threadId: "other",
215 })),
216 last,
217 ]);
218 yield* Deferred.succeed(highWater, 2);
219 expect(yield* Fiber.join(reader)).toEqual([first, overlap, last]);
220 expect(yield* PubSub.size(pubsub)).toBe(0);
221 }),
222 ),
223 );
224
225 it.effect(
226 "caps raw and projected stages while projection and client acknowledgement are blocked",
227 () =>
228 Effect.scoped(
229 Effect.gen(function* () {
230 const pubsub = yield* PubSub.unbounded<Event>();
231 const subscribed = yield* Deferred.make<PubSub.Subscription<Event>>();
232 const projectionStarted = yield* Deferred.make<void>();
233 const projectionClosed = yield* Deferred.make<void>();
234 const releaseProjection = yield* Deferred.make<void>();
235 const source = replayAndBufferLiveEvents(
236 {
237 subscribe: PubSub.subscribe(pubsub).pipe(
238 Effect.tap((subscription) => Deferred.succeed(subscribed, subscription)),
239 ),
240 latestSequence: Effect.succeed(0),
241 replay: () => Stream.empty,
242 },
243 { maxItems: 2 },
244 ).pipe(
245 Stream.mapEffect((event) =>
246 event.sequence === 3
247 ? Deferred.succeed(projectionStarted, undefined).pipe(
248 Effect.andThen(Deferred.await(releaseProjection)),
249 Effect.as(event),
250 )
251 : Effect.succeed(event),
252 ),
253 Stream.ensuring(Deferred.succeed(projectionClosed, undefined)),
254 );
255 const pull = yield* Stream.toPull(bufferLiveStream(source, { maxItems: 2 }));
256 const firstPull = yield* pull.pipe(Effect.forkScoped);
257 const subscription = yield* Deferred.await(subscribed);
258 yield* PubSub.publish(pubsub, { sequence: 1, text: "awaiting ACK" });
259 expect(yield* Fiber.join(firstPull)).toEqual([{ sequence: 1, text: "awaiting ACK" }]);
260 yield* PubSub.publish(pubsub, { sequence: 2, text: "projected tail" });
261 yield* PubSub.publish(pubsub, { sequence: 3, text: "projecting" });
262 yield* Deferred.await(projectionStarted);
263 yield* PubSub.publishAll(pubsub, [
264 { sequence: 4, text: "raw tail" },
265 { sequence: 5, text: "overflow" },
266 ]);
267 // Two projected slots and two raw slots cannot retain a fifth event.
268 yield* subscription.shutdownHook.await;
269 expect(yield* PubSub.size(pubsub)).toBe(0);
270 yield* Deferred.succeed(releaseProjection, undefined);
271 yield* Deferred.await(projectionClosed);
272 const result = yield* pull.pipe(Effect.result);
273 expect(result._tag).toBe("Failure");
274 if (result._tag === "Failure") expect(result.failure._tag).toBe("LiveStreamBufferError");
275 }),
276 ),
277 );
278});