Skip to content

Commit bffe464

Browse files
juliusmarmingejuliusmarminge
committed
fix(client-runtime): preserve thread state with existing batching
1 parent 83f3c02 commit bffe464

7 files changed

Lines changed: 147 additions & 854 deletions

File tree

‎packages/client-runtime/src/rpc/client.ts‎

Lines changed: 8 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -183,24 +183,12 @@ interface SubscriptionOptions<TTag extends EnvironmentSubscriptionRpcTag> {
183183
readonly resubscribe?: Stream.Stream<unknown, never, never>;
184184
}
185185

186-
export type DynamicSubscriptionGeneration = symbol;
187-
188-
export interface DynamicSubscriptionItem<A> {
189-
readonly generation: DynamicSubscriptionGeneration;
190-
readonly session: RpcSession;
191-
readonly value: A;
192-
}
193-
194186
function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
195187
tag: TTag,
196-
makeInput: (
197-
session: RpcSession,
198-
generation: DynamicSubscriptionGeneration,
199-
) => Effect.Effect<EnvironmentRpcInput<TTag>>,
188+
makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
200189
mapStream: (
201190
session: RpcSession,
202191
stream: Stream.Stream<EnvironmentRpcStreamValue<TTag>, EnvironmentRpcStreamFailure<TTag>>,
203-
generation: DynamicSubscriptionGeneration,
204192
) => Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>>,
205193
options?: SubscriptionOptions<TTag>,
206194
): Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>, EnvironmentSupervisor> {
@@ -234,17 +222,16 @@ function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
234222
EnvironmentRpcStreamFailure<TTag>
235223
>;
236224
const subscribeToSession = (): Stream.Stream<A, EnvironmentRpcStreamFailure<TTag>> =>
237-
Stream.suspend(() => {
238-
const generation: DynamicSubscriptionGeneration = Symbol();
239-
return Stream.unwrap(
225+
Stream.suspend(() =>
226+
Stream.unwrap(
240227
Effect.gen(function* () {
241-
const input = yield* makeInput(session, generation);
228+
const input = yield* makeInput(session);
242229
const completeObservation = yield* observer.observe({
243230
environmentId: supervisor.target.environmentId,
244231
method: tag,
245232
input,
246233
});
247-
return mapStream(session, method(input), generation).pipe(
234+
return mapStream(session, method(input)).pipe(
248235
Stream.ensuring(completeObservation),
249236
);
250237
}),
@@ -300,8 +287,8 @@ function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
300287
}
301288
return Stream.failCause(cause);
302289
}),
303-
);
304-
});
290+
),
291+
);
305292
return subscribeToSession();
306293
},
307294
}),
@@ -315,28 +302,7 @@ function subscribeDynamicMapped<TTag extends EnvironmentSubscriptionRpcTag, A>(
315302
);
316303
}
317304

