From a46b9373d46f400112b4ef79e64c820ca005a702 Mon Sep 17 00:00:00 2001 From: Arul Sharma <31745423+arul28@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:52:54 -0400 Subject: [PATCH 1/3] perf(relay): cache socket state, add /prewarm, hint DO placement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three additive, independently-deployable improvements to the tunnel relay Worker. All are backward compatible with deployed desktop and iOS clients: no auth, close-code, protocol, framing, or buffer-limit changes. - Cache each socket's attachment and its paired partner in instance memory. `normalizedAttachment()` deserialized on every frame and `partnerOf()` rescanned the connection tag on every frame; both now read from a WeakMap. Hibernation drops instance memory, so both caches fall back to the durable attachment (or the tag scan) and repopulate themselves after a wake, and a single write-through helper is now the only writer of an attachment so the cache cannot drift. The partner entry is only ever a hint, re-checked against the same predicate the scan uses. - Add `GET /prewarm/:machineKey`, returning `{ok, control}`. Reaching the object is the point: it un-hibernates it before a client needs the tunnel. Inert by construction — no storage, no alarm, no attachment migration, no signal to the host. Same auth stance as `/connect`. `/health` is answered by the Worker and never reaches a DO, so it could not serve this. - Derive a `locationHint` from `request.cf` on `claim` and `/host` only, so a new machine's object is created near the machine rather than near whichever request happened to arrive first. `/connect`, `/prewarm`, and `/pipe` never carry a hint — a travelling phone must not place a machine's object. Co-Authored-By: Claude Fable 5 --- apps/tunnel-relay/README.md | 18 +++ apps/tunnel-relay/src/relay.ts | 75 +++++++++- apps/tunnel-relay/src/tunnelDo.ts | 88 +++++++++-- apps/tunnel-relay/test/relay.test.ts | 89 +++++++++++ apps/tunnel-relay/test/relay.workerd.test.ts | 146 +++++++++++++++++++ 5 files changed, 405 insertions(+), 11 deletions(-) diff --git a/apps/tunnel-relay/README.md b/apps/tunnel-relay/README.md index 78f53bc61..45e5acf7f 100644 --- a/apps/tunnel-relay/README.md +++ b/apps/tunnel-relay/README.md @@ -65,6 +65,24 @@ Same claim + HMAC design as `apps/push-relay`: - **Client** `/connect/:machineKey` — no Worker-level auth beyond `machineKey` unguessability. This is only transport admission; ADE authorizes the socket after it reaches the host. +- **Prewarm** `GET /prewarm/:machineKey` — same stance as `/connect`. Returns + `{ok:true, control:}` saying whether a host control socket is currently + registered. Reaching the object is the point: a client that expects to connect + shortly pays the un-hibernation cost up front. It is inert — no storage, no + alarm, no attachment migration, and no signal to the host — so it can never + perturb a live tunnel. `/health` is answered by the Worker and never reaches a + DO, so it cannot be used for this. + +## Placement + +A Durable Object lives wherever the request that created it landed, for the life +of the `machineKey`. The Worker therefore derives a `locationHint` from +`request.cf` (continent, split on longitude for the two-region continents) and +passes it on `claim` and `/host` only — the machine's own routes. A phone's +location must never decide where a machine's object lives, so `/connect`, +`/prewarm`, and `/pipe` never carry a hint. Cloudflare applies a hint only when +it creates the object, so this steers new machines and leaves existing ones +exactly where they are. ## Trust model — read this diff --git a/apps/tunnel-relay/src/relay.ts b/apps/tunnel-relay/src/relay.ts index 767031e27..21da2deb6 100644 --- a/apps/tunnel-relay/src/relay.ts +++ b/apps/tunnel-relay/src/relay.ts @@ -127,7 +127,8 @@ export type TunnelRoute = | { kind: "claim"; machineKey: string } | { kind: "host"; machineKey: string } | { kind: "pipe"; machineKey: string; id: string } - | { kind: "connect"; machineKey: string }; + | { kind: "connect"; machineKey: string } + | { kind: "prewarm"; machineKey: string }; /** * Pure path router. Returns the matched tunnel route or null; the machineKey @@ -155,6 +156,9 @@ export function routeTunnelPath(pathname: string): TunnelRoute | null { if (parts.length === 2 && parts[0] === "connect") { return validKey(parts[1]) ? { kind: "connect", machineKey: parts[1] } : null; } + if (parts.length === 2 && parts[0] === "prewarm") { + return validKey(parts[1]) ? { kind: "prewarm", machineKey: parts[1] } : null; + } return null; } @@ -162,6 +166,67 @@ function validKey(value: string | undefined): value is string { return typeof value === "string" && MACHINE_KEY_PATTERN.test(value); } +/** The regions Cloudflare accepts as a Durable Object placement hint. */ +export type RelayLocationHint = + | "wnam" + | "enam" + | "sam" + | "weur" + | "eeur" + | "apac" + | "oc" + | "afr" + | "me"; + +/** The subset of `request.cf` this router reads. Absent under `wrangler dev`. */ +export type RequestPlacementGeo = { + continent?: string | null; + longitude?: string | number | null; +}; + +function parseLongitude(raw: unknown): number | null { + const value = typeof raw === "number" ? raw : typeof raw === "string" ? Number(raw.trim()) : Number.NaN; + return Number.isFinite(value) && value >= -180 && value <= 180 ? value : null; +} + +/** + * Maps the requester's coarse geography onto a Durable Object location hint. + * Only the machine's own routes (claim/host) supply one: a Durable Object is + * placed near whichever request creates it, so without a hint a phone that + * connects before the brain claims would pin the object next to the phone for + * the life of the machine key. Cloudflare ignores the hint for an object that + * already exists, so this only steers brand-new machines. + */ +export function locationHintForGeo( + geo: RequestPlacementGeo | null | undefined, +): RelayLocationHint | undefined { + const continent = typeof geo?.continent === "string" ? geo.continent.trim().toUpperCase() : ""; + const longitude = parseLongitude(geo?.longitude); + switch (continent) { + // Cloudflare's regions split North America and Europe in two, so fall back + // to the continent's larger half when the longitude is missing. + case "NA": + return longitude != null && longitude <= -100 ? "wnam" : "enam"; + case "SA": + return "sam"; + case "EU": + return longitude != null && longitude >= 20 ? "eeur" : "weur"; + case "AF": + return "afr"; + case "OC": + return "oc"; + case "AS": + return longitude != null && longitude < 65 ? "me" : "apac"; + default: + // Unknown or absent geography: let Cloudflare place the object itself. + return undefined; + } +} + +function placementHintFor(request: Request): RelayLocationHint | undefined { + return locationHintForGeo((request as { cf?: RequestPlacementGeo | null }).cf); +} + export async function handleRequest(request: Request, env: TunnelRelayEnv): Promise { const url = new URL(request.url); if (url.pathname === "/health") { @@ -180,6 +245,12 @@ export async function handleRequest(request: Request, env: TunnelRelayEnv): Prom // Every route is scoped to a single machine → a single Durable Object // instance. The DO owns the claim secret, signature verification, and the // per-connection socket pairing; the worker is a thin, stateless router. - const stub = env.TUNNEL.get(env.TUNNEL.idFromName(route.machineKey)); + const id = env.TUNNEL.idFromName(route.machineKey); + // Only the machine's own routes may influence placement — a travelling phone + // must never drag a machine's object across the planet. + const locationHint = route.kind === "claim" || route.kind === "host" + ? placementHintFor(request) + : undefined; + const stub = locationHint ? env.TUNNEL.get(id, { locationHint }) : env.TUNNEL.get(id); return stub.fetch(request); } diff --git a/apps/tunnel-relay/src/tunnelDo.ts b/apps/tunnel-relay/src/tunnelDo.ts index 5c0a4aa7d..2aff15330 100644 --- a/apps/tunnel-relay/src/tunnelDo.ts +++ b/apps/tunnel-relay/src/tunnelDo.ts @@ -117,17 +117,36 @@ export class TunnelDurableObject implements DurableObject { { frames: (string | ArrayBuffer)[]; bytes: number } >(); private readonly terminalSocketLogs = new WeakSet(); + // Hibernation drops instance memory, so both caches below are pure overhead + // removal for the hot path: every read falls back to the durable attachment + // (or a tag scan) and repopulates itself after a wake. + private readonly attachments = new WeakMap(); + private readonly partners = new WeakMap(); constructor( private readonly state: DurableObjectState, private readonly env: TunnelRelayEnv, ) {} + /** + * The only writer of a socket attachment, so the instance cache can never + * drift from the durable copy that survives hibernation. + */ + private writeAttachment(ws: WebSocket, attachment: SocketAttachment): void { + ws.serializeAttachment(attachment); + this.attachments.set(ws, attachment); + } + /** One-time deploy migration for hibernated pre-epoch socket attachments. */ private normalizedAttachment(ws: WebSocket): SocketAttachment | null { + const cached = this.attachments.get(ws); + if (cached) return cached; const attachment = ws.deserializeAttachment() as SocketAttachment | null; if (!attachment?.role) return null; - if (attachment.epoch) return attachment; + if (attachment.epoch) { + this.attachments.set(ws, attachment); + return attachment; + } const migrated = { ...attachment, epoch: LEGACY_CONTROL_EPOCH, @@ -135,10 +154,18 @@ export class TunnelDurableObject implements DurableObject { ? { legacyStateUnknown: true as const } : {}), } satisfies SocketAttachment; - ws.serializeAttachment(migrated); + this.writeAttachment(ws, migrated); return migrated; } + /** Drops both caches for a socket that has reached a terminal state. */ + private forgetSocket(ws: WebSocket): void { + const partner = this.partners.get(ws); + if (partner) this.partners.delete(partner); + this.partners.delete(ws); + this.attachments.delete(ws); + } + private epochOf(attachment: SocketAttachment | null): string | null { return attachment?.epoch ?? (attachment?.role ? LEGACY_CONTROL_EPOCH : null); } @@ -152,9 +179,34 @@ export class TunnelDurableObject implements DurableObject { if (route.kind === "host") return this.handleHost(request, url, route.machineKey); if (route.kind === "pipe") return this.handlePipe(request, url, route.machineKey, route.id); if (route.kind === "connect") return this.handleConnect(request, url); + if (route.kind === "prewarm") return this.handlePrewarm(request); return new Response("not found", { status: 404 }); } + /** + * Cheapest possible probe: simply reaching this object is what un-hibernates + * it, so a client that expects to connect shortly can pay the wake cost up + * front. Deliberately read-only — no attachment migration, no storage, no + * alarm, no signal to the host — so a prewarm can never perturb a live + * tunnel. Carries the same authorization stance as `/connect`: the + * machineKey is the secret, and nothing here is actionable without it. + */ + private handlePrewarm(request: Request): Response { + if (request.method !== "GET" && request.method !== "HEAD") { + return new Response("method not allowed", { status: 405 }); + } + let control = false; + for (const ws of this.state.getWebSockets("control")) { + if (!this.isOpen(ws)) continue; + const attachment = ws.deserializeAttachment() as SocketAttachment | null; + if (attachment?.role === "control") { + control = true; + break; + } + } + return jsonResponse({ ok: true, control }, { headers: { "cache-control": "no-store" } }); + } + private async handleClaim(request: Request): Promise { if (request.method !== "POST") return new Response("method not allowed", { status: 405 }); let secret = ""; @@ -441,7 +493,7 @@ export class TunnelDurableObject implements DurableObject { if (attachment.id) tags.push(`conn:${attachment.id}`); tags.push(`epoch:${epoch}`); this.state.acceptWebSocket(server, tags); - server.serializeAttachment({ ...attachment, epoch, ts: Date.now() } satisfies SocketAttachment); + this.writeAttachment(server, { ...attachment, epoch, ts: Date.now() } satisfies SocketAttachment); if (afterAccept) afterAccept(server); return new Response(null, { status: 101, webSocket: client }); } @@ -486,6 +538,11 @@ export class TunnelDurableObject implements DurableObject { ? "client" : null; if (!partnerRole) return null; + // The cache is only ever a hint: it is re-checked against the same + // predicate the scan uses, so a stale entry degrades to the scan below. + const cached = this.partners.get(ws); + if (cached && cached !== ws && this.isPartnerOf(cached, att, partnerRole)) return cached; + if (cached) this.partners.delete(cached); for (const candidate of this.state.getWebSockets(`conn:${att.id}`)) { if (candidate === ws || !this.isOpen(candidate)) continue; const candidateAttachment = this.normalizedAttachment(candidate); @@ -493,12 +550,23 @@ export class TunnelDurableObject implements DurableObject { candidateAttachment?.role === partnerRole && this.epochOf(candidateAttachment) === this.epochOf(att) ) { + this.partners.set(ws, candidate); + this.partners.set(candidate, ws); return candidate; } } + this.partners.delete(ws); return null; } + private isPartnerOf(candidate: WebSocket, att: SocketAttachment, partnerRole: SocketRole): boolean { + if (!this.isOpen(candidate)) return false; + const candidateAttachment = this.normalizedAttachment(candidate); + return candidateAttachment?.role === partnerRole + && candidateAttachment.id === att.id + && this.epochOf(candidateAttachment) === this.epochOf(att); + } + private controlMessageMatchesEpoch(messageEpoch: unknown, attachment: SocketAttachment): boolean { const epoch = this.epochOf(attachment); return epoch === LEGACY_CONTROL_EPOCH @@ -532,8 +600,8 @@ export class TunnelDurableObject implements DurableObject { ...cleanPartnerAttachment } = partnerAttachment; const migratedSource = { ...sourceAttachment, established: true, ts: now } satisfies SocketAttachment; - ws.serializeAttachment(migratedSource); - partner.serializeAttachment({ + this.writeAttachment(ws, migratedSource); + this.writeAttachment(partner, { ...cleanPartnerAttachment, established: true, ts: now, @@ -661,7 +729,7 @@ export class TunnelDurableObject implements DurableObject { // buffered data exists, never the ADE frame/token itself. A reconstructed // instance can then fail loudly instead of silently losing the hello. att = { ...att, legacyBuffered: true, ts: Date.now() }; - ws.serializeAttachment(att satisfies SocketAttachment); + this.writeAttachment(ws, att satisfies SocketAttachment); } } else if (att.role === "pipe") { this.closePair(ws, att, CLOSE_NOT_READY, "relay bridge not ready"); @@ -673,17 +741,19 @@ export class TunnelDurableObject implements DurableObject { private touch(ws: WebSocket, att: SocketAttachment): void { const now = Date.now(); if (now - att.ts < ACTIVITY_WRITE_THROTTLE_MS) return; - ws.serializeAttachment({ ...att, ts: now } satisfies SocketAttachment); + this.writeAttachment(ws, { ...att, ts: now } satisfies SocketAttachment); } async webSocketClose(ws: WebSocket, code: number, reason: string, _wasClean: boolean): Promise { this.logSocketTerminal(ws, "socket_closed", code); this.teardownPartner(ws, code, reason); + this.forgetSocket(ws); } async webSocketError(ws: WebSocket, _error: unknown): Promise { this.logSocketTerminal(ws, "socket_error", CLOSE_PARTNER_CLOSED); this.teardownPartner(ws, CLOSE_PARTNER_CLOSED, "partner error"); + this.forgetSocket(ws); } private logSocketTerminal(ws: WebSocket, kind: "socket_closed" | "socket_error", code: number): void { @@ -766,8 +836,8 @@ export class TunnelDurableObject implements DurableObject { legacyStateUnknown: _legacyStateUnknown, ...cleanClientAttachment } = clientAttachment; - client.serializeAttachment({ ...cleanClientAttachment, established: true, ts: now } satisfies SocketAttachment); - pipe.serializeAttachment({ ...pipeAttachment, established: true, ts: now } satisfies SocketAttachment); + this.writeAttachment(client, { ...cleanClientAttachment, established: true, ts: now } satisfies SocketAttachment); + this.writeAttachment(pipe, { ...pipeAttachment, established: true, ts: now } satisfies SocketAttachment); } private closePair( diff --git a/apps/tunnel-relay/test/relay.test.ts b/apps/tunnel-relay/test/relay.test.ts index cef633e43..3a5aa5bb3 100644 --- a/apps/tunnel-relay/test/relay.test.ts +++ b/apps/tunnel-relay/test/relay.test.ts @@ -6,6 +6,7 @@ import { generateConnectionId, handleRequest, hmacSha256Hex, + locationHintForGeo, routeTunnelPath, verifySignedQuery, type TunnelRelayEnv, @@ -201,16 +202,104 @@ describe("routeTunnelPath", () => { id: "abcdef01", }); expect(routeTunnelPath(`/connect/${MACHINE_KEY}`)).toEqual({ kind: "connect", machineKey: MACHINE_KEY }); + expect(routeTunnelPath(`/prewarm/${MACHINE_KEY}`)).toEqual({ kind: "prewarm", machineKey: MACHINE_KEY }); }); it("rejects bad machine keys and connection ids", () => { expect(routeTunnelPath("/host/not-hex")).toBeNull(); expect(routeTunnelPath(`/host/${MACHINE_KEY}/pipe/NOThex!`)).toBeNull(); expect(routeTunnelPath(`/machines/${MACHINE_KEY}/publish`)).toBeNull(); + expect(routeTunnelPath("/prewarm/not-hex")).toBeNull(); + expect(routeTunnelPath(`/prewarm/${MACHINE_KEY}/extra`)).toBeNull(); expect(routeTunnelPath("/unknown")).toBeNull(); }); }); +describe("locationHintForGeo", () => { + it("splits the two-region continents on longitude", () => { + expect(locationHintForGeo({ continent: "NA", longitude: "-122.33" })).toBe("wnam"); // Seattle + expect(locationHintForGeo({ continent: "NA", longitude: "-73.94" })).toBe("enam"); // New York + expect(locationHintForGeo({ continent: "EU", longitude: "-0.13" })).toBe("weur"); // London + expect(locationHintForGeo({ continent: "EU", longitude: "30.52" })).toBe("eeur"); // Kyiv + expect(locationHintForGeo({ continent: "AS", longitude: "55.27" })).toBe("me"); // Dubai + expect(locationHintForGeo({ continent: "AS", longitude: "139.69" })).toBe("apac"); // Tokyo + }); + + it("maps the single-region continents", () => { + expect(locationHintForGeo({ continent: "sa", longitude: "-46.63" })).toBe("sam"); + expect(locationHintForGeo({ continent: "AF", longitude: "18.42" })).toBe("afr"); + expect(locationHintForGeo({ continent: "OC", longitude: "151.21" })).toBe("oc"); + }); + + it("falls back to the larger half when the longitude is missing or junk", () => { + expect(locationHintForGeo({ continent: "NA" })).toBe("enam"); + expect(locationHintForGeo({ continent: "EU", longitude: "not-a-number" })).toBe("weur"); + expect(locationHintForGeo({ continent: "AS", longitude: 999 })).toBe("apac"); + }); + + it("declines to hint when the geography is unknown", () => { + // `request.cf` is absent under `wrangler dev` and for unrecognized regions; + // Cloudflare's own placement is the right answer there. + expect(locationHintForGeo(undefined)).toBeUndefined(); + expect(locationHintForGeo(null)).toBeUndefined(); + expect(locationHintForGeo({})).toBeUndefined(); + expect(locationHintForGeo({ continent: "XX", longitude: "10" })).toBeUndefined(); + }); +}); + +describe("Durable Object placement", () => { + function routerHarness(): { env: TunnelRelayEnv; hints: Array } { + const hints: Array = []; + const stub = { fetch: async () => new Response("routed") } as unknown as DurableObjectStub; + const env = { + TUNNEL: { + idFromName: (name: string) => name as unknown as DurableObjectId, + get: (_id: DurableObjectId, options?: { locationHint?: string }) => { + hints.push(options?.locationHint); + return stub; + }, + } as unknown as DurableObjectNamespace, + } as TunnelRelayEnv; + return { env, hints }; + } + + function geoRequest(path: string, cf?: Record): Request { + const request = new Request(`https://relay.test${path}`); + return (cf ? Object.assign(request, { cf }) : request) as Request; + } + + it("hints placement from the machine's own routes", async () => { + const { env, hints } = routerHarness(); + const cf = { continent: "NA", longitude: "-122.33" }; + + await handleRequest(geoRequest(`/machines/${MACHINE_KEY}/claim`, cf), env); + await handleRequest(geoRequest(`/host/${MACHINE_KEY}`, cf), env); + + expect(hints).toEqual(["wnam", "wnam"]); + }); + + it("never lets a phone's location place a machine's object", async () => { + const { env, hints } = routerHarness(); + // A travelling phone must not drag the object away from the machine, and a + // pipe/prewarm arrives long after placement is already settled. + const cf = { continent: "AS", longitude: "139.69" }; + + await handleRequest(geoRequest(`/connect/${MACHINE_KEY}`, cf), env); + await handleRequest(geoRequest(`/prewarm/${MACHINE_KEY}`, cf), env); + await handleRequest(geoRequest(`/host/${MACHINE_KEY}/pipe/abcdef01`, cf), env); + + expect(hints).toEqual([undefined, undefined, undefined]); + }); + + it("omits the hint entirely when the request carries no geography", async () => { + const { env, hints } = routerHarness(); + + await handleRequest(geoRequest(`/machines/${MACHINE_KEY}/claim`), env); + + expect(hints).toEqual([undefined]); + }); +}); + describe("verifySignedQuery", () => { it("accepts a correct signature within skew", async () => { const ts = String(Math.floor(Date.now() / 1000)); diff --git a/apps/tunnel-relay/test/relay.workerd.test.ts b/apps/tunnel-relay/test/relay.workerd.test.ts index 7e5461f73..ef77d9ed7 100644 --- a/apps/tunnel-relay/test/relay.workerd.test.ts +++ b/apps/tunnel-relay/test/relay.workerd.test.ts @@ -1,6 +1,7 @@ import { env } from "cloudflare:workers"; import { evictDurableObject, runInDurableObject } from "cloudflare:test"; import { afterEach, describe, expect, it } from "vitest"; +import worker from "../src"; import { buildHostSignatureBase, buildPipeSignatureBase, @@ -9,6 +10,8 @@ import { import { CLOSE_BRIDGE_REJECTED, CLOSE_CLIENT_GONE, + CLOSE_FORWARD_FAILED, + CLOSE_IDLE, LEGACY_CONTROL_EPOCH, RELAY_READY_VERSION, TunnelDurableObject, @@ -176,6 +179,29 @@ async function establishEpochV2(keyDigit: string): Promise<{ return { stub, control, client, pipe, id: openMessage.id }; } +/** Opens one more ready-v2 tunnel over an already-registered control socket. */ +async function openTunnel(args: { + stub: DurableObjectStub; + key: string; + control: WebSocket; +}): Promise<{ client: WebSocket; pipe: WebSocket; id: string }> { + const open = nextMessage(args.control); + const client = await upgrade(args.stub, `https://relay.test/connect/${args.key}?ready=${RELAY_READY_VERSION}`); + expect(await nextMessage(client)).toBe(JSON.stringify({ t: "accepted", v: RELAY_READY_VERSION })); + const { id } = JSON.parse(String(await open)) as { id: string }; + const pipe = await openPipe({ stub: args.stub, key: args.key, id, epoch: CONTROL_EPOCH }); + const ready = nextMessage(client); + args.control.send(JSON.stringify({ t: "ready", id, epoch: CONTROL_EPOCH })); + expect(await ready).toBe(JSON.stringify({ t: "ready", v: RELAY_READY_VERSION })); + return { client, pipe, id }; +} + +async function prewarm(stub: DurableObjectStub, key: string): Promise { + const response = await stub.fetch(`https://relay.test/prewarm/${key}`); + expect(response.status).toBe(200); + return response.json(); +} + describe("TunnelDurableObject in workerd", () => { it("preserves actual attachments and ordered routing for an established epoch-v2 triple", async () => { const { stub, client, pipe, id } = await establishEpochV2("a"); @@ -243,6 +269,10 @@ describe("TunnelDurableObject in workerd", () => { expect(serverControl).toBeDefined(); serverControl!.serializeAttachment({ role: "control", ts: Date.now() }); }); + // A pre-epoch attachment only ever reaches this code the way it does in + // production: the previous Worker accepted the socket, the object + // hibernated, and the deployed code wakes to an attachment it never wrote. + await evictDurableObject(stub, { webSockets: "hibernate" }); const open = nextMessage(control); const client = await upgrade(stub, `https://relay.test/connect/${key}?ready=${RELAY_READY_VERSION}`); @@ -322,4 +352,120 @@ describe("TunnelDurableObject in workerd", () => { expect(close.code).toBe(CLOSE_BRIDGE_REJECTED); expect(close.reason).toBe("relay client unavailable"); }); + + it("keeps concurrent tunnels on one machine independently paired across hibernation", async () => { + const key = machineKey("1"); + const stub = stubFor(key); + await claim(stub, key); + const control = await openControl({ stub, key, epoch: CONTROL_EPOCH }); + + const first = await openTunnel({ stub, key, control }); + const second = await openTunnel({ stub, key, control }); + expect(first.id).not.toBe(second.id); + + const atFirstPipe = nextMessage(first.pipe); + const atSecondPipe = nextMessage(second.pipe); + first.client.send("to-first"); + second.client.send("to-second"); + expect(await atFirstPipe).toBe("to-first"); + expect(await atSecondPipe).toBe("to-second"); + + // Hibernation wipes the instance's pairing cache; each side has to rebuild + // its own partner from the durable attachments, not inherit its neighbour's. + await evictDurableObject(stub, { webSockets: "hibernate" }); + + const atFirstClient = nextMessage(first.client); + const atSecondClient = nextMessage(second.client); + first.pipe.send("from-first"); + second.pipe.send("from-second"); + expect(await atFirstClient).toBe("from-first"); + expect(await atSecondClient).toBe("from-second"); + + // Tearing one pair down must not disturb the other pair's cached partner. + const firstClientClosed = nextClose(first.client); + first.pipe.close(4000, "done"); + await firstClientClosed; + + const stillRouting = nextMessage(second.pipe); + second.client.send("second-survives"); + expect(await stillRouting).toBe("second-survives"); + }); + + it("reports an unavailable partner rather than forwarding into a closing socket", async () => { + const { stub, client, pipe, id } = await establishEpochV2("5"); + + // Forward one frame first so the pair is cached in instance memory. + const warmed = nextMessage(pipe); + client.send("warm-up"); + expect(await warmed).toBe("warm-up"); + + const clientClosed = nextClose(client); + // Close the pipe and deliver the next frame in the same synchronous turn, + // before the close handler can drop the pairing. A cached partner must be + // re-checked on every read, or this frame is forwarded into a dead socket. + await runInDurableObject(stub, async (instance, state) => { + const paired = state.getWebSockets(`conn:${id}`); + const serverClient = paired.find((s) => (s.deserializeAttachment() as Attachment).role === "client"); + const serverPipe = paired.find((s) => (s.deserializeAttachment() as Attachment).role === "pipe"); + expect(serverClient).toBeDefined(); + expect(serverPipe).toBeDefined(); + serverPipe!.close(CLOSE_IDLE, "idle timeout"); + await instance.webSocketMessage(serverClient!, "into-the-void"); + }); + + const close = await clientClosed; + expect(close.code).toBe(CLOSE_FORWARD_FAILED); + expect(close.reason).toBe("relay partner unavailable"); + expect(pipe).toBeDefined(); + }); + + it("prewarms a hibernated object without disturbing its live tunnel", async () => { + const key = machineKey("2"); + const stub = stubFor(key); + await claim(stub, key); + const control = await openControl({ stub, key, epoch: CONTROL_EPOCH }); + const { client, pipe, id } = await openTunnel({ stub, key, control }); + + await evictDurableObject(stub, { webSockets: "hibernate" }); + expect(await prewarm(stub, key)).toEqual({ ok: true, control: true }); + + // The probe is inert: the pair is still established and still routes. + expect(await attachments(stub)).toEqual(expect.arrayContaining([ + expect.objectContaining({ role: "client", id, epoch: CONTROL_EPOCH, established: true }), + expect.objectContaining({ role: "pipe", id, epoch: CONTROL_EPOCH, established: true }), + ])); + const routed = nextMessage(pipe); + client.send("after-prewarm"); + expect(await routed).toBe("after-prewarm"); + }); + + it("reports whether a host control socket is registered, and rejects writes", async () => { + const key = machineKey("3"); + const stub = stubFor(key); + await claim(stub, key); + expect(await prewarm(stub, key)).toEqual({ ok: true, control: false }); + + await openControl({ stub, key, epoch: CONTROL_EPOCH }); + expect(await prewarm(stub, key)).toEqual({ ok: true, control: true }); + + await runInDurableObject(stub, (_instance, state) => { + for (const socket of state.getWebSockets("control")) socket.close(1000, "host stopped"); + }); + expect(await prewarm(stub, key)).toEqual({ ok: true, control: false }); + + const posted = await stub.fetch(`https://relay.test/prewarm/${key}`, { method: "POST" }); + expect(posted.status).toBe(405); + }); + + it("routes prewarm for an unknown machine without creating any state", async () => { + const key = machineKey("4"); + const response = await worker.fetch(new Request(`https://relay.test/prewarm/${key}`), env); + + expect(response.status).toBe(200); + expect(response.headers.get("cache-control")).toBe("no-store"); + expect(await response.json()).toEqual({ ok: true, control: false }); + + // Nothing was persisted, so the machine key is still free to be claimed. + await claim(stubFor(key), key); + }); }); From 7e521258a7701364294800e665992eca1346c962 Mon Sep 17 00:00:00 2001 From: Arul Sharma <31745423+arul28@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:11:31 -0400 Subject: [PATCH 2/3] refactor(relay): fold the partner predicate into one place MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Applied /quality dual-review findings. All behavior-preserving: - partnerOf: the cache-hit path and the tag scan now share isPartnerOf, so the "re-checked against the same predicate" invariant is structural rather than coincidental. Drops two guards verified dead — a self-entry cannot be inserted, and a stale reverse entry self-heals on its own next read. - handlePrewarm: expressed as .some(), still reading the raw attachment so it cannot trigger a migration write. - RelayLocationHint aliases Cloudflare's DurableObjectLocationHint instead of re-enumerating it, so the union cannot drift. - parseLongitude: an empty longitude is missing data, not the prime meridian. Number("") is 0, which read as a real coordinate and picked "me" for Asia. - Renamed the caches cachedAttachments/cachedPartners: the socket's own attachment stays authoritative, and the names now say so. - Tests: establishEpochV2 delegates to openTunnel, so the {t:"open"} shape is now asserted for every tunnel opened. Adds the empty-longitude case and one test pinning the socket-identity premise the caches rest on. Co-Authored-By: Claude Fable 5 --- apps/tunnel-relay/src/relay.ts | 29 +++----- apps/tunnel-relay/src/tunnelDo.ts | 61 +++++++--------- apps/tunnel-relay/test/relay.test.ts | 4 ++ apps/tunnel-relay/test/relay.workerd.test.ts | 76 ++++++++++---------- 4 files changed, 75 insertions(+), 95 deletions(-) diff --git a/apps/tunnel-relay/src/relay.ts b/apps/tunnel-relay/src/relay.ts index 21da2deb6..26d4d5a3e 100644 --- a/apps/tunnel-relay/src/relay.ts +++ b/apps/tunnel-relay/src/relay.ts @@ -167,24 +167,20 @@ function validKey(value: string | undefined): value is string { } /** The regions Cloudflare accepts as a Durable Object placement hint. */ -export type RelayLocationHint = - | "wnam" - | "enam" - | "sam" - | "weur" - | "eeur" - | "apac" - | "oc" - | "afr" - | "me"; +export type RelayLocationHint = DurableObjectLocationHint; -/** The subset of `request.cf` this router reads. Absent under `wrangler dev`. */ +/** + * The subset of `request.cf` this router reads, kept looser than Cloudflare's + * own types: `wrangler dev` omits `cf` entirely and miniflare will hand back + * whatever a test supplies. + */ export type RequestPlacementGeo = { continent?: string | null; longitude?: string | number | null; }; function parseLongitude(raw: unknown): number | null { + if (typeof raw === "string" && raw.trim() === "") return null; const value = typeof raw === "number" ? raw : typeof raw === "string" ? Number(raw.trim()) : Number.NaN; return Number.isFinite(value) && value >= -180 && value <= 180 ? value : null; } @@ -202,9 +198,10 @@ export function locationHintForGeo( ): RelayLocationHint | undefined { const continent = typeof geo?.continent === "string" ? geo.continent.trim().toUpperCase() : ""; const longitude = parseLongitude(geo?.longitude); + // Three continents span more than one Cloudflare region: the Americas and + // Europe split east/west, and western Asia is served by `me`. Each falls back + // to its most populous region when the longitude is missing. switch (continent) { - // Cloudflare's regions split North America and Europe in two, so fall back - // to the continent's larger half when the longitude is missing. case "NA": return longitude != null && longitude <= -100 ? "wnam" : "enam"; case "SA": @@ -223,10 +220,6 @@ export function locationHintForGeo( } } -function placementHintFor(request: Request): RelayLocationHint | undefined { - return locationHintForGeo((request as { cf?: RequestPlacementGeo | null }).cf); -} - export async function handleRequest(request: Request, env: TunnelRelayEnv): Promise { const url = new URL(request.url); if (url.pathname === "/health") { @@ -249,7 +242,7 @@ export async function handleRequest(request: Request, env: TunnelRelayEnv): Prom // Only the machine's own routes may influence placement — a travelling phone // must never drag a machine's object across the planet. const locationHint = route.kind === "claim" || route.kind === "host" - ? placementHintFor(request) + ? locationHintForGeo((request as { cf?: RequestPlacementGeo | null }).cf) : undefined; const stub = locationHint ? env.TUNNEL.get(id, { locationHint }) : env.TUNNEL.get(id); return stub.fetch(request); diff --git a/apps/tunnel-relay/src/tunnelDo.ts b/apps/tunnel-relay/src/tunnelDo.ts index 2aff15330..d906d1d7f 100644 --- a/apps/tunnel-relay/src/tunnelDo.ts +++ b/apps/tunnel-relay/src/tunnelDo.ts @@ -117,11 +117,11 @@ export class TunnelDurableObject implements DurableObject { { frames: (string | ArrayBuffer)[]; bytes: number } >(); private readonly terminalSocketLogs = new WeakSet(); - // Hibernation drops instance memory, so both caches below are pure overhead - // removal for the hot path: every read falls back to the durable attachment - // (or a tag scan) and repopulates itself after a wake. - private readonly attachments = new WeakMap(); - private readonly partners = new WeakMap(); + // Pure overhead removal for the hot path; the socket's own attachment stays + // authoritative. Hibernation drops instance memory, so every read falls back + // to the durable copy (or a tag scan) and repopulates itself after a wake. + private readonly cachedAttachments = new WeakMap(); + private readonly cachedPartners = new WeakMap(); constructor( private readonly state: DurableObjectState, @@ -134,17 +134,17 @@ export class TunnelDurableObject implements DurableObject { */ private writeAttachment(ws: WebSocket, attachment: SocketAttachment): void { ws.serializeAttachment(attachment); - this.attachments.set(ws, attachment); + this.cachedAttachments.set(ws, attachment); } /** One-time deploy migration for hibernated pre-epoch socket attachments. */ private normalizedAttachment(ws: WebSocket): SocketAttachment | null { - const cached = this.attachments.get(ws); + const cached = this.cachedAttachments.get(ws); if (cached) return cached; const attachment = ws.deserializeAttachment() as SocketAttachment | null; if (!attachment?.role) return null; if (attachment.epoch) { - this.attachments.set(ws, attachment); + this.cachedAttachments.set(ws, attachment); return attachment; } const migrated = { @@ -160,10 +160,10 @@ export class TunnelDurableObject implements DurableObject { /** Drops both caches for a socket that has reached a terminal state. */ private forgetSocket(ws: WebSocket): void { - const partner = this.partners.get(ws); - if (partner) this.partners.delete(partner); - this.partners.delete(ws); - this.attachments.delete(ws); + const partner = this.cachedPartners.get(ws); + if (partner) this.cachedPartners.delete(partner); + this.cachedPartners.delete(ws); + this.cachedAttachments.delete(ws); } private epochOf(attachment: SocketAttachment | null): string | null { @@ -195,15 +195,9 @@ export class TunnelDurableObject implements DurableObject { if (request.method !== "GET" && request.method !== "HEAD") { return new Response("method not allowed", { status: 405 }); } - let control = false; - for (const ws of this.state.getWebSockets("control")) { - if (!this.isOpen(ws)) continue; - const attachment = ws.deserializeAttachment() as SocketAttachment | null; - if (attachment?.role === "control") { - control = true; - break; - } - } + const control = this.state.getWebSockets("control").some((ws) => ( + this.isOpen(ws) && (ws.deserializeAttachment() as SocketAttachment | null)?.role === "control" + )); return jsonResponse({ ok: true, control }, { headers: { "cache-control": "no-store" } }); } @@ -538,24 +532,17 @@ export class TunnelDurableObject implements DurableObject { ? "client" : null; if (!partnerRole) return null; - // The cache is only ever a hint: it is re-checked against the same - // predicate the scan uses, so a stale entry degrades to the scan below. - const cached = this.partners.get(ws); - if (cached && cached !== ws && this.isPartnerOf(cached, att, partnerRole)) return cached; - if (cached) this.partners.delete(cached); + // The cache is only ever a hint, re-checked against the very predicate the + // scan uses, so a socket that closed since it was cached falls through. + const cached = this.cachedPartners.get(ws); + if (cached && this.isPartnerOf(cached, att, partnerRole)) return cached; for (const candidate of this.state.getWebSockets(`conn:${att.id}`)) { - if (candidate === ws || !this.isOpen(candidate)) continue; - const candidateAttachment = this.normalizedAttachment(candidate); - if ( - candidateAttachment?.role === partnerRole - && this.epochOf(candidateAttachment) === this.epochOf(att) - ) { - this.partners.set(ws, candidate); - this.partners.set(candidate, ws); - return candidate; - } + if (candidate === ws || !this.isPartnerOf(candidate, att, partnerRole)) continue; + this.cachedPartners.set(ws, candidate); + this.cachedPartners.set(candidate, ws); + return candidate; } - this.partners.delete(ws); + this.cachedPartners.delete(ws); return null; } diff --git a/apps/tunnel-relay/test/relay.test.ts b/apps/tunnel-relay/test/relay.test.ts index 3a5aa5bb3..34e796361 100644 --- a/apps/tunnel-relay/test/relay.test.ts +++ b/apps/tunnel-relay/test/relay.test.ts @@ -235,6 +235,10 @@ describe("locationHintForGeo", () => { expect(locationHintForGeo({ continent: "NA" })).toBe("enam"); expect(locationHintForGeo({ continent: "EU", longitude: "not-a-number" })).toBe("weur"); expect(locationHintForGeo({ continent: "AS", longitude: 999 })).toBe("apac"); + // An empty longitude is missing data, not the prime meridian — `Number("")` + // is 0, which would otherwise read as a real coordinate and pick `me`. + expect(locationHintForGeo({ continent: "AS", longitude: "" })).toBe("apac"); + expect(locationHintForGeo({ continent: "AS", longitude: " " })).toBe("apac"); }); it("declines to hint when the geography is unknown", () => { diff --git a/apps/tunnel-relay/test/relay.workerd.test.ts b/apps/tunnel-relay/test/relay.workerd.test.ts index ef77d9ed7..5c75b947c 100644 --- a/apps/tunnel-relay/test/relay.workerd.test.ts +++ b/apps/tunnel-relay/test/relay.workerd.test.ts @@ -154,46 +154,39 @@ async function attachments(stub: DurableObjectStub): Promise { )); } -async function establishEpochV2(keyDigit: string): Promise<{ +/** Opens one ready-v2 tunnel over an already-registered control socket. */ +async function openTunnel(args: { stub: DurableObjectStub; + key: string; control: WebSocket; - client: WebSocket; - pipe: WebSocket; - id: string; -}> { - const key = machineKey(keyDigit); - const stub = stubFor(key); - await claim(stub, key); - const control = await openControl({ stub, key, epoch: CONTROL_EPOCH }); - - const open = nextMessage(control); - const client = await upgrade(stub, `https://relay.test/connect/${key}?ready=${RELAY_READY_VERSION}`); +}): Promise<{ client: WebSocket; pipe: WebSocket; id: string }> { + const open = nextMessage(args.control); + const client = await upgrade(args.stub, `https://relay.test/connect/${args.key}?ready=${RELAY_READY_VERSION}`); expect(await nextMessage(client)).toBe(JSON.stringify({ t: "accepted", v: RELAY_READY_VERSION })); const openMessage = JSON.parse(String(await open)) as { t: string; id: string; epoch: string; readyVersion: number }; expect(openMessage).toMatchObject({ t: "open", epoch: CONTROL_EPOCH, readyVersion: RELAY_READY_VERSION }); - const pipe = await openPipe({ stub, key, id: openMessage.id, epoch: CONTROL_EPOCH }); + const pipe = await openPipe({ stub: args.stub, key: args.key, id: openMessage.id, epoch: CONTROL_EPOCH }); const ready = nextMessage(client); - control.send(JSON.stringify({ t: "ready", id: openMessage.id, epoch: CONTROL_EPOCH })); + args.control.send(JSON.stringify({ t: "ready", id: openMessage.id, epoch: CONTROL_EPOCH })); expect(await ready).toBe(JSON.stringify({ t: "ready", v: RELAY_READY_VERSION })); - return { stub, control, client, pipe, id: openMessage.id }; + return { client, pipe, id: openMessage.id }; } -/** Opens one more ready-v2 tunnel over an already-registered control socket. */ -async function openTunnel(args: { - stub: DurableObjectStub; +/** A claimed machine with a control socket and one established tunnel. */ +async function establishEpochV2(keyDigit: string): Promise<{ key: string; + stub: DurableObjectStub; control: WebSocket; -}): Promise<{ client: WebSocket; pipe: WebSocket; id: string }> { - const open = nextMessage(args.control); - const client = await upgrade(args.stub, `https://relay.test/connect/${args.key}?ready=${RELAY_READY_VERSION}`); - expect(await nextMessage(client)).toBe(JSON.stringify({ t: "accepted", v: RELAY_READY_VERSION })); - const { id } = JSON.parse(String(await open)) as { id: string }; - const pipe = await openPipe({ stub: args.stub, key: args.key, id, epoch: CONTROL_EPOCH }); - const ready = nextMessage(client); - args.control.send(JSON.stringify({ t: "ready", id, epoch: CONTROL_EPOCH })); - expect(await ready).toBe(JSON.stringify({ t: "ready", v: RELAY_READY_VERSION })); - return { client, pipe, id }; + client: WebSocket; + pipe: WebSocket; + id: string; +}> { + const key = machineKey(keyDigit); + const stub = stubFor(key); + await claim(stub, key); + const control = await openControl({ stub, key, epoch: CONTROL_EPOCH }); + return { key, stub, control, ...(await openTunnel({ stub, key, control })) }; } async function prewarm(stub: DurableObjectStub, key: string): Promise { @@ -354,12 +347,8 @@ describe("TunnelDurableObject in workerd", () => { }); it("keeps concurrent tunnels on one machine independently paired across hibernation", async () => { - const key = machineKey("1"); - const stub = stubFor(key); - await claim(stub, key); - const control = await openControl({ stub, key, epoch: CONTROL_EPOCH }); - - const first = await openTunnel({ stub, key, control }); + const { key, stub, control, ...firstTunnel } = await establishEpochV2("1"); + const first = firstTunnel; const second = await openTunnel({ stub, key, control }); expect(first.id).not.toBe(second.id); @@ -391,6 +380,18 @@ describe("TunnelDurableObject in workerd", () => { expect(await stillRouting).toBe("second-survives"); }); + it("hands back a stable socket object, which is what makes caching possible", async () => { + const { stub, id } = await establishEpochV2("6"); + + // The instance caches attachment and partner state keyed by the socket + // object. If the runtime ever returned a fresh wrapper per lookup, both + // caches would silently never hit and every other test would still pass. + await runInDurableObject(stub, (_instance, state) => { + expect(state.getWebSockets(`conn:${id}`)[0]).toBe(state.getWebSockets(`conn:${id}`)[0]); + expect(state.getWebSockets("control")[0]).toBe(state.getWebSockets("control")[0]); + }); + }); + it("reports an unavailable partner rather than forwarding into a closing socket", async () => { const { stub, client, pipe, id } = await establishEpochV2("5"); @@ -416,15 +417,10 @@ describe("TunnelDurableObject in workerd", () => { const close = await clientClosed; expect(close.code).toBe(CLOSE_FORWARD_FAILED); expect(close.reason).toBe("relay partner unavailable"); - expect(pipe).toBeDefined(); }); it("prewarms a hibernated object without disturbing its live tunnel", async () => { - const key = machineKey("2"); - const stub = stubFor(key); - await claim(stub, key); - const control = await openControl({ stub, key, epoch: CONTROL_EPOCH }); - const { client, pipe, id } = await openTunnel({ stub, key, control }); + const { key, stub, client, pipe, id } = await establishEpochV2("2"); await evictDurableObject(stub, { webSockets: "hibernate" }); expect(await prewarm(stub, key)).toEqual({ ok: true, control: true }); From 2664be2c406b35bcbd0a31ea17ca786d60560426 Mon Sep 17 00:00:00 2001 From: Arul Sharma <31745423+arul28@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:13:28 -0400 Subject: [PATCH 3/3] docs(sync): record relay Durable Object placement Co-Authored-By: Claude Fable 5 --- docs/features/sync-and-multi-device/README.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/features/sync-and-multi-device/README.md b/docs/features/sync-and-multi-device/README.md index d630813c7..32f4956cb 100644 --- a/docs/features/sync-and-multi-device/README.md +++ b/docs/features/sync-and-multi-device/README.md @@ -884,7 +884,12 @@ Canonical files (`apps/ade-cli/src/services/sync/`): without them; there is no stored enablement or user kill-switch. The store derives the controller-facing `wss:///connect/` URL and the canonical host/pipe - HMAC signing strings shared with the `apps/tunnel-relay` worker. + HMAC signing strings shared with the `apps/tunnel-relay` worker. The claim + and host requests also decide where the relay's Durable Object is placed: + the worker derives a location hint from the requesting machine's geography + so the object is created near the machine rather than near whichever request + arrives first. Cloudflare honors a hint only at creation, so an existing + machineKey keeps its original placement. - `syncTunnelClientService.ts` — the brain-side tunnel client. When the machine has a current ADE account lease it keeps an outbound WebSocket registered with the relay worker (HMAC-signed host/pipe upgrades,