From 07ba25fc31652ecb8d6417cfefdf9ea09192627e Mon Sep 17 00:00:00 2001 From: Makisuo Date: Sat, 31 Oct 2026 02:24:25 +0100 Subject: [PATCH 1/2] perf(sync): carry every sync source over one socket Chrome opens WebSockets to a host one handshake at a time, and each sync source had its own socket, so a load opened the user's object, then each org, in turn: ~6s to channels for a user in 10 orgs (prod trace, 2026-10-31 00:51:30), each upgrade ~60ms on the server and the rest queued in the browser. The org collections are ready only once every org answered, so the sidebar waited for the last handshake. Now the client opens one socket, `GET /sync/mux`, and each source as a numbered stream on it. The api Worker opens each source's object socket as before and relays frames unparsed (`@hazel/sync-do` `mux.ts`): effect `Socket`s for both sides, a `FiberMap` of relays keyed by stream, each relay's scope owning its object socket. The session starts with the response (`forkDetach({ startImmediately })`): workerd drops frames that arrive before a listener does. On the client, `SyncMux` is each connection's `WebSocket` constructor, so reconnects, resumes, `Reauth` and close codes are unchanged. The per-source route stays for older clients. Co-Authored-By: Claude Opus 5.5 --- apps/backend/src/worker/sync-route.test.ts | 31 ++- apps/backend/src/worker/sync-route.ts | 210 +++++++++++--- apps/web-foldkit/src/core/db/sync.ts | 9 +- libs/effect-db-collection/src/hub.ts | 19 +- libs/effect-db-collection/src/index.ts | 1 + libs/effect-db-collection/src/mux.test.ts | 136 +++++++++ libs/effect-db-collection/src/mux.ts | 257 ++++++++++++++++++ .../src/testing/fake-websocket.ts | 11 +- packages/domain/src/sync/index.ts | 1 + packages/domain/src/sync/mux.ts | 59 ++++ packages/sync-do/src/index.ts | 1 + packages/sync-do/src/mux.test.ts | 133 +++++++++ packages/sync-do/src/mux.ts | 157 +++++++++++ packages/sync-do/test/workerd/mux.test.ts | 126 +++++++++ packages/sync-do/test/workerd/worker.ts | 38 ++- 15 files changed, 1137 insertions(+), 52 deletions(-) create mode 100644 libs/effect-db-collection/src/mux.test.ts create mode 100644 libs/effect-db-collection/src/mux.ts create mode 100644 packages/domain/src/sync/mux.ts create mode 100644 packages/sync-do/src/mux.test.ts create mode 100644 packages/sync-do/src/mux.ts create mode 100644 packages/sync-do/test/workerd/mux.test.ts diff --git a/apps/backend/src/worker/sync-route.test.ts b/apps/backend/src/worker/sync-route.test.ts index 112637520..373c065c8 100644 --- a/apps/backend/src/worker/sync-route.test.ts +++ b/apps/backend/src/worker/sync-route.test.ts @@ -4,7 +4,7 @@ */ import { Database } from "@hazel/db" import { CurrentUser, SessionLoadError } from "@hazel/domain" -import { SYNC_SUBPROTOCOL, SyncCloseCode } from "@hazel/domain/sync" +import { SYNC_MUX_SUBPROTOCOL, SYNC_SUBPROTOCOL, SyncCloseCode } from "@hazel/domain/sync" import type { UserId } from "@hazel/schema" import { Effect, Layer, Option } from "effect" import { HttpRouter } from "effect/http" @@ -81,6 +81,13 @@ const build = (options: { }), ), { + /** `GET /sync/mux`, offering `protocols`. */ + mux: (protocols: string) => + handler( + new Request("http://api.test/sync/mux", { + headers: { upgrade: "websocket", "sec-websocket-protocol": protocols }, + }), + ), /** `POST /sync/warm`, with `token` as the bearer (`null`: none). */ warm: (query: string, token: string | null = jwt) => handler( @@ -126,6 +133,28 @@ const recordingNamespace = () => { } } +describe("multiplexed sync socket", () => { + it("closes a bad or missing token with 4001, echoing the mux subprotocol", async () => { + const send = build({ verify: () => Effect.succeed(Option.none()), env: {} }) + const response = await send.mux(`${SYNC_MUX_SUBPROTOCOL}, ${jwt}`) + expect(response.headers.get("sec-websocket-protocol")).toBe(SYNC_MUX_SUBPROTOCOL) + await send.mux(`${SYNC_SUBPROTOCOL}, ${jwt}`) + expect(closed).toEqual([ + { code: SyncCloseCode.Reauth, reason: "invalid_token" }, + { code: SyncCloseCode.Reauth, reason: "missing_token" }, + ]) + }) + + it("closes with 1011 (not 4001) when the session store is down", async () => { + const send = build({ + verify: () => Effect.fail(new SessionLoadError({ message: "down", detail: "db" })), + env: {}, + }) + await send.mux(`${SYNC_MUX_SUBPROTOCOL}, ${jwt}`) + expect(closed).toEqual([{ code: 1011, reason: "unavailable" }]) + }) +}) + describe("sync route", () => { it("forwards the credential's expiry as the principal's expiresAt", async () => { const expiresAt = Date.now() + 300_000 diff --git a/apps/backend/src/worker/sync-route.ts b/apps/backend/src/worker/sync-route.ts index eee77da53..b13aa19d0 100644 --- a/apps/backend/src/worker/sync-route.ts +++ b/apps/backend/src/worker/sync-route.ts @@ -16,17 +16,23 @@ * * A Connect guest's mounted channel streams from its host channel's object (`channel_index`). * + * `GET /sync/mux` carries every source of a client over one socket (`hazel.mux.v1, `): the + * client opens each source with its own token, and the route opens that object's socket the same + * way and relays its frames (`@hazel/sync-do` `mux.ts`). A browser opens WebSockets to a host one handshake + * at a time, so a socket per source costs a load one round of handshakes per source. + * * `POST /sync/warm?channel=` wakes the objects a load is about to open, so a cold one starts * while the web app's script still loads (`index.html`): the caller's `UserObject`, and the * channel's object when the directory knows it (an unknown id must not create one). It carries * the access token as a bearer and answers 204 once they are up; the page does not wait for it. */ -import { SYNC_SUBPROTOCOL, SyncCloseCode } from "@hazel/domain/sync" -import { toHttpServerResponse } from "@hazel/sync-do" +import { SYNC_MUX_SUBPROTOCOL, SYNC_SUBPROTOCOL, SyncCloseCode, type SyncMuxOpen } from "@hazel/domain/sync" +import { type MuxOpenResult, refuseMuxStream, startMuxSession, toHttpServerResponse } from "@hazel/sync-do" import { type Context, Effect, Option, Schema } from "effect" import { HttpRouter, HttpServerRequest, HttpServerResponse } from "effect/http" import { HazelSession } from "../services/hazel-session" -import { ChannelId } from "@hazel/schema" +import { ChannelId, type UserId } from "@hazel/schema" +import type { Socket } from "effect/socket" import { Directory } from "../services/directory" import { isObjectKind, OBJECT_NAMESPACES, type ObjectKind } from "./objects/kinds" import { encodeSyncPrincipal, OBJECT_ID_HEADER, PRINCIPAL_HEADER, WARM_HEADER } from "./objects/principal" @@ -53,32 +59,50 @@ class SyncForwardError extends Schema.TaggedError()("SyncForwa const SYNC_PATH = /^\/sync\/(org|channel|user)\/([^/?]+)/ -/** `hazel.v1, ` → the token. */ -export const tokenFromProtocols = (header: string | undefined): string | undefined => { +/** `, ` (`hazel.v1`, or `hazel.mux.v1` on `/sync/mux`) → the token. */ +export const tokenFromProtocols = ( + header: string | undefined, + protocol: string = SYNC_SUBPROTOCOL, +): string | undefined => { const offered = (header ?? "").split(",").map((entry) => entry.trim()) - if (offered[0] !== SYNC_SUBPROTOCOL) return undefined + if (offered[0] !== protocol) return undefined const token = offered[1] return token === undefined || token.length === 0 ? undefined : token } const isJwt = (token: string) => token.split(".").length === 3 +const makePair = () => + new (globalThis as unknown as { WebSocketPair: new () => Record<0 | 1, WorkerdSocket> }).WebSocketPair() + /** A 101 whose socket is closed right away with `code`, so the browser sees the reason. */ -export const acceptAndClose = (code: number, reason: string) => { - const Pair = (globalThis as unknown as { WebSocketPair: new () => Record<0 | 1, WorkerdSocket> }) - .WebSocketPair - const pair = new Pair() +export const acceptAndClose = (code: number, reason: string, protocol: string = SYNC_SUBPROTOCOL) => { + const pair = makePair() pair[1].accept() pair[1].close(code, reason) return toHttpServerResponse( new Response(null, { status: 101, webSocket: pair[0], - headers: { "Sec-WebSocket-Protocol": SYNC_SUBPROTOCOL }, + headers: { "Sec-WebSocket-Protocol": protocol }, } as unknown as ResponseInit), ) } +/** Where a socket goes once its token checks out, or the close code it gets instead. */ +type SocketTarget = + | { readonly _tag: "Refused"; readonly code: number; readonly reason: string } + | { + readonly _tag: "Target" + readonly userId: UserId + readonly object: { fetch(request: Request): Promise } + /** The principal and object id, which the object trusts (it is reachable only through bindings). */ + readonly headers: { readonly [PRINCIPAL_HEADER]: string; readonly [OBJECT_ID_HEADER]: string } + readonly span: Record + } + +const refusedTarget = (code: number, reason: string): SocketTarget => ({ _tag: "Refused", code, reason }) + /** The object a sync socket for `kind`/`id` opens: a Connect guest mount's host channel. */ export const syncObjectId = (kind: ObjectKind, id: string) => { const channelId = kind === "channel" ? decodeChannelId(id) : Option.none() @@ -177,27 +201,16 @@ export const makeSyncRoute = (env: Record) => }).pipe(Effect.withSpan("SyncRoute.warm")), ) - return yield* router.add( - "GET", - "/sync/*", + /** + * The object a socket for `kind`/`id` opens with `token`, and the principal headers it gets; + * or the close code the client gets instead. + */ + const socketTarget = (kind: ObjectKind, id: string, token: string) => Effect.gen(function* () { - const request = yield* HttpServerRequest.HttpServerRequest - const match = SYNC_PATH.exec(pathOf(request.url)) - const kind = match?.[1] - const id = match?.[2] - if (kind === undefined || id === undefined || !isObjectKind(kind)) { - return HttpServerResponse.text("Not found", { status: 404 }) - } - if (request.headers.upgrade?.toLowerCase() !== "websocket") { - return HttpServerResponse.text("Expected a WebSocket upgrade", { status: 426 }) - } - const token = tokenFromProtocols(request.headers["sec-websocket-protocol"]) - if (token === undefined) return acceptAndClose(SyncCloseCode.Reauth, "missing_token") - const verified = yield* verifySocketToken(session, token) - if (verified._tag === "Unavailable") return acceptAndClose(UNAVAILABLE, "unavailable") + if (verified._tag === "Unavailable") return refusedTarget(UNAVAILABLE, "unavailable") if (Option.isNone(verified.user)) - return acceptAndClose(SyncCloseCode.Reauth, "invalid_token") + return refusedTarget(SyncCloseCode.Reauth, "invalid_token") const { user, expiresAt } = verified.user.value const objectId = yield* syncObjectId(kind, id).pipe( @@ -211,15 +224,131 @@ export const makeSyncRoute = (env: Record) => }).pipe(Effect.as(Option.none())), ), ) - if (Option.isNone(objectId)) return acceptAndClose(UNAVAILABLE, "unavailable") + if (Option.isNone(objectId)) return refusedTarget(UNAVAILABLE, "unavailable") - const namespace = env[OBJECT_NAMESPACES[kind]] as Namespace | undefined - if (namespace === undefined || typeof namespace.getByName !== "function") { + const namespace = namespaceOf(kind) + if (namespace === undefined) { yield* Effect.logError("Sync route: object namespace is not bound", { namespace: OBJECT_NAMESPACES[kind], }) - return acceptAndClose(UNAVAILABLE, "unavailable") + return refusedTarget(UNAVAILABLE, "unavailable") + } + // The object closes the socket with 4001 at `expiresAt`; the client reconnects with a + // fresh token (or sends `Reauth` first). + return { + _tag: "Target", + userId: user.id, + object: namespace.getByName(objectId.value), + headers: { + [PRINCIPAL_HEADER]: encodeSyncPrincipal({ userId: user.id, expiresAt }), + [OBJECT_ID_HEADER]: objectId.value, + }, + span: { "object.kind": kind, "object.id": objectId.value }, + } satisfies SocketTarget as SocketTarget + }) + + /** One stream of a `/sync/mux` socket: its source object's socket, accepted here for the relay. */ + const openMuxStream = (userId: UserId) => (stream: typeof SyncMuxOpen.Type) => + Effect.gen(function* () { + const { kind, sourceId: id, token } = stream + const target = yield* socketTarget(kind, id, token) + if (target._tag === "Refused") return refuseMuxStream(target.code, target.reason) + // One socket is one user's: a source cannot be opened with someone else's token. + if (target.userId !== userId) + return refuseMuxStream(SyncCloseCode.Forbidden, "wrong_user") + const request = new Request(`https://object/sync/${kind}/${encodeURIComponent(id)}`, { + headers: { + upgrade: "websocket", + "sec-websocket-protocol": SYNC_SUBPROTOCOL, + ...target.headers, + }, + }) + const response = yield* Effect.tryPromise({ + try: () => target.object.fetch(request), + catch: (cause) => + new SyncForwardError({ objectId: target.headers[OBJECT_ID_HEADER], cause }), + }).pipe(Effect.withSpan("SyncRoute.forward", { kind: "client", attributes: target.span })) + const socket = ( + response as Response & { webSocket?: (WorkerdSocket & Socket.WebSocketLike) | null } + ).webSocket + if (socket === undefined || socket === null) { + yield* Effect.logError("Sync object answered without a socket", { + source: `${kind}:${id}`, + status: response.status, + }) + return refuseMuxStream(UNAVAILABLE, "unavailable") + } + socket.accept() + return { _tag: "Opened", socket } satisfies MuxOpenResult as MuxOpenResult + }).pipe( + Effect.catchTag("SyncForwardError", (error) => + Effect.logError("Sync upgrade could not reach the object", { + objectId: error.objectId, + cause: String(error.cause), + }).pipe(Effect.as(refuseMuxStream(UNAVAILABLE, "unavailable"))), + ), + Effect.withSpan("SyncMux.open", { + attributes: { "sync.source": `${stream.kind}:${stream.sourceId}` }, + }), + ) + + yield* router.add( + "GET", + "/sync/mux", + Effect.gen(function* () { + const request = yield* HttpServerRequest.HttpServerRequest + if (request.headers.upgrade?.toLowerCase() !== "websocket") { + return HttpServerResponse.text("Expected a WebSocket upgrade", { status: 426 }) + } + const token = tokenFromProtocols( + request.headers["sec-websocket-protocol"], + SYNC_MUX_SUBPROTOCOL, + ) + if (token === undefined) { + return acceptAndClose(SyncCloseCode.Reauth, "missing_token", SYNC_MUX_SUBPROTOCOL) + } + // The socket's token only says whose it is; each source is opened with its own. + const verified = yield* verifySocketToken(session, token) + if (verified._tag === "Unavailable") { + return acceptAndClose(UNAVAILABLE, "unavailable", SYNC_MUX_SUBPROTOCOL) + } + if (Option.isNone(verified.user)) { + return acceptAndClose(SyncCloseCode.Reauth, "invalid_token", SYNC_MUX_SUBPROTOCOL) + } + const pair = makePair() + const server = pair[1] as WorkerdSocket & Socket.WebSocketLike + server.accept() + yield* startMuxSession(server, openMuxStream(verified.user.value.user.id)) + return toHttpServerResponse( + new Response(null, { + status: 101, + webSocket: pair[0], + headers: { "Sec-WebSocket-Protocol": SYNC_MUX_SUBPROTOCOL }, + } as unknown as ResponseInit), + ) + }).pipe(Effect.withSpan("SyncRoute.muxUpgrade")), + ) + + return yield* router.add( + "GET", + "/sync/*", + Effect.gen(function* () { + const request = yield* HttpServerRequest.HttpServerRequest + const match = SYNC_PATH.exec(pathOf(request.url)) + const kind = match?.[1] + const id = match?.[2] + if (kind === undefined || id === undefined || !isObjectKind(kind)) { + return HttpServerResponse.text("Not found", { status: 404 }) } + if (request.headers.upgrade?.toLowerCase() !== "websocket") { + return HttpServerResponse.text("Expected a WebSocket upgrade", { status: 426 }) + } + const token = tokenFromProtocols(request.headers["sec-websocket-protocol"]) + if (token === undefined) return acceptAndClose(SyncCloseCode.Reauth, "missing_token") + + const target = yield* socketTarget(kind, id, token) + if (target._tag === "Refused") return acceptAndClose(target.code, target.reason) + const web = yield* HttpServerRequest.toWeb(request).pipe( Effect.map(Option.some), // `RequestError`: the upgrade request itself could not be converted. @@ -231,20 +360,15 @@ export const makeSyncRoute = (env: Record) => ) if (Option.isNone(web)) return acceptAndClose(UNAVAILABLE, "unavailable") const headers = new Headers(web.value.headers) - // The object closes the socket with 4001 at `expiresAt`; the client reconnects with - // a fresh token. - headers.set(PRINCIPAL_HEADER, encodeSyncPrincipal({ userId: user.id, expiresAt })) - headers.set(OBJECT_ID_HEADER, objectId.value) + for (const [name, value] of Object.entries(target.headers)) headers.set(name, value) const forwarded = new Request(web.value, { headers }) return yield* Effect.tryPromise({ - try: () => namespace.getByName(objectId.value).fetch(forwarded), - catch: (cause) => new SyncForwardError({ objectId: objectId.value, cause }), + try: () => target.object.fetch(forwarded), + catch: (cause) => + new SyncForwardError({ objectId: target.headers[OBJECT_ID_HEADER], cause }), }).pipe( Effect.map(toHttpServerResponse), - Effect.withSpan("SyncRoute.forward", { - kind: "client", - attributes: { "object.kind": kind, "object.id": objectId.value }, - }), + Effect.withSpan("SyncRoute.forward", { kind: "client", attributes: target.span }), Effect.catchTag("SyncForwardError", (error) => Effect.logError("Sync upgrade could not reach the object", { objectId: error.objectId, diff --git a/apps/web-foldkit/src/core/db/sync.ts b/apps/web-foldkit/src/core/db/sync.ts index 88dc611d1..a0399f4a0 100644 --- a/apps/web-foldkit/src/core/db/sync.ts +++ b/apps/web-foldkit/src/core/db/sync.ts @@ -1,8 +1,8 @@ /** * The app's one sync hub (`infra/durable-object-sync-plan.md` §3) and the source sets the - * collections follow. Sockets open `/sync//` with the access token; - * a token a source refuses (4001) is invalidated so the next read exchanges again, and a refused - * exchange dispatches `SESSION_REJECTED_EVENT` as before. + * collections follow. Every source streams over one socket, `/sync/mux`, each + * opened with the access token; a token a source refuses (4001) is invalidated so the next read + * exchanges again, and a refused exchange dispatches `SESSION_REJECTED_EVENT` as before. */ import type { SourceKey, SourceSet } from "@hazel/effect-db-collection" import { SyncHub, sourceSet } from "@hazel/effect-db-collection" @@ -14,6 +14,9 @@ export const syncHub = new SyncHub({ baseUrl: backendUrl, getToken: readAccessToken, invalidateToken: invalidateAccessToken, + // One socket for every source: a browser opens WebSockets to a host one handshake at a time, + // so a socket per org made a load wait one handshake per org, in turn. + multiplex: true, }) /** `user:`, set by the connection manager once `user.me` answered. */ diff --git a/libs/effect-db-collection/src/hub.ts b/libs/effect-db-collection/src/hub.ts index 1cafa74f7..1a93468c5 100644 --- a/libs/effect-db-collection/src/hub.ts +++ b/libs/effect-db-collection/src/hub.ts @@ -7,6 +7,7 @@ import { Effect, type Scope } from "effect" import { SyncConnection, type SyncConnectionStatus, type SyncRunTable } from "./connection" import type { SyncSourceTerminalError } from "./errors" import { dispatchErrorStateChanged } from "./events" +import { SyncMux } from "./mux" import { type BackoffConfig, consoleLogger, @@ -27,6 +28,11 @@ export interface SyncHubConfig { readonly invalidateToken?: (token: string) => void /** Injected in tests; defaults to `globalThis.WebSocket`. */ readonly WebSocket?: WebSocketConstructor + /** + * Carry every source over one socket (`/sync/mux`, `./mux.ts`) instead of a socket per + * source: a browser opens one WebSocket handshake at a time per host. @default false + */ + readonly multiplex?: boolean readonly backoff?: BackoffConfig /** @default 25000 */ readonly heartbeatIntervalMs?: number @@ -90,6 +96,8 @@ interface Entry { export class SyncHub { readonly log: SyncLogger private readonly config: SyncHubConfig + /** The sources' shared socket, when multiplexing. */ + private readonly mux: SyncMux | undefined private readonly entries = new Map() private readonly sinks = new Map>() private readonly statusListeners = new Set<(source: SourceKey, status: SyncConnectionStatus) => void>() @@ -101,6 +109,13 @@ export class SyncHub { constructor(config: SyncHubConfig) { this.config = config this.log = config.log ?? consoleLogger + this.mux = config.multiplex + ? new SyncMux({ + baseUrl: config.baseUrl, + ...(config.WebSocket ? { WebSocket: config.WebSocket } : {}), + log: this.log, + }) + : undefined if (typeof window !== "undefined" && typeof window.addEventListener === "function") { window.addEventListener("online", this.onOnline) } @@ -207,6 +222,7 @@ export class SyncHub { window.removeEventListener("online", this.onOnline) } for (const [source, entry] of this.entries) this.evict(source, entry) + this.mux?.close() this.statusListeners.clear() } @@ -220,7 +236,8 @@ export class SyncHub { baseUrl: this.config.baseUrl, getToken: this.config.getToken, ...(this.config.invalidateToken ? { invalidateToken: this.config.invalidateToken } : {}), - ...(this.config.WebSocket ? { WebSocket: this.config.WebSocket } : {}), + ...(this.mux ? { WebSocket: this.mux.WebSocket } : {}), + ...(!this.mux && this.config.WebSocket ? { WebSocket: this.config.WebSocket } : {}), ...(this.config.backoff ? { backoff: this.config.backoff } : {}), ...(this.config.heartbeatIntervalMs ? { heartbeatIntervalMs: this.config.heartbeatIntervalMs } diff --git a/libs/effect-db-collection/src/index.ts b/libs/effect-db-collection/src/index.ts index a7de01088..55644902b 100644 --- a/libs/effect-db-collection/src/index.ts +++ b/libs/effect-db-collection/src/index.ts @@ -30,6 +30,7 @@ export { type SyncRunInfo, type SyncTableSink, } from "./hub" +export { SyncMux, type SyncMuxOptions } from "./mux" export { HISTORY_PAGE_SIZE, type HistoryDemand, type HistoryState } from "./history" export { extractPartitionValues } from "./subset" // Events diff --git a/libs/effect-db-collection/src/mux.test.ts b/libs/effect-db-collection/src/mux.test.ts new file mode 100644 index 000000000..b05ded676 --- /dev/null +++ b/libs/effect-db-collection/src/mux.test.ts @@ -0,0 +1,136 @@ +import { encodeMuxData, type SyncServerFrame } from "@hazel/domain/sync" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" +import { SyncHub } from "./hub" +import { type FakeSocket, fakeSockets } from "./testing/fake-websocket" + +const silent = () => {} + +const setup = () => { + const sockets = fakeSockets() + let tokenCount = 0 + const invalidateToken = vi.fn() + const hub = new SyncHub({ + baseUrl: "https://sync.test", + getToken: async () => `token-${++tokenCount}`, + invalidateToken, + WebSocket: sockets.WebSocket, + multiplex: true, + backoff: { initialDelayMs: 1000, maxDelayMs: 30_000, multiplier: 2, jitter: false }, + releaseGraceMs: 0, + log: silent, + }) + return { sockets, hub, invalidateToken } +} + +/** Lets the token reads resolve and the streams get created. */ +const flush = () => vi.advanceTimersByTimeAsync(0) + +const send = (socket: FakeSocket, id: number, frame: SyncServerFrame) => + socket.receiveRaw(encodeMuxData(id, JSON.stringify(frame))) + +/** Opens stream `id` and sends it a snapshot at `seq`. */ +const snapshot = (socket: FakeSocket, id: number, seq: number) => { + socket.receive({ _tag: "Opened", id }) + send(socket, id, { _tag: "Begin", reset: true }) + send(socket, id, { _tag: "Commit", seq } as SyncServerFrame) + send(socket, id, { _tag: "Ready" }) +} + +const opens = (socket: FakeSocket) => socket.sent.filter((frame) => frame._tag === "Open") + +beforeEach(() => { + vi.useFakeTimers() +}) +afterEach(() => { + vi.useRealTimers() +}) + +describe("SyncMux", () => { + it("carries every source over one socket, opened with the mux subprotocol", async () => { + const { sockets, hub } = setup() + const user = hub.acquire("user", "u") + const org = hub.acquire("org", "o") + await flush() + expect(sockets.all).toHaveLength(1) + const socket = sockets.all[0]! + expect(socket.url).toBe("wss://sync.test/sync/mux") + expect(socket.protocols).toEqual(["hazel.mux.v1", "token-1"]) + + socket.open() + expect(opens(socket)).toEqual([ + { _tag: "Open", id: 0, kind: "user", sourceId: "u", token: "token-1" }, + { _tag: "Open", id: 1, kind: "org", sourceId: "o", token: "token-2" }, + ]) + socket.receive({ _tag: "Opened", id: 1 }) + expect(socket.sentRaw).toContain( + encodeMuxData(1, JSON.stringify({ _tag: "Hello", since: null, paging: true })), + ) + + snapshot(socket, 0, 4) + snapshot(socket, 1, 9) + expect(user.connection.status.state).toBe("ready") + expect(org.connection.lastSeq).toBe(9) + }) + + it("closes a stream with the code its source refused it with", async () => { + const { sockets, hub } = setup() + const org = hub.acquire("org", "o") + await flush() + const socket = sockets.all[0]! + socket.open() + socket.receive({ _tag: "Closed", id: 0, code: 4003, reason: "forbidden" }) + expect(org.connection.error?._tag).toBe("SyncSourceForbiddenError") + }) + + it("reconnects every source through a new socket when the shared one drops", async () => { + const { sockets, hub } = setup() + const user = hub.acquire("user", "u") + const org = hub.acquire("org", "o") + await flush() + const first = sockets.all[0]! + first.open() + snapshot(first, 0, 4) + snapshot(first, 1, 9) + + first.serverClose(1006) + expect(user.connection.status.state).toBe("error") + await vi.advanceTimersByTimeAsync(1000) + expect(sockets.all).toHaveLength(2) + const second = sockets.all[1]! + second.open() + // New streams, resuming where each source left off. + expect(opens(second).map((frame) => frame.id)).toEqual([2, 3]) + second.receive({ _tag: "Opened", id: 3 }) + expect(second.sentRaw).toContain( + encodeMuxData(3, JSON.stringify({ _tag: "Hello", since: 9, paging: true })), + ) + expect(org.connection.status.state).toBe("connecting") + }) + + it("ignores a closed stream's late frames and closes the socket once idle", async () => { + const { sockets, hub } = setup() + const org = hub.acquire("org", "o") + await flush() + const socket = sockets.all[0]! + socket.open() + snapshot(socket, 0, 4) + + org.release() + expect(socket.sent.at(-1)).toEqual({ _tag: "Close", id: 0 }) + send(socket, 0, { _tag: "Begin", reset: true }) + expect(socket.isClosed).toBe(false) + await vi.advanceTimersByTimeAsync(10_000) + expect(socket.closedWith).toEqual({ code: 1000, reason: "idle" }) + }) + + it("passes a refused socket token to every stream, which exchange a fresh one", async () => { + const { sockets, hub, invalidateToken } = setup() + hub.acquire("org", "o") + await flush() + sockets.all[0]!.serverClose(4001, "invalid_token") + expect(invalidateToken).toHaveBeenCalledWith("token-1") + await flush() + expect(sockets.all).toHaveLength(2) + expect(sockets.all[1]!.protocols).toEqual(["hazel.mux.v1", "token-2"]) + }) +}) diff --git a/libs/effect-db-collection/src/mux.ts b/libs/effect-db-collection/src/mux.ts new file mode 100644 index 000000000..0e7680c55 --- /dev/null +++ b/libs/effect-db-collection/src/mux.ts @@ -0,0 +1,257 @@ +/** + * Every source connection over one WebSocket (`/sync/mux`, `@hazel/domain/sync` `mux.ts`). A + * browser opens WebSockets to a host one handshake at a time, so a socket per source costs a load + * one handshake per source, in turn; through the mux a load pays for one. + * + * The mux hands each {@link SyncConnection} a stream that behaves as its own WebSocket would (it is + * the connection's `WebSocket` constructor), so reconnects, resumes, `Reauth` and close codes work + * unchanged: a stream opens once its source's object accepted it, and closes with the code the + * source's own socket would have. When the shared socket goes, every stream closes with its code, + * and each connection reconnects on its own backoff, through the next shared socket. + */ +import { + decodeMuxData, + encodeMuxData, + SYNC_MUX_SUBPROTOCOL, + type SyncMuxClientControl, + SyncMuxServerControl, + type SyncSourceKind, +} from "@hazel/domain/sync" +import { Result, Schema } from "effect" +import { consoleLogger, type SyncLogger, type WebSocketConstructor, type WebSocketLike } from "./types" + +const CONNECTING = 0 +const OPEN = 1 +const CLOSED = 3 + +/** A socket the server never answered: its streams close with this, and reconnect. */ +const ABNORMAL = 1006 + +const decodeControl = Schema.decodeUnknownResult(Schema.fromJsonString(SyncMuxServerControl)) + +const SOURCE_PATH = /\/sync\/(org|channel|user)\/([^/?]+)$/ + +export interface SyncMuxOptions { + /** The sync Worker's origin, `http(s)://` or `ws(s)://`. */ + readonly baseUrl: string + /** Injected in tests; defaults to `globalThis.WebSocket`. */ + readonly WebSocket?: WebSocketConstructor + /** How long the shared socket stays open without streams, so a reconnect reuses it. @default 10000 */ + readonly idleCloseMs?: number + readonly log?: SyncLogger +} + +/** One source's stream: what its connection sees as its socket. */ +class MuxStream implements WebSocketLike { + readyState = CONNECTING + onopen: ((event: unknown) => void) | null = null + onmessage: ((event: { readonly data: unknown }) => void) | null = null + onclose: ((event: { readonly code: number; readonly reason: string }) => void) | null = null + onerror: ((event: unknown) => void) | null = null + + constructor( + private readonly mux: SyncMux, + readonly id: number, + readonly kind: SyncSourceKind, + readonly sourceId: string, + readonly token: string, + ) {} + + send(data: string): void { + if (this.readyState !== OPEN) throw new Error("The sync stream is not open") + this.mux.sendData(this.id, data) + } + + close(): void { + if (this.readyState === CLOSED) return + this.readyState = CLOSED + this.mux.closeStream(this) + } + + opened(): void { + if (this.readyState !== CONNECTING) return + this.readyState = OPEN + this.onopen?.({}) + } + + deliver(frame: string): void { + if (this.readyState === OPEN) this.onmessage?.({ data: frame }) + } + + closed(code: number, reason: string): void { + if (this.readyState === CLOSED) return + this.readyState = CLOSED + this.onclose?.({ code, reason }) + } +} + +export class SyncMux { + /** Pass as a connection's `WebSocket`: `new` opens a stream to the URL's source. */ + readonly WebSocket: WebSocketConstructor + + private readonly options: SyncMuxOptions + private readonly log: SyncLogger + private readonly url: string + private socket: WebSocketLike | null = null + private isOpen = false + /** Never reused on a socket (nor across them), so a late frame finds no stream. */ + private nextId = 0 + private readonly streams = new Map() + private idleTimer: ReturnType | undefined + + constructor(options: SyncMuxOptions) { + this.options = options + this.log = options.log ?? consoleLogger + this.url = `${options.baseUrl.replace(/\/+$/, "").replace(/^http/, "ws")}/sync/mux` + const mux = this + // A constructor that returns an object: `new` yields the stream. + this.WebSocket = function MuxWebSocket(url: string, protocols: string | Array) { + return mux.openStream(url, protocols) + } as unknown as WebSocketConstructor + } + + /** Closes the shared socket; its streams close as if it dropped. */ + close(): void { + this.clearIdleTimer() + const ws = this.socket + if (ws !== null) this.detach(ws, 1000, "closed") + this.closeStreams(1000, "closed") + } + + sendData(id: number, frame: string): void { + this.socket?.send(encodeMuxData(id, frame)) + } + + closeStream(stream: MuxStream): void { + if (this.streams.get(stream.id) !== stream) return + this.streams.delete(stream.id) + if (this.isOpen) this.sendControl({ _tag: "Close", id: stream.id }) + if (this.streams.size === 0) this.scheduleIdleClose() + } + + private openStream(url: string, protocols: string | Array): MuxStream { + const match = SOURCE_PATH.exec(url) + const token = Array.isArray(protocols) ? protocols[1] : undefined + if (match === null || token === undefined) throw new Error(`Not a sync source socket: ${url}`) + const stream = new MuxStream( + this, + this.nextId++, + match[1] as SyncSourceKind, + decodeURIComponent(match[2]!), + token, + ) + this.streams.set(stream.id, stream) + this.clearIdleTimer() + if (this.socket === null) this.connect(token) + else if (this.isOpen) this.announce(stream) + return stream + } + + /** Opens the shared socket with the token of the stream that needed it. */ + private connect(token: string): void { + const Ctor = this.options.WebSocket ?? (globalThis.WebSocket as unknown as WebSocketConstructor) + let ws: WebSocketLike + try { + ws = new Ctor(this.url, [SYNC_MUX_SUBPROTOCOL, token]) + } catch (cause) { + this.log("warning", "Could not open the sync socket", { cause }) + // The connections retry on their backoff, after this call returned their streams. + queueMicrotask(() => this.closeStreams(ABNORMAL, "unavailable")) + return + } + this.socket = ws + this.isOpen = false + ws.onopen = () => { + if (this.socket !== ws) return + this.isOpen = true + for (const stream of this.streams.values()) this.announce(stream) + } + ws.onmessage = (event) => { + if (this.socket !== ws) return + if (typeof event.data === "string") this.onFrame(event.data) + } + ws.onclose = (event) => { + if (this.socket !== ws) return + this.socket = null + this.isOpen = false + this.clearIdleTimer() + this.closeStreams(event.code, event.reason) + } + // A close event always follows. + ws.onerror = () => {} + } + + private announce(stream: MuxStream): void { + this.sendControl({ + _tag: "Open", + id: stream.id, + kind: stream.kind, + sourceId: stream.sourceId, + token: stream.token, + }) + } + + private onFrame(data: string): void { + const routed = decodeMuxData(data) + if (routed !== undefined) { + this.streams.get(routed.id)?.deliver(routed.frame) + return + } + const decoded = decodeControl(data) + if (Result.isFailure(decoded)) { + this.log("warning", "Unknown sync mux frame", { cause: decoded.failure }) + return + } + const control = decoded.success + const stream = this.streams.get(control.id) + if (stream === undefined) return + if (control._tag === "Opened") return stream.opened() + this.streams.delete(control.id) + if (this.streams.size === 0) this.scheduleIdleClose() + stream.closed(control.code, control.reason) + } + + private sendControl(frame: SyncMuxClientControl): void { + try { + this.socket?.send(JSON.stringify(frame)) + } catch (cause) { + this.log("warning", "Sync mux send failed", { cause }) + } + } + + private closeStreams(code: number, reason: string): void { + const streams = [...this.streams.values()] + this.streams.clear() + for (const stream of streams) stream.closed(code, reason) + } + + private scheduleIdleClose(): void { + this.clearIdleTimer() + this.idleTimer = setTimeout(() => { + this.idleTimer = undefined + const ws = this.socket + if (ws !== null && this.streams.size === 0) this.detach(ws, 1000, "idle") + }, this.options.idleCloseMs ?? 10_000) + } + + private clearIdleTimer(): void { + if (this.idleTimer !== undefined) clearTimeout(this.idleTimer) + this.idleTimer = undefined + } + + private detach(ws: WebSocketLike, code: number, reason: string): void { + if (this.socket === ws) { + this.socket = null + this.isOpen = false + } + ws.onopen = null + ws.onmessage = null + ws.onclose = null + ws.onerror = null + try { + ws.close(code, reason) + } catch { + // Already closed. + } + } +} diff --git a/libs/effect-db-collection/src/testing/fake-websocket.ts b/libs/effect-db-collection/src/testing/fake-websocket.ts index 788b6effc..1fba73f66 100644 --- a/libs/effect-db-collection/src/testing/fake-websocket.ts +++ b/libs/effect-db-collection/src/testing/fake-websocket.ts @@ -9,6 +9,8 @@ export class FakeSocket implements WebSocketLike { onclose: ((event: { readonly code: number; readonly reason: string }) => void) | null = null onerror: ((event: unknown) => void) | null = null readonly sent: Array<{ readonly _tag: string; readonly [key: string]: unknown }> = [] + /** Every frame as sent; `sent` holds the JSON ones parsed. */ + readonly sentRaw: Array = [] closedWith: { readonly code?: number; readonly reason?: string } | undefined constructor( @@ -17,7 +19,8 @@ export class FakeSocket implements WebSocketLike { ) {} send(data: string): void { - this.sent.push(JSON.parse(data)) + this.sentRaw.push(data) + if (data.startsWith("{")) this.sent.push(JSON.parse(data)) } close(code?: number, reason?: string): void { @@ -33,7 +36,11 @@ export class FakeSocket implements WebSocketLike { } receive(frame: SyncServerFrame | Record): void { - this.onmessage?.({ data: JSON.stringify(frame) }) + this.receiveRaw(JSON.stringify(frame)) + } + + receiveRaw(data: string): void { + this.onmessage?.({ data }) } /** A full run: `Begin`, one `Rows` per table, `Commit`. */ diff --git a/packages/domain/src/sync/index.ts b/packages/domain/src/sync/index.ts index 56e16de8b..b4731ba9f 100644 --- a/packages/domain/src/sync/index.ts +++ b/packages/domain/src/sync/index.ts @@ -1 +1,2 @@ export * from "./protocol" +export * from "./mux" diff --git a/packages/domain/src/sync/mux.ts b/packages/domain/src/sync/mux.ts new file mode 100644 index 000000000..9ff551839 --- /dev/null +++ b/packages/domain/src/sync/mux.ts @@ -0,0 +1,59 @@ +/** + * The multiplexed sync socket (`GET /sync/mux`): every source of a client over one WebSocket, so a + * load pays one handshake instead of one per source (browsers open WebSockets to a host one + * handshake at a time). The api Worker opens a socket to each source's object and relays frames. + * + * The client opens a source as a stream with an id of its own (never reused on the socket, so a + * late frame of a stream it closed cannot reach the source's next one). Each stream carries the + * per-source protocol (`./protocol.ts`), unchanged and unparsed: a data frame is `\n`, + * both ways. A control frame is a JSON object (it starts with `{`): the client opens and closes + * streams, the Worker says a stream opened or closed, with the close code a per-source socket + * would have had. + */ +import { Schema } from "effect" +import { SyncSourceKind } from "./protocol" + +/** The subprotocol a client offers on `GET /sync/mux`; the access token rides as the second entry. */ +export const SYNC_MUX_SUBPROTOCOL = "hazel.mux.v1" + +const StreamId = Schema.Int.check(Schema.isGreaterThanOrEqualTo(0)) + +/** + * Open stream `id` to `kind`/`id` with `token` (verified per stream: it carries the expiry the + * object enforces, and a `Reauth` on the stream moves it). + */ +export const SyncMuxOpen = Schema.TaggedStruct("Open", { + id: StreamId, + kind: SyncSourceKind, + sourceId: Schema.String, + token: Schema.String, +}) + +/** Close stream `id`. */ +export const SyncMuxClose = Schema.TaggedStruct("Close", { id: StreamId }) + +export const SyncMuxClientControl = Schema.Union([SyncMuxOpen, SyncMuxClose]) +export type SyncMuxClientControl = typeof SyncMuxClientControl.Type + +/** The source's object accepted stream `id`: its frames follow. */ +export const SyncMuxOpened = Schema.TaggedStruct("Opened", { id: StreamId }) + +/** Stream `id` closed, with the code and reason its own socket would have closed with. */ +export const SyncMuxClosed = Schema.TaggedStruct("Closed", { + id: StreamId, + code: Schema.Int, + reason: Schema.String, +}) + +export const SyncMuxServerControl = Schema.Union([SyncMuxOpened, SyncMuxClosed]) +export type SyncMuxServerControl = typeof SyncMuxServerControl.Type + +export const encodeMuxData = (id: number, frame: string): string => `${id}\n${frame}` + +/** A data frame's stream and frame; `undefined` for a control frame (or garbage). */ +export const decodeMuxData = (data: string): { readonly id: number; readonly frame: string } | undefined => { + const newline = data.indexOf("\n") + if (newline <= 0 || newline > 10) return undefined + const id = Number(data.slice(0, newline)) + return Number.isInteger(id) && id >= 0 ? { id, frame: data.slice(newline + 1) } : undefined +} diff --git a/packages/sync-do/src/index.ts b/packages/sync-do/src/index.ts index 86873d137..40d9ffc18 100644 --- a/packages/sync-do/src/index.ts +++ b/packages/sync-do/src/index.ts @@ -7,6 +7,7 @@ export * from "./alarm.ts" export * from "./changes.ts" export * from "./errors.ts" export * from "./http.ts" +export * from "./mux.ts" export * from "./outbox.ts" export * from "./runtime.ts" export * from "./source.ts" diff --git a/packages/sync-do/src/mux.test.ts b/packages/sync-do/src/mux.test.ts new file mode 100644 index 000000000..1d7049f02 --- /dev/null +++ b/packages/sync-do/src/mux.test.ts @@ -0,0 +1,133 @@ +/** + * `GET /sync/mux`'s relay: streams open and close on the client's control frames, frames pass through + * with their stream id, and an object's close reaches the client with its code. + */ +import { encodeMuxData, type SyncMuxClientControl, type SyncMuxOpen } from "@hazel/domain/sync" +import { Effect, Fiber } from "effect" +import { describe, expect, it } from "vitest" +import { type MuxOpenResult, refuseMuxStream, runMuxSession } from "./mux.ts" + +class FakeSocket extends EventTarget { + readyState = 1 + readonly sent: Array = [] + closedWith: { code: number; reason: string } | undefined + + send(data: string) { + if (this.readyState !== 1) throw new Error("closed") + this.sent.push(data) + } + + close(code = 1000, reason = "") { + if (this.readyState === 3) return + this.readyState = 3 + this.closedWith = { code, reason } + } + + /** A frame from the peer. */ + receive(data: string) { + this.dispatchEvent(new MessageEvent("message", { data })) + } + + /** The peer closed the socket. */ + peerClose(code: number, reason = "") { + this.readyState = 3 + this.dispatchEvent(Object.assign(new Event("close"), { code, reason })) + } +} + +const flush = async () => { + for (let i = 0; i < 5; i++) await new Promise((resolve) => setTimeout(resolve, 0)) +} + +const control = (frame: SyncMuxClientControl) => JSON.stringify(frame) + +const start = (open: (stream: typeof SyncMuxOpen.Type) => Effect.Effect) => { + const client = new FakeSocket() + const fiber = Effect.runFork(runMuxSession(client, open).pipe(Effect.scoped)) + return { + client, + fiber, + controls: () => client.sent.filter((frame) => frame.startsWith("{")).map((f) => JSON.parse(f)), + } +} + +describe("runMuxSession", () => { + it("opens a source and relays its frames both ways", async () => { + const upstream = new FakeSocket() + const opened: Array<[string, string]> = [] + const { client, controls } = start((stream) => { + opened.push([`${stream.kind}:${stream.sourceId}`, stream.token]) + return Effect.succeed({ _tag: "Opened", socket: upstream }) + }) + client.receive(control({ _tag: "Open", id: 1, kind: "org", sourceId: "o1", token: "t1" })) + await flush() + expect(opened).toEqual([["org:o1", "t1"]]) + expect(controls()).toEqual([{ _tag: "Opened", id: 1 }]) + + client.receive(encodeMuxData(1, '{"_tag":"Hello","since":null}')) + upstream.receive('{"_tag":"Ready"}') + await flush() + expect(upstream.sent).toEqual(['{"_tag":"Hello","since":null}']) + expect(client.sent).toContain(encodeMuxData(1, '{"_tag":"Ready"}')) + }) + + it("tells the client a refused source's close code", async () => { + const { client, controls } = start(() => Effect.succeed(refuseMuxStream(4001, "invalid_token"))) + client.receive(control({ _tag: "Open", id: 1, kind: "user", sourceId: "u1", token: "bad" })) + await flush() + expect(controls()).toEqual([{ _tag: "Closed", id: 1, code: 4001, reason: "invalid_token" }]) + }) + + it("forwards the object's close and stops relaying the source", async () => { + const upstream = new FakeSocket() + const { client, controls } = start(() => Effect.succeed({ _tag: "Opened", socket: upstream })) + client.receive(control({ _tag: "Open", id: 1, kind: "channel", sourceId: "c1", token: "t" })) + await flush() + upstream.peerClose(4003, "forbidden") + await flush() + expect(controls()).toContainEqual({ _tag: "Closed", id: 1, code: 4003, reason: "forbidden" }) + client.receive(encodeMuxData(1, '{"_tag":"Ping"}')) + await flush() + expect(upstream.sent).toEqual([]) + }) + + it("closes a source's object socket when the client closes the source", async () => { + const upstream = new FakeSocket() + const { client } = start(() => Effect.succeed({ _tag: "Opened", socket: upstream })) + client.receive(control({ _tag: "Open", id: 1, kind: "org", sourceId: "o1", token: "t" })) + await flush() + client.receive(control({ _tag: "Close", id: 1 })) + await flush() + expect(upstream.closedWith).toEqual({ code: 1000, reason: "closed" }) + }) + + it("keeps a source's streams apart: a closed stream's frames never reach the next", async () => { + const sockets = [new FakeSocket(), new FakeSocket()] + let next = 0 + const { client, controls } = start(() => Effect.succeed({ _tag: "Opened", socket: sockets[next++]! })) + client.receive(control({ _tag: "Open", id: 1, kind: "org", sourceId: "o1", token: "t" })) + client.receive(control({ _tag: "Close", id: 1 })) + client.receive(control({ _tag: "Open", id: 2, kind: "org", sourceId: "o1", token: "t" })) + await flush() + expect(sockets[0]!.closedWith?.code).toBe(1000) + client.receive(encodeMuxData(2, '{"_tag":"Ping"}')) + sockets[1]!.receive('{"_tag":"Pong"}') + await flush() + expect(sockets[0]!.sent).toEqual([]) + expect(sockets[1]!.sent).toEqual(['{"_tag":"Ping"}']) + expect(client.sent).toContain(encodeMuxData(2, '{"_tag":"Pong"}')) + expect(controls()).toContainEqual({ _tag: "Opened", id: 2 }) + }) + + it("closes every object socket when the client goes", async () => { + const sockets = [new FakeSocket(), new FakeSocket()] + let next = 0 + const { client, fiber } = start(() => Effect.succeed({ _tag: "Opened", socket: sockets[next++]! })) + client.receive(control({ _tag: "Open", id: 1, kind: "org", sourceId: "o1", token: "t" })) + client.receive(control({ _tag: "Open", id: 2, kind: "org", sourceId: "o2", token: "t" })) + await flush() + client.peerClose(1001, "going away") + await Effect.runPromise(Fiber.await(fiber)) + expect(sockets.map((socket) => socket.closedWith?.code)).toEqual([1000, 1000]) + }) +}) diff --git a/packages/sync-do/src/mux.ts b/packages/sync-do/src/mux.ts new file mode 100644 index 000000000..ee9a991ef --- /dev/null +++ b/packages/sync-do/src/mux.ts @@ -0,0 +1,157 @@ +/** + * The multiplexed sync socket's relay (`GET /sync/mux`, `@hazel/domain/sync` `mux.ts`): one client + * socket, and per stream the client opens, a socket to that stream's source object. + * + * - The client socket and each object socket are effect `Socket`s (`Socket.fromWebSocket`): a + * pull loop per socket, and a close is a `SocketCloseError` carrying its code. + * - The relays live in a `FiberMap` keyed by stream id: an `Open` runs one, a `Close` interrupts + * it, and the session's scope ends them all when the client goes. Each relay's scope owns its + * object socket, so interrupting it closes that socket. + * - Frames pass through unparsed: the relay only prefixes the stream id (`encodeMuxData`). + * + * The session runs detached from the upgrade request (the response must return for the socket to + * open), started by {@link startMuxSession}. Each fiber has its own dispatcher, and the sockets' + * events resume it, so its work stays in the request's I/O context, which workerd keeps alive while + * the client socket is open. + */ +import { + decodeMuxData, + encodeMuxData, + SyncMuxClientControl, + type SyncMuxOpen, + type SyncMuxServerControl, +} from "@hazel/domain/sync" +import { Array as Arr, Effect, FiberMap, Option, Schema } from "effect" +import { Socket } from "effect/socket" + +const utf8 = new TextDecoder() + +/** What opening a stream's object socket gave: the socket, or the close code the client gets. */ +export type MuxOpenResult = + | { readonly _tag: "Opened"; readonly socket: Socket.WebSocketLike } + | { readonly _tag: "Refused"; readonly code: number; readonly reason: string } + +export const refuseMuxStream = (code: number, reason: string): MuxOpenResult => ({ + _tag: "Refused", + code, + reason, +}) + +/** A relay that failed for a reason other than its object closing it; the client retries. */ +const RELAY_FAILED = 1011 + +const decodeControl = Schema.decodeUnknownOption(Schema.fromJsonString(SyncMuxClientControl)) + +const closeQuietly = (socket: Socket.WebSocketLike, code: number, reason: string) => + Effect.sync(() => { + try { + socket.close(code, reason) + } catch { + // Already closed. + } + }) + +/** + * Relays `client` until it closes. `open` opens one stream's object socket (accepted, so open) + * after checking the stream's token. + */ +export const runMuxSession = Effect.fnUntraced(function* ( + client: Socket.WebSocketLike, + open: (stream: typeof SyncMuxOpen.Type) => Effect.Effect, +) { + const clientSocket = yield* Socket.fromWebSocket(Effect.succeed(client)) + const { pull } = yield* clientSocket.reader + const { write, writeAll } = yield* clientSocket.writer + const relays = yield* FiberMap.make() + // Where a client data frame for each stream goes: the object socket of its running relay. + const upstreams = new Map() + + const sendControl = (frame: SyncMuxServerControl) => write(JSON.stringify(frame)).pipe(Effect.ignore) + + const relay = (stream: typeof SyncMuxOpen.Type) => + Effect.gen(function* () { + const { id } = stream + const opened = yield* open(stream) + if (opened._tag === "Refused") { + return yield* sendControl({ _tag: "Closed", id, code: opened.code, reason: opened.reason }) + } + const upstream = yield* Socket.fromWebSocket( + Effect.acquireRelease(Effect.succeed(opened.socket), (socket) => + closeQuietly(socket, 1000, "closed"), + ), + ) + const { pull: pullUpstream } = yield* upstream.reader + upstreams.set(id, opened.socket) + yield* Effect.addFinalizer(() => + Effect.sync(() => { + if (upstreams.get(id) === opened.socket) upstreams.delete(id) + }), + ) + yield* sendControl({ _tag: "Opened", id }) + // Fails when either side goes: the object closed its socket, or the client is gone. + const error = yield* pullUpstream.pipe( + Effect.flatMap((frames) => + writeAll( + Arr.map(frames, (frame) => + encodeMuxData(id, typeof frame === "string" ? frame : utf8.decode(frame)), + ), + ), + ), + Effect.forever, + Effect.flip, + ) + const reason = error.reason + yield* sendControl( + reason._tag === "SocketCloseError" + ? { _tag: "Closed", id, code: reason.code, reason: reason.closeReason ?? "" } + : { _tag: "Closed", id, code: RELAY_FAILED, reason: "relay_failed" }, + ) + }).pipe(Effect.scoped) + + const onFrame = (data: string | Uint8Array) => + Effect.gen(function* () { + if (typeof data !== "string") return + const routed = decodeMuxData(data) + if (routed !== undefined) { + const upstream = upstreams.get(routed.id) + if (upstream === undefined) return + try { + upstream.send(routed.frame) + } catch { + // The object socket is closing; its relay reports the close. + } + return + } + const control = decodeControl(data) + if (Option.isNone(control)) { + return yield* Effect.logWarning("Sync mux frame unreadable") + } + if (control.value._tag === "Open") { + yield* FiberMap.run(relays, control.value.id, relay(control.value)) + } else { + yield* FiberMap.remove(relays, control.value.id) + } + }) + + // Ends when the client closes (a `SocketError`); the scope then ends every relay. + yield* pull.pipe( + Effect.flatMap((frames) => Effect.forEach(frames, onFrame, { discard: true })), + Effect.forever, + Effect.catchTag("SocketError", () => Effect.void), + ) +}) + +/** + * Runs {@link runMuxSession} detached from the upgrade request. It starts at once, so it listens + * before the response goes out: workerd delivers the client's first frames as soon as the socket + * opens, and drops those that arrive before a listener does. + */ +export const startMuxSession = ( + client: Socket.WebSocketLike, + open: (stream: typeof SyncMuxOpen.Type) => Effect.Effect, +): Effect.Effect => + runMuxSession(client, open).pipe( + Effect.scoped, + Effect.forkDetach({ startImmediately: true }), + Effect.asVoid, + ) diff --git a/packages/sync-do/test/workerd/mux.test.ts b/packages/sync-do/test/workerd/mux.test.ts new file mode 100644 index 000000000..3cbb97377 --- /dev/null +++ b/packages/sync-do/test/workerd/mux.test.ts @@ -0,0 +1,126 @@ +import { env, runInDurableObject, SELF } from "cloudflare:test" +import { decodeMuxData, encodeMuxData, type SyncMuxClientControl } from "@hazel/domain/sync" +import { describe, expect, it } from "vitest" +import type { Env, ToyObject } from "./worker.ts" + +const testEnv = env as unknown as Env + +type Frame = { readonly _tag: string } & Record + +/** A frame off the shared socket: a control frame (`id` undefined) or a stream's. */ +interface Received { + readonly id: number | undefined + readonly frame: Frame +} + +const write = async (name: string, id: string) => + ( + (await ( + await testEnv.TOY.get(testEnv.TOY.idFromName(name)).fetch(`http://toy/write?id=${id}`) + ).json()) as { + seq: number + } + ).seq + +const socketsOf = (name: string) => + runInDurableObject( + testEnv.TOY.get(testEnv.TOY.idFromName(name)), + (_: ToyObject, state) => state.getWebSockets().length, + ) + +/** Opens the shared socket through the test Worker's relay and collects what it carries. */ +const connectMux = async () => { + const response = await SELF.fetch("http://worker/mux", { headers: { Upgrade: "websocket" } }) + expect(response.status).toBe(101) + const socket = response.webSocket! + socket.accept() + const received: Array = [] + const waiters: Array<() => void> = [] + socket.addEventListener("message", (event) => { + const data = event.data as string + const routed = decodeMuxData(data) + received.push( + routed === undefined + ? { id: undefined, frame: JSON.parse(data) } + : { id: routed.id, frame: JSON.parse(routed.frame) }, + ) + for (const wake of waiters.splice(0)) wake() + }) + const waitFor = async (predicate: (received: Received) => boolean) => { + const deadline = Date.now() + 5_000 + while (true) { + const found = received.find(predicate) + if (found !== undefined) return found + if (Date.now() > deadline) throw new Error(`Timed out; got ${JSON.stringify(received)}`) + await new Promise((resolve) => { + waiters.push(resolve) + setTimeout(resolve, 100) + }) + } + } + return { + socket, + received, + waitFor, + control: (frame: SyncMuxClientControl) => socket.send(JSON.stringify(frame)), + send: (id: number, frame: Frame) => socket.send(encodeMuxData(id, JSON.stringify(frame))), + } +} + +const until = async (check: () => Promise) => { + const deadline = Date.now() + 5_000 + while (!(await check())) { + if (Date.now() > deadline) throw new Error("Timed out") + await new Promise((resolve) => setTimeout(resolve, 50)) + } +} + +describe("the multiplexed socket on workerd", () => { + it("streams several objects over one socket, live, and closes their sockets with it", async () => { + await write("mux-a", "a1") + await write("mux-b", "b1") + const mux = await connectMux() + mux.control({ _tag: "Open", id: 1, kind: "org", sourceId: "mux-a", token: "bob" }) + mux.control({ _tag: "Open", id: 2, kind: "org", sourceId: "mux-b", token: "bob" }) + await mux.waitFor(({ id, frame }) => id === undefined && frame._tag === "Opened" && frame.id === 2) + await mux.waitFor(({ id, frame }) => id === undefined && frame._tag === "Opened" && frame.id === 1) + + mux.send(1, { _tag: "Hello", since: null }) + mux.send(2, { _tag: "Hello", since: null }) + await mux.waitFor(({ id, frame }) => id === 1 && frame._tag === "Ready") + await mux.waitFor(({ id, frame }) => id === 2 && frame._tag === "Ready") + const rowsOf = (stream: number) => + mux.received + .filter(({ id, frame }) => id === stream && frame._tag === "Rows") + .map(({ frame }) => frame.upserts) + expect(rowsOf(1)).toEqual([[{ id: "a1", name: "a1", private: false }]]) + expect(rowsOf(2)).toEqual([[{ id: "b1", name: "b1", private: false }]]) + + // A write after the snapshot reaches its stream, and only its stream. + const seq = await write("mux-b", "b2") + await mux.waitFor(({ id, frame }) => id === 2 && frame._tag === "Commit" && frame.seq === seq) + expect(rowsOf(1)).toHaveLength(1) + + mux.send(1, { _tag: "Ping" }) + await mux.waitFor(({ id, frame }) => id === 1 && frame._tag === "Pong") + + // Closing a stream closes its object's socket; the other stays. + mux.control({ _tag: "Close", id: 1 }) + await until(async () => (await socketsOf("mux-a")) === 0) + expect(await socketsOf("mux-b")).toBe(1) + + // The client going closes every object socket. + mux.socket.close(1000, "done") + await until(async () => (await socketsOf("mux-b")) === 0) + }) + + it("reports an object's close to its stream with the object's code", async () => { + const mux = await connectMux() + mux.control({ _tag: "Open", id: 7, kind: "org", sourceId: "mux-c", token: "bob" }) + await mux.waitFor(({ frame }) => frame._tag === "Opened") + // The toy source closes a socket whose principal expired with 4001. + mux.control({ _tag: "Open", id: 8, kind: "org", sourceId: "mux-c", token: "eve&exp=1" }) + const closed = await mux.waitFor(({ frame }) => frame._tag === "Closed" && frame.id === 8) + expect(closed.frame).toEqual({ _tag: "Closed", id: 8, code: 4001, reason: "token_expired" }) + }) +}) diff --git a/packages/sync-do/test/workerd/worker.ts b/packages/sync-do/test/workerd/worker.ts index 2a919411f..faad23455 100644 --- a/packages/sync-do/test/workerd/worker.ts +++ b/packages/sync-do/test/workerd/worker.ts @@ -6,7 +6,16 @@ import * as SqliteDo from "@effect/sql-sqlite-do/SqliteClient" import { DurableObject } from "cloudflare:workers" import { Effect, ManagedRuntime } from "effect" import { SqlClient } from "effect/sql/SqlClient" -import { enqueue, makeSyncedObject, type SyncObjectState, type SyncSocket } from "../../src/index.ts" +import type { SyncMuxOpen } from "@hazel/domain/sync" +import { + enqueue, + makeSyncedObject, + type MuxOpenResult, + refuseMuxStream, + startMuxSession, + type SyncObjectState, + type SyncSocket, +} from "../../src/index.ts" import { makeToyControls, makeToySource, TOY_DDL, type ToyPrincipal } from "../../src/test/toy-source.ts" export interface Env { @@ -110,6 +119,31 @@ export class ToyObject extends DurableObject { } } +/** + * `/mux`: the multiplexed socket's relay as the api Worker runs it, detached from the upgrade + * request. A stream opens the toy object named by its `sourceId`, as the user named by its token. + */ +const muxUpgrade = (env: Env) => + Effect.gen(function* () { + const pair = new WebSocketPair() + const server = pair[1] + server.accept() + const open = (stream: typeof SyncMuxOpen.Type) => + Effect.promise(async (): Promise => { + const stub = env.TOY.get(env.TOY.idFromName(stream.sourceId)) + const response = await stub.fetch(`http://toy/connect?user=${stream.token}`, { + headers: { Upgrade: "websocket", "Sec-WebSocket-Protocol": "hazel.v1, secret-token" }, + }) + const socket = response.webSocket + if (socket === null) return refuseMuxStream(1011, "unavailable") + socket.accept() + return { _tag: "Opened", socket: socket as never } + }) + yield* startMuxSession(server as never, open) + return new Response(null, { status: 101, webSocket: pair[0] }) + }) + export default { - fetch: () => new Response("toy"), + fetch: (request: Request, env: Env) => + new URL(request.url).pathname === "/mux" ? Effect.runPromise(muxUpgrade(env)) : new Response("toy"), } From 920af6f2a54f0842f9d5c4a513887e6abbd8e290 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Sat, 31 Oct 2026 02:30:28 +0100 Subject: [PATCH 2/2] fix(sync): close a stream's object socket when the stream closes mid-open The relay registered the socket's close only after `open` returned, so a `Close` (or the client going) while the object was still answering interrupted the relay and left the object's socket open. Opening is now the acquire of an `acquireRelease`: uninterruptible, so the interrupt waits for the socket and the release closes it. Co-Authored-By: Claude Opus 5.5 --- packages/sync-do/src/mux.test.ts | 16 +++++++++++++++- packages/sync-do/src/mux.ts | 12 ++++++------ 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/packages/sync-do/src/mux.test.ts b/packages/sync-do/src/mux.test.ts index 1d7049f02..8b2f97314 100644 --- a/packages/sync-do/src/mux.test.ts +++ b/packages/sync-do/src/mux.test.ts @@ -3,7 +3,7 @@ * with their stream id, and an object's close reaches the client with its code. */ import { encodeMuxData, type SyncMuxClientControl, type SyncMuxOpen } from "@hazel/domain/sync" -import { Effect, Fiber } from "effect" +import { Deferred, Effect, Fiber } from "effect" import { describe, expect, it } from "vitest" import { type MuxOpenResult, refuseMuxStream, runMuxSession } from "./mux.ts" @@ -119,6 +119,20 @@ describe("runMuxSession", () => { expect(controls()).toContainEqual({ _tag: "Opened", id: 2 }) }) + it("closes the socket of a stream closed while its object was still answering", async () => { + const upstream = new FakeSocket() + const answer = Effect.runSync(Deferred.make()) + const { client, controls } = start(() => Deferred.await(answer)) + client.receive(control({ _tag: "Open", id: 1, kind: "org", sourceId: "o1", token: "t" })) + await flush() + client.receive(control({ _tag: "Close", id: 1 })) + await flush() + Effect.runSync(Deferred.succeed(answer, { _tag: "Opened", socket: upstream })) + await flush() + expect(upstream.closedWith).toEqual({ code: 1000, reason: "closed" }) + expect(controls()).toEqual([]) + }) + it("closes every object socket when the client goes", async () => { const sockets = [new FakeSocket(), new FakeSocket()] let next = 0 diff --git a/packages/sync-do/src/mux.ts b/packages/sync-do/src/mux.ts index ee9a991ef..857a1ad85 100644 --- a/packages/sync-do/src/mux.ts +++ b/packages/sync-do/src/mux.ts @@ -71,15 +71,15 @@ export const runMuxSession = Effect.fnUntraced(function* ( const relay = (stream: typeof SyncMuxOpen.Type) => Effect.gen(function* () { const { id } = stream - const opened = yield* open(stream) + // Opening is the acquire, so it is uninterruptible: a stream closed (or a client gone) + // while its object answers waits for the socket, and the release closes it. + const opened = yield* Effect.acquireRelease(open(stream), (result) => + result._tag === "Opened" ? closeQuietly(result.socket, 1000, "closed") : Effect.void, + ) if (opened._tag === "Refused") { return yield* sendControl({ _tag: "Closed", id, code: opened.code, reason: opened.reason }) } - const upstream = yield* Socket.fromWebSocket( - Effect.acquireRelease(Effect.succeed(opened.socket), (socket) => - closeQuietly(socket, 1000, "closed"), - ), - ) + const upstream = yield* Socket.fromWebSocket(Effect.succeed(opened.socket)) const { pull: pullUpstream } = yield* upstream.reader upstreams.set(id, opened.socket) yield* Effect.addFinalizer(() =>