diff --git a/relay/src/channel-do.ts b/relay/src/channel-do.ts index 98019be..0513187 100644 --- a/relay/src/channel-do.ts +++ b/relay/src/channel-do.ts @@ -1,4 +1,5 @@ import { MAX_SSE_CONNECTIONS_PER_CHANNEL } from "@bridgehook/shared"; +import { SseHub } from "./sse-hub.js"; import { type SyncResult, SyncWaiters, clampSyncTimeout } from "./sync.js"; /** @@ -15,8 +16,10 @@ import { type SyncResult, SyncWaiters, clampSyncTimeout } from "./sync.js"; * response once an executor reports it) */ export class ChannelDO implements DurableObject { - private sseWriters: Set = new Set(); - private encoder = new TextEncoder(); + private hub = new SseHub({ + maxConnections: MAX_SSE_CONNECTIONS_PER_CHANNEL, + headers: { "Access-Control-Allow-Origin": "*" }, + }); private waiters = new SyncWaiters(); constructor( @@ -29,7 +32,7 @@ export class ChannelDO implements DurableObject { const path = url.pathname; if (request.method === "GET" && path.endsWith("/events")) { - return this.handleSSE(request); + return this.hub.open(request.signal); } if (request.method === "POST" && path.endsWith("/notify")) { @@ -43,41 +46,6 @@ export class ChannelDO implements DurableObject { return new Response("Not Found", { status: 404 }); } - private handleSSE(request: Request): Response { - // Enforce per-channel connection cap to prevent resource exhaustion - if (this.sseWriters.size >= MAX_SSE_CONNECTIONS_PER_CHANNEL) { - return new Response("Too many connections", { - status: 429, - headers: { "Access-Control-Allow-Origin": "*" }, - }); - } - - const { readable, writable } = new TransformStream(); - const writer = writable.getWriter(); - this.sseWriters.add(writer); - - // Send connected event (best-effort) - writer - .write(this.encoder.encode(`data: ${JSON.stringify({ type: "connected" })}\n\n`)) - .catch((err) => console.error("SSE initial write failed:", err)); - - // Cleanup on disconnect - request.signal.addEventListener("abort", () => { - this.sseWriters.delete(writer); - writer.close().catch(() => { - /* already closed */ - }); - }); - - return new Response(readable, { - headers: { - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "Access-Control-Allow-Origin": "*", - }, - }); - } - private async handleWait(request: Request): Promise { const body = (await request.json().catch(() => null)) as { eventId?: unknown; @@ -97,38 +65,8 @@ export class ChannelDO implements DurableObject { // whether or not anyone is watching the SSE stream. this.settleWaiter(payload); - if (this.sseWriters.size === 0) { - return new Response(JSON.stringify({ pushed: 0 }), { - status: 200, - headers: { "Content-Type": "application/json" }, - }); - } - - const message = this.encoder.encode(`data: ${payload}\n\n`); - - // Fan out in parallel — a slow client must not block others. - // Remove writers whose write rejects (client disconnected). - const writers = Array.from(this.sseWriters); - const results = await Promise.allSettled(writers.map((w) => w.write(message))); - - let pushed = 0; - results.forEach((r, i) => { - if (r.status === "fulfilled") { - pushed++; - } else { - const w = writers[i]; - this.sseWriters.delete(w); - console.error("SSE writer failed, removed:", r.reason); - w.close().catch(() => { - /* already closed */ - }); - } - }); - - return new Response(JSON.stringify({ pushed, connected: this.sseWriters.size }), { - status: 200, - headers: { "Content-Type": "application/json" }, - }); + const pushed = this.hub.size === 0 ? 0 : await this.hub.broadcast(payload); + return Response.json({ pushed, connected: this.hub.size }); } private settleWaiter(payload: string): void { diff --git a/relay/src/notify-lifetime.test.ts b/relay/src/notify-lifetime.test.ts index 6aaf039..19edc44 100644 --- a/relay/src/notify-lifetime.test.ts +++ b/relay/src/notify-lifetime.test.ts @@ -24,7 +24,7 @@ function sourceFiles(dir: string): string[] { * when it awaits or stores the promise for waitUntil, or when the line before * it opens a `background(` call. */ -export function keptAlive(lines: string[], i: number): boolean { +function keptAlive(lines: string[], i: number): boolean { let open = i; while (open > 0 && lines[open].trim().startsWith(".")) open--; const head = lines[open].trim(); diff --git a/relay/src/sse-hub.test.ts b/relay/src/sse-hub.test.ts new file mode 100644 index 0000000..c6bbaab --- /dev/null +++ b/relay/src/sse-hub.test.ts @@ -0,0 +1,102 @@ +import { describe, expect, it } from "vitest"; +import { SseHub } from "./sse-hub.js"; + +const decoder = new TextDecoder(); + +/** A client that keeps reading its stream, collecting what arrives. */ +function reader(res: Response) { + const r = (res.body as ReadableStream).getReader(); + const got: string[] = []; + let closed = false; + const loop = (async () => { + try { + for (;;) { + const { value, done } = await r.read(); + if (done) break; + got.push(decoder.decode(value)); + } + } catch { + /* aborted */ + } + closed = true; + })(); + return { + got, + cancel: () => r.cancel(), + closed: () => closed, + loop, + }; +} + +const tick = (ms = 0) => new Promise((r) => setTimeout(r, ms)); + +describe("SseHub", () => { + it("sends connected, then broadcasts to every reading client", async () => { + const hub = new SseHub({ maxConnections: 4, heartbeatMs: 0 }); + const a = reader(hub.open()); + const b = reader(hub.open()); + expect(await hub.broadcast('{"type":"webhook"}')).toBe(2); + await tick(); + for (const c of [a, b]) { + expect(c.got.join("")).toContain('"type":"connected"'); + expect(c.got.join("")).toContain('data: {"type":"webhook"}'); + } + }); + + it("forgets a client that cancelled its stream", async () => { + const hub = new SseHub({ maxConnections: 4, heartbeatMs: 0 }); + const a = reader(hub.open()); + await tick(); + await a.cancel(); + await hub.broadcast("{}"); + await tick(); + expect(hub.size).toBe(0); + }); + + it("drops a client that stopped reading, without holding up the others", async () => { + const hub = new SseHub({ maxConnections: 4, writeTimeoutMs: 50, heartbeatMs: 0 }); + hub.open(); // never read: its writes stay pending + const live = reader(hub.open()); + const started = Date.now(); + const pushed = await hub.broadcast('{"n":1}'); + expect(Date.now() - started).toBeLessThan(1000); + expect(pushed).toBe(1); + expect(hub.size).toBe(1); + await tick(); + expect(live.got.join("")).toContain('{"n":1}'); + }); + + it("never refuses a new client at the cap: the oldest makes room", async () => { + const hub = new SseHub({ maxConnections: 2, heartbeatMs: 0 }); + const first = reader(hub.open()); + reader(hub.open()); + const newest = reader(hub.open()); + expect(hub.size).toBe(2); + await tick(10); + expect(first.closed()).toBe(true); + expect(await hub.broadcast('{"x":1}')).toBe(2); + await tick(); + expect(newest.got.join("")).toContain('{"x":1}'); + }); + + it("drops a client on abort", async () => { + const hub = new SseHub({ maxConnections: 4, heartbeatMs: 0 }); + const ctl = new AbortController(); + reader(hub.open(ctl.signal)); + expect(hub.size).toBe(1); + ctl.abort(); + expect(hub.size).toBe(0); + }); + + it("heartbeats find a dead client even when nothing is broadcast", async () => { + const hub = new SseHub({ maxConnections: 4, writeTimeoutMs: 30, heartbeatMs: 20 }); + hub.open(); // never read + const live = reader(hub.open()); + await tick(200); + expect(hub.size).toBe(1); + expect(live.got.join("")).toContain(": heartbeat"); + await live.cancel(); + await tick(100); + expect(hub.size).toBe(0); + }); +}); diff --git a/relay/src/sse-hub.ts b/relay/src/sse-hub.ts new file mode 100644 index 0000000..c4a1780 --- /dev/null +++ b/relay/src/sse-hub.ts @@ -0,0 +1,152 @@ +/** + * The Server-Sent Events connections a Durable Object holds (per user in + * UserDO, per channel in ChannelDO), and how dead ones are let go. + * + * A client that goes away is not reliably reported: the request's abort + * signal does not always cross the Worker → Durable Object hop, and a write + * to a stream nobody reads does not reject, it stays pending. Before this + * hub (HAL-2722) such connections stayed counted forever, so a user whose + * extension had restarted eight times could no longer open a live stream + * ("Too many connections") until the next deploy, and one pending write held + * up every notify and heartbeat for the others. Now: + * + * - every write has a deadline; a connection that misses it is dropped; + * - a cancelled stream is noticed through the writer's `closed` promise as + * well as the abort signal; + * - at the cap the oldest connection is closed to make room, because a + * client that is connecting is the one most likely to be alive; + * - a heartbeat runs while anything is connected, so dead connections are + * found within one interval even when no webhook arrives. + */ + +export interface SseHubOptions { + /** Most open connections; the oldest is closed to admit a new one. */ + maxConnections: number; + /** A write that has not completed by then marks its connection dead. */ + writeTimeoutMs?: number; + /** Comment frame interval while anything is connected; 0 disables. */ + heartbeatMs?: number; + /** CORS for the stream response. */ + headers?: Record; +} + +const encoder = new TextEncoder(); +const CONNECTED = encoder.encode(`data: ${JSON.stringify({ type: "connected" })}\n\n`); +const HEARTBEAT = encoder.encode(": heartbeat\n\n"); + +interface Connection { + writer: WritableStreamDefaultWriter; +} + +export class SseHub { + private connections = new Set(); + private heartbeat: ReturnType | null = null; + private readonly writeTimeoutMs: number; + private readonly heartbeatMs: number; + + constructor(private readonly opts: SseHubOptions) { + this.writeTimeoutMs = opts.writeTimeoutMs ?? 10_000; + this.heartbeatMs = opts.heartbeatMs ?? 20_000; + } + + get size(): number { + return this.connections.size; + } + + /** Opens a stream for a client; `signal` is the client's request signal. */ + open(signal?: AbortSignal): Response { + while (this.connections.size >= this.opts.maxConnections) { + const oldest = this.connections.values().next().value; + if (!oldest) break; + this.drop(oldest, "evicted for a newer connection"); + } + + const { readable, writable } = new TransformStream(); + const conn: Connection = { writer: writable.getWriter() }; + this.connections.add(conn); + + conn.writer.closed.then( + () => this.forget(conn), + () => this.forget(conn), + ); + signal?.addEventListener("abort", () => this.drop(conn, "client aborted")); + + void this.write(conn, CONNECTED); + this.startHeartbeat(); + + return new Response(readable, { + headers: { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + ...(this.opts.headers ?? {}), + }, + }); + } + + /** Sends one `data:` frame to every connection; resolves with how many took it. */ + async broadcast(payload: string): Promise { + return this.fanOut(encoder.encode(`data: ${payload}\n\n`)); + } + + private async fanOut(frame: Uint8Array): Promise { + const results = await Promise.all(Array.from(this.connections, (c) => this.write(c, frame))); + return results.filter(Boolean).length; + } + + /** True when the write completed in time; otherwise the connection is dropped. */ + private async write(conn: Connection, frame: Uint8Array): Promise { + let timer: ReturnType | undefined; + const deadline = new Promise<"timeout">((resolve) => { + timer = setTimeout(() => resolve("timeout"), this.writeTimeoutMs); + }); + // A write that loses the race is rejected later by drop()'s abort(); + // handle that here so it is not an unhandled rejection. + const written = conn.writer.write(frame); + written.catch(() => {}); + try { + const outcome = await Promise.race([written.then(() => "ok" as const), deadline]); + if (outcome === "timeout") { + this.drop(conn, "write timed out"); + return false; + } + return true; + } catch { + this.drop(conn, "write failed"); + return false; + } finally { + clearTimeout(timer); + } + } + + private drop(conn: Connection, reason: string): void { + if (!this.connections.has(conn)) return; + this.forget(conn); + // Which path releases connections in production (abort, a failed or + // late write, or eviction) shows in the Worker logs. + console.log(`sse-hub: dropped a connection (${reason}); ${this.connections.size} left`); + // abort() rejects pending writes, so nothing stays awaited on it. + conn.writer.abort(reason).catch(() => { + /* already closed */ + }); + } + + private forget(conn: Connection): void { + this.connections.delete(conn); + if (this.connections.size === 0) this.stopHeartbeat(); + } + + private startHeartbeat(): void { + if (this.heartbeat !== null || this.heartbeatMs <= 0) return; + this.heartbeat = setInterval(() => { + void this.fanOut(HEARTBEAT); + }, this.heartbeatMs); + } + + private stopHeartbeat(): void { + if (this.heartbeat !== null) { + clearInterval(this.heartbeat); + this.heartbeat = null; + } + } +} diff --git a/relay/src/user-do.ts b/relay/src/user-do.ts index a76d82b..ce5e0e7 100644 --- a/relay/src/user-do.ts +++ b/relay/src/user-do.ts @@ -1,3 +1,5 @@ +import { SseHub } from "./sse-hub.js"; + /** * Durable Object for per-USER SSE connections (cross-channel). * @@ -21,19 +23,26 @@ * polling/SSE remains the canonical path). */ const HEARTBEAT_INTERVAL_MS = 20_000; -const MAX_SSE_CONNECTIONS_PER_USER = 8; +/** + * At the cap the oldest connection is closed to admit a new one (SseHub). + * Set well above what one person keeps open (dashboard tabs, the extension, + * a few CLIs), so eviction removes connections that are already dead rather + * than live clients evicting each other in a reconnect loop. + */ +const MAX_SSE_CONNECTIONS_PER_USER = 32; export class UserDO implements DurableObject { - private sseWriters: Set = new Set(); - private encoder = new TextEncoder(); - private heartbeatHandle: ReturnType | null = null; + private hub = new SseHub({ + maxConnections: MAX_SSE_CONNECTIONS_PER_USER, + heartbeatMs: HEARTBEAT_INTERVAL_MS, + headers: { "Access-Control-Allow-Origin": "*" }, + }); constructor( private state: DurableObjectState, private env: unknown, ) { - // Touch references so TS doesn't complain about unused params; the - // state/env objects are kept for future expansion (alarms / hibernation). + // Kept for future expansion (alarms / hibernation). void this.state; void this.env; } @@ -43,101 +52,14 @@ export class UserDO implements DurableObject { const path = url.pathname; if (request.method === "GET" && path.endsWith("/stream")) { - return this.handleSSE(request); + return this.hub.open(request.signal); } if (request.method === "POST" && path.endsWith("/notify")) { - return this.handleNotify(request); + const payload = await request.text(); + const pushed = this.hub.size === 0 ? 0 : await this.hub.broadcast(payload); + return Response.json({ pushed, connected: this.hub.size }); } return new Response("Not Found", { status: 404 }); } - - private handleSSE(request: Request): Response { - if (this.sseWriters.size >= MAX_SSE_CONNECTIONS_PER_USER) { - return new Response("Too many connections", { - status: 429, - headers: { "Access-Control-Allow-Origin": "*" }, - }); - } - - const { readable, writable } = new TransformStream(); - const writer = writable.getWriter(); - this.sseWriters.add(writer); - - // Initial frame so the client confirms it's connected. - writer - .write(this.encoder.encode(`data: ${JSON.stringify({ type: "connected" })}\n\n`)) - .catch((err) => console.error("UserDO initial write failed:", err)); - - // Start heartbeat on first connection. The interval ticks in the DO - // isolate's event loop and writes a comment frame to every writer — - // dead writers fall out via the same cleanup path as real notifies. - if (this.heartbeatHandle === null) { - this.heartbeatHandle = setInterval(() => { - this.heartbeat().catch((err) => console.error("heartbeat error:", err)); - }, HEARTBEAT_INTERVAL_MS); - } - - request.signal.addEventListener("abort", () => { - this.sseWriters.delete(writer); - writer.close().catch(() => { - /* already closed */ - }); - if (this.sseWriters.size === 0 && this.heartbeatHandle !== null) { - clearInterval(this.heartbeatHandle); - this.heartbeatHandle = null; - } - }); - - return new Response(readable, { - headers: { - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "Access-Control-Allow-Origin": "*", - "X-Accel-Buffering": "no", - }, - }); - } - - private async handleNotify(request: Request): Promise { - if (this.sseWriters.size === 0) { - return new Response(JSON.stringify({ pushed: 0 }), { - status: 200, - headers: { "Content-Type": "application/json" }, - }); - } - const payload = await request.text(); - const message = this.encoder.encode(`data: ${payload}\n\n`); - const pushed = await this.fanOut(message); - return new Response(JSON.stringify({ pushed, connected: this.sseWriters.size }), { - status: 200, - headers: { "Content-Type": "application/json" }, - }); - } - - private async heartbeat(): Promise { - if (this.sseWriters.size === 0) return; - // SSE comment frame — ignored by EventSource clients but keeps the - // stream alive across intermediaries. - const frame = this.encoder.encode(": heartbeat\n\n"); - await this.fanOut(frame); - } - - private async fanOut(message: Uint8Array): Promise { - const writers = Array.from(this.sseWriters); - const results = await Promise.allSettled(writers.map((w) => w.write(message))); - let pushed = 0; - results.forEach((r, i) => { - if (r.status === "fulfilled") { - pushed++; - } else { - const w = writers[i]; - this.sseWriters.delete(w); - w.close().catch(() => { - /* already closed */ - }); - } - }); - return pushed; - } }