Relay worker micro-optimizations: socket-state caching, /prewarm, DO placement

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:<bool>}` 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
 

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,13 +156,70 @@ 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;
 }
 
 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 = DurableObjectLocationHint;
+
+/**
+ * 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;
+}
+
+/**
+ * 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);
+  // 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) {
+    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;
+  }
+}
+
 export async function handleRequest(request: Request, env: TunnelRelayEnv): Promise<Response> {
   const url = new URL(request.url);
   if (url.pathname === "/health") {
@@ -180,6 +238,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"
+    ? locationHintForGeo((request as { cf?: RequestPlacementGeo | null }).cf)
+    : undefined;
+  const stub = locationHint ? env.TUNNEL.get(id, { locationHint }) : env.TUNNEL.get(id);
   return stub.fetch(request);
 }

apps/tunnel-relay/src/tunnelDo.ts

@@ -117,28 +117,55 @@ export class TunnelDurableObject implements DurableObject {
     { frames: (string | ArrayBuffer)[]; bytes: number }
   >();
   private readonly terminalSocketLogs = new WeakSet<WebSocket>();
+  // 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<WebSocket, SocketAttachment>();
+  private readonly cachedPartners = new WeakMap<WebSocket, WebSocket>();
 
   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.cachedAttachments.set(ws, attachment);
+  }
+
   /** One-time deploy migration for hibernated pre-epoch socket attachments. */
   private normalizedAttachment(ws: WebSocket): SocketAttachment | null {
+    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) return attachment;
+    if (attachment.epoch) {
+      this.cachedAttachments.set(ws, attachment);
+      return attachment;
+    }
     const migrated = {
       ...attachment,
       epoch: LEGACY_CONTROL_EPOCH,
       ...(attachment.role !== "control" && !attachment.established
         ? { 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.cachedPartners.get(ws);
+    if (partner) this.cachedPartners.delete(partner);
+    this.cachedPartners.delete(ws);
+    this.cachedAttachments.delete(ws);
+  }
+
   private epochOf(attachment: SocketAttachment | null): string | null {
     return attachment?.epoch ?? (attachment?.role ? LEGACY_CONTROL_EPOCH : null);
   }
@@ -152,9 +179,28 @@ 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 });
+    }
+    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" } });
+  }
+
   private async handleClaim(request: Request): Promise<Response> {
     if (request.method !== "POST") return new Response("method not allowed", { status: 405 });
     let secret = "";
@@ -441,7 +487,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,19 +532,28 @@ export class TunnelDurableObject implements DurableObject {
         ? "client"
         : null;
     if (!partnerRole) return null;
+    // 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)
-      ) {
-        return candidate;
-      }
+      if (candidate === ws || !this.isPartnerOf(candidate, att, partnerRole)) continue;
+      this.cachedPartners.set(ws, candidate);
+      this.cachedPartners.set(candidate, ws);
+      return candidate;
     }
+    this.cachedPartners.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 +587,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 +716,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 +728,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<void> {
     this.logSocketTerminal(ws, "socket_closed", code);
     this.teardownPartner(ws, code, reason);
+    this.forgetSocket(ws);
   }
 
   async webSocketError(ws: WebSocket, _error: unknown): Promise<void> {
     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 +823,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(