perf(sync): carry every sync source over one socket

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

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, <token>`): 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=<id>` 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<SyncForwardError>()("SyncForwa
 
 const SYNC_PATH = /^\/sync\/(org|channel|user)\/([^/?]+)/
 
-/** `hazel.v1, <token>` → the token. */
-export const tokenFromProtocols = (header: string | undefined): string | undefined => {
+/** `<protocol>, <token>` (`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<Response> }
+			/** 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<string, string>
+	  }
+
+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<string, unknown>) =>
 				}).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<string, unknown>) =>
 							}).pipe(Effect.as(Option.none<string>())),
 						),
 					)
-					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<string, unknown>) =>
 					)
 					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,

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 `<VITE_BACKEND_URL>/sync/<kind>/<id>` 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, `<VITE_BACKEND_URL>/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:<me>`, set by the connection manager once `user.me` answered. */