318-
export function subscribeDynamicWithGeneration<TTag extends EnvironmentSubscriptionRpcTag>(
319-
tag: TTag,
320-
makeInput: (
321-
session: RpcSession,
322-
generation: DynamicSubscriptionGeneration,
323-
) => Effect.Effect<EnvironmentRpcInput<TTag>>,
324-
options?: SubscriptionOptions<TTag>,
325-
): Stream.Stream<
326-
DynamicSubscriptionItem<EnvironmentRpcStreamValue<TTag>>,
327-
EnvironmentRpcStreamFailure<TTag>,
328-
EnvironmentSupervisor
329-
> {
330-
return subscribeDynamicMapped(
331-
tag,
332-
makeInput,
333-
(session, stream, generation) =>
334-
stream.pipe(Stream.map((value) => ({ generation, session, value }))),
335-
options,
336-
);
337-
}
338-
339-
function subscribeDynamic<TTag extends EnvironmentSubscriptionRpcTag>(
305+
export function subscribeDynamic<TTag extends EnvironmentSubscriptionRpcTag>(
340306
tag: TTag,
341307
makeInput: (session: RpcSession) => Effect.Effect<EnvironmentRpcInput<TTag>>,
342308
options?: SubscriptionOptions<TTag>,

‎packages/client-runtime/src/state/shell-sync.test.ts‎

Lines changed: 0 additions & 253 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ import * as Queue from "effect/Queue";
1212
import * as Ref from "effect/Ref";
1313
import * as Stream from "effect/Stream";
1414
import * as SubscriptionRef from "effect/SubscriptionRef";
15-
import * as TestClock from "effect/testing/TestClock";
1615

1716
import {
1817
AVAILABLE_CONNECTION_STATE,
@@ -153,85 +152,6 @@ describe("environment shell synchronization", () => {
153152
}),
154153
);
155154

156-
it.effect("keeps the failure visible while an already buffered shell batch drains", () =>
157-
Effect.gen(function* () {
158-
const items: OrchestrationShellStreamItem[] = [
159-
{ kind: "snapshot", snapshot: LIVE_SHELL_SNAPSHOT },
160-
...Array.from({ length: 300 }, (_, index) => ({
161-
kind: "thread-upserted" as const,
162-
sequence: index + 2,
163-
thread: { id: `thread-${index}` } as never,
164-
})),
165-
{ kind: "synchronized" },
166-
];
167-
let attempts = 0;
168-
const client = {
169-
[ORCHESTRATION_WS_METHODS.subscribeShell]: () =>
170-
attempts++ === 0
171-
? Stream.concat(Stream.fromArray(items), Stream.fail(new Error("stream failed")))
172-
: Stream.concat(
173-
Stream.succeed<OrchestrationShellStreamItem>({ kind: "synchronized" }),
174-
Stream.never,
175-
),
176-
} as unknown as WsRpcProtocolClient;
177-
const supervisorState = yield* SubscriptionRef.make(AVAILABLE_CONNECTION_STATE);
178-
const activeSession = yield* SubscriptionRef.make<Option.Option<RpcSession.RpcSession>>(
179-
Option.some(session(client)),
180-
);
181-
const supervisor = EnvironmentSupervisor.EnvironmentSupervisor.of({
182-
target: TARGET,
183-
state: supervisorState,
184-
session: activeSession,
185-
prepared: yield* SubscriptionRef.make(Option.some(PREPARED)),
186-
connect: Effect.void,
187-
disconnect: Effect.void,
188-
retryNow: Effect.void,
189-
} satisfies EnvironmentSupervisor.EnvironmentSupervisor["Service"]);
190-
const cache = Persistence.EnvironmentCacheStore.of({
191-
loadShell: () => Effect.succeed(Option.none()),
192-
saveShell: () => Effect.void,
193-
loadThread: () => Effect.succeed(Option.none()),
194-
saveThread: () => Effect.void,
195-
removeThread: () => Effect.void,
196-
loadServerConfig: () => Effect.succeed(Option.none()),
197-
saveServerConfig: () => Effect.void,
198-
loadVcsRefs: () => Effect.succeed(Option.none()),
199-
saveVcsRefs: () => Effect.void,
200-
removeVcsRefs: () => Effect.void,
201-
clearVcsRefs: () => Effect.void,
202-
clear: () => Effect.void,
203-
});
204-
// Cold cache with no HTTP snapshot available → falls back to the
205-
// socket-embedded snapshot.
206-
const snapshotLoader = ShellSnapshotLoader.of({
207-
load: () => Effect.succeed(Option.none()),
208-
});
209-
const shellState = yield* makeEnvironmentShellState().pipe(
210-
Effect.provideService(EnvironmentSupervisor.EnvironmentSupervisor, supervisor),
211-
Effect.provideService(Persistence.EnvironmentCacheStore, cache),
212-
Effect.provideService(ShellSnapshotLoader, snapshotLoader),
213-
);
214-
215-
const drained = yield* SubscriptionRef.changes(shellState).pipe(
216-
Stream.filter((value) => Option.getOrNull(value.snapshot)?.snapshotSequence === 301),
217-
Stream.runHead,
218-
Effect.map(Option.getOrThrow),
219-
);
220-
expect(drained.error).toEqual(Option.some("Could not synchronize environment data."));
221-
expect(drained.status).toBe("cached");
222-
expect(Option.getOrThrow(drained.snapshot).threads).toHaveLength(300);
223-
yield* TestClock.adjust("250 millis");
224-
const recovered = yield* SubscriptionRef.changes(shellState).pipe(
225-
Stream.filter((value) => value.status === "live"),
226-
Stream.runHead,
227-
Effect.map(Option.getOrThrow),
228-
);
229-
expect(attempts).toBe(2);
230-
expect(recovered.error).toEqual(Option.none());
231-
expect(Option.getOrThrow(recovered.snapshot).snapshotSequence).toBe(301);
232-
}),
233-
);
234-
235155
it.live.each([
236156
{ bufferSize: Infinity, expectedSequences: [51] },
237157
// RpcClient defaults to a 16-event buffer, which splits larger server chunks.
@@ -531,177 +451,4 @@ describe("environment shell synchronization", () => {
531451
expect(yield* Ref.get(loaderCalls)).toBe(2);
532452
}),
533453
);
534-
535-
it.effect("folds a backlogged shell event flood into a bounded number of publications", () =>
536-
Effect.gen(function* () {
537-
const events = yield* Queue.unbounded<OrchestrationShellStreamItem>();
538-
const client = {
539-
[ORCHESTRATION_WS_METHODS.subscribeShell]: () => Stream.fromQueue(events),
540-
} as unknown as WsRpcProtocolClient;
541-
const supervisorState = yield* SubscriptionRef.make(AVAILABLE_CONNECTION_STATE);
542-
const activeSession = yield* SubscriptionRef.make<Option.Option<RpcSession.RpcSession>>(
543-
Option.some(session(client)),
544-
);
545-
const supervisor = EnvironmentSupervisor.EnvironmentSupervisor.of({
546-
target: TARGET,
547-
state: supervisorState,
548-
session: activeSession,
549-
prepared: yield* SubscriptionRef.make(Option.some(PREPARED)),
550-
connect: Effect.void,
551-
disconnect: Effect.void,
552-
retryNow: Effect.void,
553-
} satisfies EnvironmentSupervisor.EnvironmentSupervisor["Service"]);
554-
const cache = Persistence.EnvironmentCacheStore.of({
555-
loadShell: () => Effect.succeed(Option.none()),
556-
saveShell: () => Effect.void,
557-
loadThread: () => Effect.succeed(Option.none()),
558-
saveThread: () => Effect.void,
559-
removeThread: () => Effect.void,
560-
loadServerConfig: () => Effect.succeed(Option.none()),
561-
saveServerConfig: () => Effect.void,
562-
loadVcsRefs: () => Effect.succeed(Option.none()),
563-
saveVcsRefs: () => Effect.void,
564-
removeVcsRefs: () => Effect.void,
565-
clearVcsRefs: () => Effect.void,
566-
clear: () => Effect.void,
567-
});
568-
const snapshotLoader = ShellSnapshotLoader.of({
569-
load: () => Effect.succeed(Option.none()),
570-
});
571-
const shellState = yield* makeEnvironmentShellState().pipe(
572-
Effect.provideService(EnvironmentSupervisor.EnvironmentSupervisor, supervisor),
573-
Effect.provideService(Persistence.EnvironmentCacheStore, cache),
574-
Effect.provideService(ShellSnapshotLoader, snapshotLoader),
575-
);
576-
const publications = yield* Ref.make(0);
577-
yield* SubscriptionRef.changes(shellState).pipe(
578-
Stream.runForEach(() => Ref.update(publications, (count) => count + 1)),
579-
Effect.forkScoped,
580-
);
581-
582-
yield* Queue.offer(events, { kind: "snapshot", snapshot: LIVE_SHELL_SNAPSHOT });
583-
yield* SubscriptionRef.changes(shellState).pipe(
584-
Stream.filter((state) => Option.isSome(state.snapshot)),
585-
Stream.runHead,
586-
);
587-
const before = yield* Ref.get(publications);
588-
589-
// 400 upserts across 25 threads, offered as one backlog batch.
590-
const flood = Array.from({ length: 400 }, (_, index): OrchestrationShellStreamItem => ({
591-
kind: "thread-upserted",
592-
sequence: index + 2,
593-
thread: { id: `thread-${index % 25}` } as never,
594-
}));
595-
yield* Queue.offerAll(events, flood);
596-
const settled = yield* SubscriptionRef.changes(shellState).pipe(
597-
Stream.filter(
598-
(state) => Option.isSome(state.snapshot) && state.snapshot.value.snapshotSequence === 401,
599-
),
600-
Stream.runHead,
601-
);
602-
603-
// Every upsert applied: all 25 threads landed at the final sequence.
604-
const snapshot = Option.getOrThrow(Option.getOrThrow(settled).snapshot);
605-
expect(snapshot.threads).toHaveLength(25);
606-
// The backlog folded into batches instead of one publication per event.
607-
for (let attempt = 0; attempt < 10; attempt += 1) {
608-
yield* Effect.yieldNow;
609-
}
610-
expect((yield* Ref.get(publications)) - before).toBeLessThan(50);
611-
}),
612-
);
613-
it.effect(
614-
"yields between shell slices and discards the remaining backlog after a session switch",
615-
() =>
616-
Effect.gen(function* () {
617-
const events = yield* Queue.unbounded<OrchestrationShellStreamItem>();
618-
const client = {
619-
[ORCHESTRATION_WS_METHODS.subscribeShell]: () => Stream.fromQueue(events),
620-
} as unknown as WsRpcProtocolClient;
621-
const supervisorState = yield* SubscriptionRef.make(AVAILABLE_CONNECTION_STATE);
622-
const activeSession = yield* SubscriptionRef.make<Option.Option<RpcSession.RpcSession>>(
623-
Option.some(session(client)),
624-
);
625-
const supervisor = EnvironmentSupervisor.EnvironmentSupervisor.of({
626-
target: TARGET,
627-
state: supervisorState,
628-
session: activeSession,
629-
prepared: yield* SubscriptionRef.make(Option.some(PREPARED)),
630-
connect: Effect.void,
631-
disconnect: Effect.void,
632-
retryNow: Effect.void,
633-
} satisfies EnvironmentSupervisor.EnvironmentSupervisor["Service"]);
634-
const cache = Persistence.EnvironmentCacheStore.of({
635-
loadShell: () => Effect.succeed(Option.none()),
636-
saveShell: () => Effect.void,
637-
loadThread: () => Effect.succeed(Option.none()),
638-
saveThread: () => Effect.void,
639-
removeThread: () => Effect.void,
640-
loadServerConfig: () => Effect.succeed(Option.none()),
641-
saveServerConfig: () => Effect.void,
642-
loadVcsRefs: () => Effect.succeed(Option.none()),
643-
saveVcsRefs: () => Effect.void,
644-
removeVcsRefs: () => Effect.void,
645-
clearVcsRefs: () => Effect.void,
646-
clear: () => Effect.void,
647-
});
648-
const snapshotLoader = ShellSnapshotLoader.of({
649-
load: () => Effect.succeed(Option.none()),
650-
});
651-
const shellState = yield* makeEnvironmentShellState().pipe(
652-
Effect.provideService(EnvironmentSupervisor.EnvironmentSupervisor, supervisor),
653-
Effect.provideService(Persistence.EnvironmentCacheStore, cache),
654-
Effect.provideService(ShellSnapshotLoader, snapshotLoader),
655-
);
656-
yield* Queue.offer(events, { kind: "snapshot", snapshot: LIVE_SHELL_SNAPSHOT });
657-
yield* SubscriptionRef.changes(shellState).pipe(
658-
Stream.filter((state) => Option.isSome(state.snapshot)),
659-
Stream.runHead,
660-
);
661-
const replacement = session({
662-
[ORCHESTRATION_WS_METHODS.subscribeShell]: () =>
663-
Stream.make({
664-
kind: "snapshot" as const,
665-
snapshot: { ...LIVE_SHELL_SNAPSHOT, snapshotSequence: 1000 },
666-
}).pipe(Stream.concat(Stream.never)),
667-
} as unknown as WsRpcProtocolClient);
668-
const appliedSequences: number[] = [];
669-
let switched = false;
670-
yield* SubscriptionRef.changes(shellState).pipe(
671-
Stream.runForEach((state) =>
672-
Effect.gen(function* () {
673-
if (Option.isNone(state.snapshot)) return;
674-
const sequence = state.snapshot.value.snapshotSequence;
675-
if (sequence <= 1 || sequence >= 1000) return;
676-
appliedSequences.push(sequence);
677-
if (!switched) {
678-
switched = true;
679-
yield* SubscriptionRef.set(activeSession, Option.some(replacement));
680-
}
681-
}),
682-
),
683-
Effect.forkScoped,
684-
);
685-
686-
// 400 upserts across 25 threads, offered as one backlog batch.
687-
const flood = Array.from({ length: 400 }, (_, index): OrchestrationShellStreamItem => ({
688-
kind: "thread-upserted",
689-
sequence: index + 2,
690-
thread: { id: `thread-${index % 25}` } as never,
691-
}));
692-
yield* Queue.offerAll(events, flood);
693-
const settled = yield* SubscriptionRef.changes(shellState).pipe(
694-
Stream.filter(
695-
(state) =>
696-
Option.isSome(state.snapshot) && state.snapshot.value.snapshotSequence === 1000,
697-
),
698-
Stream.runHead,
699-
);
700-
701-
const snapshot = Option.getOrThrow(Option.getOrThrow(settled).snapshot);
702-
expect(snapshot.threads).toEqual(LIVE_SHELL_SNAPSHOT.threads);
703-
expect(appliedSequences.length).toBeGreaterThan(0);
704-
expect(Math.max(...appliedSequences)).toBeLessThanOrEqual(129);
705-
}),
706-
);
707454
});

0 commit comments

Comments
 (0)