1import * as Cause from "effect/Cause";
2import * as Clock from "effect/Clock";
3import * as Config from "effect/Config";
4import * as Effect from "effect/Effect";
5import * as Context from "effect/Context";
6import * as Deferred from "effect/Deferred";
7import * as Exit from "effect/Exit";
8import * as Queue from "effect/Queue";
9import * as Schema from "effect/Schema";
10import * as Stream from "effect/Stream";
11
12const MAX_EVENTS = 64;
13const MAX_BYTES = 1024 * 1024;
14const encodeQueue = Schema.encodeEffect(Schema.fromJsonString(Schema.Unknown));
15
16export const enabled = Config.Boolean("T3CODE_PROVIDER_COMMIT_BATCHING").pipe(
17 Config.withDefault(true),
18 Effect.orDie,
19);
20
21const isSnapshot = (event: { readonly type: string }) =>
22 event.type === "node.updated" ||
23 event.type === "message.updated" ||
24 event.type === "turn_item.updated" ||
25 event.type === "plan.updated";
26
27export class CommitBatchError extends Schema.TaggedError<CommitBatchError>()("CommitBatchError", {
28 cause: Schema.Defect(),
29}) {}
30export const isCommitBatchError = Schema.is(CommitBatchError);
31
32interface Job {
33 readonly owner: { active: boolean };
34 readonly queuedAt: number;
35 readonly snapshot: boolean;
36 readonly bytes: number;
37 readonly execute: Effect.Effect<void, CommitBatchError>;
38 readonly settle: (
39 exit: Exit.Exit<void, CommitBatchError>,
40 preserveFailure: boolean,
41 ) => Effect.Effect<void>;
42 readonly cancel: Effect.Effect<void>;
43}
44interface Batch {
45 readonly events: ReadonlyArray<Job>;
46 readonly bytes: number;
47 readonly reason: "ready" | "limit" | "control";
48}
49
50function collect(queue: Queue.Queue<Job>) {
51 return Stream.fromEffectRepeat(Queue.takeBetween(queue, 1, MAX_EVENTS)).pipe(
52 Stream.flatMap((ready) => {
53 const batches: Array<Batch> = [];
54 let events: Array<Job> = [];
55 let bytes = 0;
56 const flush = (reason: Batch["reason"]) => {
57 if (events.length > 0) batches.push({ events, bytes, reason });
58 events = [];
59 bytes = 0;
60 };
61 for (const job of ready) {
62 if (!job.snapshot) {
63 flush("control");
64 batches.push({ events: [job], bytes: 0, reason: "control" });
65 } else {
66 if (events.length > 0 && bytes + job.bytes > MAX_BYTES) flush("limit");
67 events.push(job);
68 bytes += job.bytes;
69 if (events.length >= MAX_EVENTS || bytes >= MAX_BYTES) flush("limit");
70 }
71 }
72 flush("ready");
73 return Stream.fromArray(batches);
74 }),
75 );
76}
77
78const CurrentWorker = Context.Reference<number | undefined>(
79 "t3/ProviderCommitBatch/CurrentWorker",
80 {
81 defaultValue: () => undefined,
82 },
83);
84
85export const make = Effect.fn("ProviderCommitBatch.make")(function* (
86 commit: (work: Effect.Effect<void, CommitBatchError>) => Effect.Effect<void, CommitBatchError>,
87 enabled: boolean,
88 isolate: (
89 work: Effect.Effect<void, CommitBatchError>,
90 ) => Effect.Effect<void, CommitBatchError> = (work) => work,
91) {
92 const queue = yield* Queue.bounded<Job>(MAX_EVENTS);
93 const pending = new Set<Job>();
94 let closed = false;
95 const shutdown = Effect.gen(function* () {
96 closed = true;
97 yield* Queue.shutdown(queue);
98 for (const job of pending) yield* job.cancel;
99 });
100 yield* Effect.addFinalizer(() => shutdown);
101 // These helpers run once per raw event; the shared transaction owns the tracing span.
102 const enqueue = Effect.fnUntraced(function* <A, E, R>(
103 effect: Effect.Effect<A, E, R>,
104 token: { active: boolean },
105 snapshot: boolean,
106 bytes: number,
107 ) {
108 const caller = yield* Effect.context<R>();
109 const reply = yield* Deferred.make<A, E | CommitBatchError>();
110 let ownFailure: Cause.Cause<E> | undefined;
111 let result: A;
112 const job: Job = {
113 owner: token,
114 queuedAt: yield* Clock.currentTimeMillis,
115 snapshot,
116 bytes,
117 execute: effect.pipe(
118 Effect.tapCause((cause) =>
119 Effect.sync(() => {
120 ownFailure = cause;
121 }),
122 ),
123 Effect.mapError((cause) => new CommitBatchError({ cause })),
124 Effect.provideContext(caller),
125 Effect.tap((value) =>
126 Effect.sync(() => {
127 result = value;
128 }),
129 ),
130 Effect.asVoid,
131 ),
132 settle: Effect.fnUntraced(function* (
133 exit: Exit.Exit<void, CommitBatchError>,
134 preserveFailure: boolean,
135 ) {
136 pending.delete(job);
137 if (Exit.isFailure(exit)) {
138 token.active = false;
139 yield* Deferred.failCause(
140 reply,
141 preserveFailure && ownFailure !== undefined ? ownFailure : exit.cause,
142 );
143 } else yield* Deferred.succeed(reply, result);
144 }),
145 cancel: Effect.gen(function* () {
146 pending.delete(job);
147 yield* Deferred.interrupt(reply);
148 }),
149 };
150 yield* Effect.uninterruptibleMask((restore) =>
151 Effect.gen(function* () {
152 pending.add(job);
153 const accepted = yield* restore(Queue.offer(queue, job)).pipe(
154 Effect.onError(() => job.cancel),
155 );
156 if (!accepted) yield* job.cancel;
157 }),
158 );
159 return Deferred.await(reply);
160 });
161 if (enabled) {
162 yield* collect(queue).pipe(
163 Stream.runForEach((batch) =>
164 Effect.uninterruptibleMask((restore) =>
165 Effect.gen(function* () {
166 const jobs = batch.events.filter((job) => job.owner.active);
167 for (const job of batch.events) if (!job.owner.active) yield* job.cancel;
168 if (jobs.length === 0) return;
169 yield* Effect.annotateCurrentSpan({
170 "orchestration_v2.work_item_count": jobs.length,
171 "orchestration_v2.consumer_count": new Set(jobs.map((job) => job.owner)).size,
172 "orchestration_v2.raw_json_bytes": batch.bytes,
173 "orchestration_v2.batch_flush_reason": batch.reason,
174 "orchestration_v2.backlog_batching": true,
175 "orchestration_v2.batch_wait_ms":
176 (yield* Clock.currentTimeMillis) - Math.min(...jobs.map((job) => job.queuedAt)),
177 });
178 const outcomes = new Map<Job, Exit.Exit<void, CommitBatchError>>();
179 const work = Effect.forEach(
180 jobs,
181 (job) =>
182 Effect.gen(function* () {
183 if (!job.owner.active) return yield* job.cancel;
184 const outcome = yield* (
185 jobs.length > 1 && job.bytes > 0 ? isolate(job.execute) : job.execute
186 ).pipe(Effect.exit);
187 outcomes.set(job, outcome);
188 if (Exit.isFailure(outcome)) {
189 job.owner.active = false;
190 if (jobs.length === 1) return yield* Effect.failCause(outcome.cause);
191 }
192 }),
193 { discard: true },
194 );
195 const exit = yield* restore(
196 (jobs.every((job) => job.bytes === 0) ? work : commit(work)).pipe(
197 Effect.provideService(CurrentWorker, yield* Effect.fiberId),
198 ),
199 ).pipe(Effect.exit);
200 for (const job of jobs) {
201 const outcome = outcomes.get(job) ?? Exit.void;
202 yield* job.settle(
203 Exit.isSuccess(exit) ? outcome : exit,
204 Exit.isSuccess(exit) ||
205 (Exit.isFailure(exit) &&
206 Exit.isFailure(outcome) &&
207 Cause.squash(exit.cause) === Cause.squash(outcome.cause)),
208 );
209 }
210 }),
211 ).pipe(Effect.withSpan("orchestrationV2.ProviderCommitBatch.flush")),
212 ),
213 Effect.ensuring(shutdown),
214 Effect.forkScoped,
215 );
216 }
217 const barrier = Effect.fnUntraced(function* <A, E, R>(effect: Effect.Effect<A, E, R>) {
218 if (closed) return yield* Effect.interrupt;
219 if (!enabled || pending.size === 0 || (yield* CurrentWorker) === (yield* Effect.fiberId))
220 return yield* effect;
221 const token = { active: true };
222 return yield* enqueue(Effect.exit(effect), token, false, 0).pipe(
223 Effect.flatten,
224 Effect.orDie,
225 Effect.flatten,
226 Effect.ensuring(
227 Effect.sync(() => {
228 token.active = false;
229 }),
230 ),
231 );
232 });
233 const consume = Effect.fnUntraced(function* <A extends { readonly type: string }, E, R, E2, R2>(
234 source: Stream.Stream<A, E, R>,
235 process: (event: A) => Effect.Effect<void, E2, R2>,
236 after: Effect.Effect<boolean, E2, R2>,
237 canBatch: Effect.Effect<boolean, never, R2> = Effect.succeed(true),
238 ) {
239 if (!enabled)
240 return yield* source.pipe(
241 Stream.mapEffect((event) => process(event).pipe(Effect.andThen(after))),
242 Stream.takeUntil(Boolean),
243 Stream.runDrain,
244 );
245 const token = { active: true };
246 const signals = source.pipe(
247 Stream.map((event) => ({ _tag: "Event" as const, event })),
248 Stream.catchCause((cause) => Stream.succeed({ _tag: "Failure" as const, cause })),
249 Stream.concat(Stream.succeed({ _tag: "End" as const })),
250 );
251 let resume: Effect.Effect<void> = Effect.void;
252 yield* signals.pipe(
253 Stream.mapEffect((signal) =>
254 Effect.gen(function* () {
255 yield* resume;
256 const processed = yield* Deferred.make<void>();
257 const event = signal._tag === "Event" ? signal.event : undefined;
258 // EOF only acknowledges prior work; it can share a commit with other consumers.
259 const snapshot =
260 signal._tag === "End" ||
261 (event !== undefined && isSnapshot(event) && (yield* canBatch));
262 const bytes =
263 snapshot && event !== undefined ? Buffer.byteLength(yield* encodeQueue(event)) : 0;
264 const work: Effect.Effect<boolean, E | E2, R2> =
265 signal._tag === "Failure"
266 ? Effect.failCause(signal.cause)
267 : event === undefined
268 ? Effect.succeed(true)
269 : process(event).pipe(Effect.as(false));
270 const reply = snapshot
271 ? yield* enqueue(work, token, true, bytes)
272 : (yield* enqueue(Effect.void, token, false, 0)).pipe(Effect.andThen(work));
273 if (!snapshot) resume = Deferred.await(processed);
274 return { reply, processed, check: !snapshot };
275 }),
276 ),
277 Stream.buffer({ capacity: MAX_EVENTS }),
278 Stream.mapEffect(({ reply, processed, check }) =>
279 reply.pipe(
280 Effect.flatMap((end) =>
281 end ? Effect.succeed(true) : check ? after : Effect.succeed(false),
282 ),
283 Effect.tap((stop) =>
284 Effect.sync(() => {
285 if (stop) token.active = false;
286 }),
287 ),
288 Effect.ensuring(Deferred.succeed(processed, undefined)),
289 ),
290 ),
291 Stream.takeUntil(Boolean),
292 Stream.runDrain,
293 Effect.ensuring(
294 Effect.sync(() => {
295 token.active = false;
296 }),
297 ),
298 );
299 });
300 return { barrier, consume };
301});