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..8b2f97314 --- /dev/null +++ b/packages/sync-do/src/mux.test.ts @@ -0,0 +1,147 @@ +/** + * `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 { Deferred, 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 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 + 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..857a1ad85 --- /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 + // 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.succeed(opened.socket)) + 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"), }