From 3cb4c661520ae4e7916a0cc28d5b61a1f9851108 Mon Sep 17 00:00:00 2001 From: mbroton Date: Fri, 13 Nov 2026 01:53:41 +0100 Subject: [PATCH 1/2] fix(worker): keep serving while a recycled browser finishes its sessions A recycle drained the worker first: it took no new sessions until all open sessions ended or DRAIN_TIMEOUT expired, so one slow session stopped the worker's capacity at each recycle. The worker now starts the new browser at once, sends new sessions to it, and closes the old browser after its last session ends or DRAIN_TIMEOUT expires. The session budget counts accepted sessions, so the recycle starts when the budget is spent. If the budget runs out again while the old browser is still open, the worker renews the budget and keeps the current browser, so it never runs more than two browsers. SIGTERM still drains and exits. --- README.md | 10 +- server/internal/httpapi/server.go | 3 - website/src/content/docs/docs/architecture.md | 22 +- .../src/content/docs/docs/configuration.md | 4 +- worker/README.md | 2 +- worker/src/gateway.test.ts | 99 +++++- worker/src/gateway.ts | 55 ++- worker/src/worker.test.ts | 330 ++++++++++++++---- worker/src/worker.ts | 143 +++++--- 9 files changed, 531 insertions(+), 137 deletions(-) diff --git a/README.md b/README.md index f596cb3..b03ec32 100644 --- a/README.md +++ b/README.md @@ -193,10 +193,12 @@ flowchart TD - **PostgreSQL** holds all durable state: sessions, workers, and API keys are rows. Sessions of dead workers are closed out automatically, so capacity recovers without intervention. -- **Recycling**: after a configurable number of sessions, a worker drains - and replaces its browser without restarting the container. Sessions still - running after `DRAIN_TIMEOUT` are closed. Selection concentrates load on the - longest-serving worker, so recycles tend to happen one worker at a time. +- **Recycling**: after a configurable number of sessions, a worker replaces + its browser without restarting the container. New sessions go to the new + browser at once, and open sessions finish on the old one. Sessions still + open on the old browser after `DRAIN_TIMEOUT` are closed. Selection + concentrates load on the longest-serving worker, so recycles tend to + happen one worker at a time. ## Production deployment diff --git a/server/internal/httpapi/server.go b/server/internal/httpapi/server.go index bd4d02d..666bb42 100644 --- a/server/internal/httpapi/server.go +++ b/server/internal/httpapi/server.go @@ -362,9 +362,6 @@ func registerWorkerRoutes( } type registeredWorker struct { Worker - // Session budget for this worker (0 = unlimited). The worker drains - // itself once it has served this many sessions, without waiting for - // the heartbeat that would tell it to. MaxLifetimeSessions int64 `json:"max_lifetime_sessions" doc:"Sessions this worker may serve before it must recycle; 0 means unlimited"` } type registerOutput struct { diff --git a/website/src/content/docs/docs/architecture.md b/website/src/content/docs/docs/architecture.md index e590e5f..e9fe7ac 100644 --- a/website/src/content/docs/docs/architecture.md +++ b/website/src/content/docs/docs/architecture.md @@ -58,12 +58,22 @@ arrival order. Each server replica has its own queue. See ## Worker recycling and failures -By default, a worker drains after 50 lifetime session claims. Draining stops new -sessions and waits for existing sessions for up to `DRAIN_TIMEOUT`. The worker -then replaces its browser process and resumes service without restarting the -container. A shutdown signal such as `SIGTERM` drains and stops the worker instead. - -When the drain timeout expires, remaining sessions are closed. The supplied +By default, a worker recycles its browser after 50 lifetime session claims. It +starts a new browser and sends new sessions to it at once, without restarting +the container. Open sessions finish on the old browser, which closes after its +last session ends. A slow session therefore does not stop the worker from +taking new sessions. While both browsers run, the worker needs memory for one +more browser. + +A worker runs at most two browsers. If it reaches the session limit again +before the old browser closes, the current browser keeps serving and recycles +at a later limit. + +A shutdown signal such as `SIGTERM` stops new sessions, waits for open +sessions, and then stops the worker. + +When `DRAIN_TIMEOUT` expires on an old browser or on a worker that is shutting +down, remaining sessions are closed. The supplied Compose file sets this timeout to **30 seconds**; the worker's standalone default is 300 seconds. Set it for the time your tasks need to finish. See the [worker settings](/browserthing/docs/configuration/#worker). diff --git a/website/src/content/docs/docs/configuration.md b/website/src/content/docs/docs/configuration.md index d766f24..a76f9a0 100644 --- a/website/src/content/docs/docs/configuration.md +++ b/website/src/content/docs/docs/configuration.md @@ -17,7 +17,7 @@ Server duration values use Go duration notation, such as `10s` or `5m`. | `DEFAULT_BROWSER_TYPE` | `chromium` | Browser selected when a connection omits `?browser=`. | | `MAX_QUEUE_SIZE` | `100` | Maximum waiting requests per server replica. `0` disables queueing. | | `QUEUE_WAIT_TIMEOUT` | `30s` | Maximum wait for a worker slot. | -| `MAX_LIFETIME_SESSIONS` | `50` | Session claims before a worker drains and recycles. `0` disables recycling. | +| `MAX_LIFETIME_SESSIONS` | `50` | Session claims before a worker recycles its browser. `0` disables recycling. | | `WORKER_HEARTBEAT_TTL` | `30s` | Time before an available worker with no heartbeat becomes stalled. | | `SESSION_HEARTBEAT_TTL` | `30s` | Time before a pending or running session with no heartbeat expires. | | `SESSION_HEARTBEAT_INTERVAL` | `10s` | How often a live relay renews its session. | @@ -57,7 +57,7 @@ Worker time values are integer **seconds**. | `MAX_SLOTS` | `5` | Concurrent sessions per worker, from 1 to 1024. | | `HEADLESS` | `true` | Run the browser without a visible window. | | `HEARTBEAT_INTERVAL` | `5` | Interval between worker heartbeats. | -| `DRAIN_TIMEOUT` | `300` | Maximum wait for active sessions during recycling or shutdown. Remaining sessions close when it expires. | +| `DRAIN_TIMEOUT` | `300` | Maximum wait for open sessions on the old browser after a recycle, or during shutdown. Remaining sessions close when it expires. | | `LOG_LEVEL` | `info` | One of `debug`, `info`, `warn`, or `error`. | | `LOG_FORMAT` | `json` | One of `json` or `text`. | diff --git a/worker/README.md b/worker/README.md index 28568ec..2c39e12 100644 --- a/worker/README.md +++ b/worker/README.md @@ -24,7 +24,7 @@ The right value depends on what the sessions do: When in doubt, keep the default and add workers instead of raising `MAX_SLOTS`: capacity grows the same either way, and each worker bounds how -many sessions a browser crash or recycle can take down with it. +many sessions a browser crash can take down with it. ## Scaling beyond one machine diff --git a/worker/src/gateway.test.ts b/worker/src/gateway.test.ts index a611572..3643391 100644 --- a/worker/src/gateway.test.ts +++ b/worker/src/gateway.test.ts @@ -7,6 +7,7 @@ import WebSocket, { WebSocketServer } from 'ws'; import { SessionGateway } from './gateway.js'; const sessionId = '11111111-1111-4111-8111-111111111111'; +const secondSessionId = '22222222-2222-4222-8222-222222222222'; test('relays text and binary frames and tracks the session map', async t => { const fixture = await createFixture(); @@ -195,9 +196,103 @@ test('shutting down takes precedence over duplicate session rejection', async t assert.match(response, /^HTTP\/1\.1 503 /); }); +test('emits sessionOpen once for each accepted session', async t => { + const fixture = await createFixture(); + t.after(() => fixture.close()); + const opened: string[] = []; + fixture.gateway.on('sessionOpen', id => opened.push(id)); + + assert.equal(await refusedStatus(fixture.gatewayUrl, 'not-a-uuid'), 400); + const first = await connect(fixture.gatewayUrl, sessionId); + assert.equal(await refusedStatus(fixture.gatewayUrl, sessionId), 409); + const second = await connect(fixture.gatewayUrl, secondSessionId); + await waitFor(() => opened.length === 2); + + assert.deepEqual(opened, [sessionId, secondSessionId]); + first.terminate(); + second.terminate(); +}); + +test('a replaced browser endpoint keeps its sessions until they close', async t => { + const fixture = await createFixture(); + t.after(() => fixture.close()); + const oldEndpoint = fixture.browserEndpoint; + const newEndpoint = `${oldEndpoint}/replacement`; + const oldClient = await connect(fixture.gatewayUrl, sessionId); + await fixture.nextUpstreamRequest(); + + fixture.gateway.setBrowserEndpoint(newEndpoint); + const newClient = await connect(fixture.gatewayUrl, secondSessionId); + t.after(() => newClient.terminate()); + assert.equal((await fixture.nextUpstreamRequest()).url, '/replacement'); + + let oldIdle = false; + const oldIdlePromise = fixture.gateway.whenEndpointIdle(oldEndpoint).then(() => { + oldIdle = true; + }); + const oldEcho = nextMessage(oldClient); + oldClient.send('still served'); + assert.deepEqual(await oldEcho, { data: Buffer.from('still served'), isBinary: false }); + assert.equal(oldIdle, false); + + oldClient.close(1000, 'done'); + await bounded(oldIdlePromise, 500); + assert.deepEqual(fixture.gateway.activeSessionIds, [secondSessionId]); + + let newIdle = false; + void fixture.gateway.whenEndpointIdle(newEndpoint).then(() => { + newIdle = true; + }); + await delay(20); + assert.equal(newIdle, false); +}); + +test('an endpoint is not idle while a dial to it is still connecting', async t => { + const fixture = await createFixture(100); + t.after(() => fixture.close()); + const oldEndpoint = fixture.browserEndpoint; + const dialing = connect(fixture.gatewayUrl, sessionId); + await fixture.nextUpstreamRequest(); + + fixture.gateway.setBrowserEndpoint(`${oldEndpoint}/replacement`); + let oldIdle = false; + const oldIdlePromise = fixture.gateway.whenEndpointIdle(oldEndpoint).then(() => { + oldIdle = true; + }); + const client = await dialing; + assert.equal(oldIdle, false); + + client.close(); + await bounded(oldIdlePromise, 500); +}); + +test('closing the sessions of an endpoint leaves other endpoints open', async t => { + const fixture = await createFixture(); + t.after(() => fixture.close()); + const oldEndpoint = fixture.browserEndpoint; + const oldClient = await connect(fixture.gatewayUrl, sessionId); + fixture.gateway.setBrowserEndpoint(`${oldEndpoint}/replacement`); + const newClient = await connect(fixture.gatewayUrl, secondSessionId); + t.after(() => newClient.terminate()); + await waitFor(() => fixture.gateway.activeConnectionCount === 2); + const oldClosed = waitForClose(oldClient); + + fixture.gateway.closeEndpointSessions(oldEndpoint, 1001, 'worker is recycling'); + + const close = await bounded(oldClosed, 500); + assert.equal(close.code, 1001); + assert.equal(close.reason.toString(), 'worker is recycling'); + await bounded(fixture.gateway.whenEndpointIdle(oldEndpoint), 500); + assert.deepEqual(fixture.gateway.activeSessionIds, [secondSessionId]); + const newEcho = nextMessage(newClient); + newClient.send('new browser'); + assert.deepEqual(await newEcho, { data: Buffer.from('new browser'), isBinary: false }); +}); + async function createFixture(upstreamHandshakeDelayMs = 0): Promise<{ gateway: SessionGateway; gatewayUrl: string; + browserEndpoint: string; nextUpstreamRequest: () => Promise; nextUpstreamSocket: () => Promise; nextUpstreamMessage: () => Promise<{ data: Buffer; isBinary: boolean }>; @@ -247,13 +342,15 @@ async function createFixture(upstreamHandshakeDelayMs = 0): Promise<{ }); }); - const gateway = new SessionGateway(`ws://127.0.0.1:${browserAddress.port}`, 0, '127.0.0.1'); + const browserEndpoint = `ws://127.0.0.1:${browserAddress.port}`; + const gateway = new SessionGateway(browserEndpoint, 0, '127.0.0.1'); await gateway.start(); let closePromise: Promise | null = null; return { gateway, gatewayUrl: `ws://127.0.0.1:${gateway.listeningPort}`, + browserEndpoint, nextUpstreamRequest: () => { const request = requests.shift(); return request ? Promise.resolve(request) : new Promise(resolve => requestWaiters.push(resolve)); diff --git a/worker/src/gateway.ts b/worker/src/gateway.ts index 8245109..7b15b09 100644 --- a/worker/src/gateway.ts +++ b/worker/src/gateway.ts @@ -6,14 +6,17 @@ import WebSocket, { WebSocketServer, type RawData } from 'ws'; interface SessionSockets { client: WebSocket; upstream: WebSocket; + endpoint: string; } interface PendingSession { clientSocket: Duplex; upstream: WebSocket; + endpoint: string; } interface GatewayEvents { + sessionOpen: [sessionId: string]; sessionClose: [sessionId: string]; } @@ -33,6 +36,7 @@ export class SessionGateway extends EventEmitter { private readonly webSocketServer = new WebSocketServer({ noServer: true }); private readonly sessions = new Map(); private readonly pendingSessions = new Map(); + private readonly idleWaiters: Array<{ endpoint: string; resolve: () => void }> = []; private shuttingDown = false; private shutdownPromise: Promise | null = null; private listenPort: number | null = null; @@ -54,6 +58,22 @@ export class SessionGateway extends EventEmitter { this.browserEndpoint = endpoint; } + /** Resolves when no open or connecting session uses the browser endpoint. */ + whenEndpointIdle(endpoint: string): Promise { + if (!this.usesEndpoint(endpoint)) { + return Promise.resolve(); + } + return new Promise(resolve => this.idleWaiters.push({ endpoint, resolve })); + } + + closeEndpointSessions(endpoint: string, code: number, reason: string): void { + for (const [sessionId, session] of this.sessions) { + if (session.endpoint === endpoint) { + this.closeSession(sessionId, code, reason, session); + } + } + } + get activeSessionIds(): string[] { return [...this.sessions.keys()]; } @@ -122,6 +142,7 @@ export class SessionGateway extends EventEmitter { rejectUpgrade(session.clientSocket, 503, 'Worker is shutting down'); } } + this.releaseIdleWaiters(); const closed = Promise.all([ closeWebSocketServer(this.webSocketServer), @@ -153,11 +174,12 @@ export class SessionGateway extends EventEmitter { return; } + const endpoint = this.browserEndpoint; let upstream: WebSocket; try { // Node validates forwarded header values synchronously inside the // WebSocket constructor; a throw here must not reach uncaughtException. - upstream = new WebSocket(this.browserEndpoint, { + upstream = new WebSocket(endpoint, { headers: forwardedHeaders(request.headers), handshakeTimeout: upstreamHandshakeTimeoutMs, }); @@ -173,7 +195,7 @@ export class SessionGateway extends EventEmitter { this.closeSession(sessionId, 1011, 'browser connection failed', acceptedSession); } }); - this.pendingSessions.set(sessionId, { clientSocket: socket, upstream }); + this.pendingSessions.set(sessionId, { clientSocket: socket, upstream, endpoint }); const clearPending = (): boolean => { if (this.pendingSessions.get(sessionId)?.upstream !== upstream) { @@ -182,6 +204,7 @@ export class SessionGateway extends EventEmitter { this.pendingSessions.delete(sessionId); socket.off('close', abortDial); upstream.off('open', acceptClient); + this.releaseIdleWaiters(); return true; }; const abortDial = () => { @@ -209,7 +232,7 @@ export class SessionGateway extends EventEmitter { } try { this.webSocketServer.handleUpgrade(request, socket, head, client => { - const session = { client, upstream }; + const session = { client, upstream, endpoint }; acceptedSession = session; client.once('error', () => this.closeSession( sessionId, @@ -220,6 +243,7 @@ export class SessionGateway extends EventEmitter { this.sessions.set(sessionId, session); this.pipeSession(sessionId, session); clearPending(); + this.emit('sessionOpen', sessionId); }); } catch { rejectClient(400, 'Invalid WebSocket upgrade'); @@ -260,8 +284,33 @@ export class SessionGateway extends EventEmitter { return; } this.sessions.delete(sessionId); + this.releaseIdleWaiters(); this.emit('sessionClose', sessionId); } + + private usesEndpoint(endpoint: string): boolean { + for (const session of this.sessions.values()) { + if (session.endpoint === endpoint) { + return true; + } + } + for (const session of this.pendingSessions.values()) { + if (session.endpoint === endpoint) { + return true; + } + } + return false; + } + + private releaseIdleWaiters(): void { + for (let index = this.idleWaiters.length - 1; index >= 0; index -= 1) { + const waiter = this.idleWaiters[index]!; + if (!this.usesEndpoint(waiter.endpoint)) { + this.idleWaiters.splice(index, 1); + waiter.resolve(); + } + } + } } function forwardedHeaders(headers: IncomingHttpHeaders): IncomingHttpHeaders { diff --git a/worker/src/worker.test.ts b/worker/src/worker.test.ts index f675132..f25ce25 100644 --- a/worker/src/worker.test.ts +++ b/worker/src/worker.test.ts @@ -199,7 +199,8 @@ test('heartbeat requests time out and the loop keeps ticking', async t => { assert.ok(heartbeatAttempts >= 3); }); -test('heartbeat draining status enters drain mode', async t => { +test('heartbeat draining status recycles the browser', async t => { + let recycleRequests = 0; const server = await startControlPlane(async (request, response) => { if (request.url === '/internal/workers') { sendJSON(response, 201, { id: 'worker-1' }); @@ -209,16 +210,18 @@ test('heartbeat draining status enters drain mode', async t => { sendJSON(response, 200, { status: 'draining', stale_session_ids: [] }); return; } + if (request.url?.endsWith('/recycle')) { + recycleRequests += 1; + } sendJSON(response, 200, { id: 'worker-1' }); }); t.after(() => server.close()); const worker = createWorker(server.url, new FakeGateway([activeId])); await worker.start(); - await waitFor(() => worker.state === 'draining'); + await waitFor(() => recycleRequests >= 1); + assert.notEqual(worker.state, 'draining'); await worker.shutdown(0); - - assert.equal(worker.state, 'shutting_down'); }); test('a draining worker re-asserts drain intent after an available heartbeat', async () => { @@ -226,12 +229,11 @@ test('a draining worker re-asserts drain intent after an available heartbeat', a const gateway = new FakeGateway([activeId]); const worker = createInjectedWorker(controlPlane, gateway, 5_000); await worker.start(); - await worker.requestDrain('SIGTERM'); + await worker.requestDrain('SIGTERM', true); - await waitFor(() => controlPlane.statuses.includes('draining')); + await waitFor(() => controlPlane.statuses.filter(status => status === 'draining').length >= 2); + assert.equal(worker.state, 'draining'); await worker.shutdown(0); - - assert.equal(controlPlane.statuses[0], 'draining'); }); test('heartbeat 404 registers a new worker ID', async t => { @@ -288,11 +290,12 @@ test('heartbeat 404 re-asserts drain status immediately after registration', asy }; const worker = createInjectedWorker(controlPlane, new FakeGateway([activeId]), 5_000); await worker.start(); - await worker.requestDrain('SIGTERM'); + await worker.requestDrain('SIGTERM', true); await waitFor(() => events.includes('status:worker-2:draining')); - assert.deepEqual(events.slice(0, 4), [ + assert.deepEqual(events.slice(0, 5), [ 'register:worker-1', + 'status:worker-1:draining', 'heartbeat:worker-1', 'register:worker-2', 'status:worker-2:draining', @@ -339,44 +342,164 @@ test('a control-plane drain recycles the browser in place instead of exiting', a await worker.shutdown(0); }); -test('drain timeout forces the browser swap with active connections', async () => { +test('a recycle keeps serving while sessions on the old browser finish', async () => { const controlPlane = new FakeControlPlane(); const gateway = new FakeGateway([activeId]); - const worker = createInjectedWorker(controlPlane, gateway, 30); + const browsers: FakeBrowser[] = []; + const worker = createRecordingWorker(controlPlane, gateway, browsers, 5_000); + await worker.start(); + + await bounded(worker.requestDrain('control plane'), 500); + assert.equal(worker.state, 'running'); + assert.deepEqual(controlPlane.recycles, ['worker-1']); + assert.equal(browsers.length, 2); + assert.equal(browsers[0]!.closed, false); + assert.deepEqual(gateway.closedSessions, []); + + gateway.addSession(secondId); + await delay(30); + assert.equal(browsers[0]!.closed, false); + + gateway.closeSession(activeId); + await waitFor(() => browsers[0]!.closed); + assert.equal(browsers[1]!.closed, false); + assert.deepEqual(gateway.activeSessionIds, [secondId]); + await worker.shutdown(0); +}); + +test('the old browser closes with its sessions when DRAIN_TIMEOUT expires', async () => { + const controlPlane = new FakeControlPlane(); + const gateway = new FakeGateway([activeId]); + const browsers: FakeBrowser[] = []; + const worker = createRecordingWorker(controlPlane, gateway, browsers, 30); await worker.start(); await worker.requestDrain('control plane'); - await waitFor(() => controlPlane.recycles.length === 1); + gateway.addSession(secondId); + await waitFor(() => browsers[0]!.closed); assert.equal(worker.state, 'running'); - assert.equal(gateway.activeConnectionCount, 0); assert.deepEqual(gateway.closedSessions, [{ sessionId: activeId, code: 1001 }]); - assert.deepEqual(controlPlane.recycles, ['worker-1']); + assert.deepEqual(gateway.activeSessionIds, [secondId]); + assert.equal(browsers[1]!.closed, false); + await worker.shutdown(0); +}); + +test('a spent session budget recycles at once while its sessions are still open', async () => { + const controlPlane = new FakeControlPlane(2); + const gateway = new FakeGateway(); + const browsers: FakeBrowser[] = []; + const worker = createRecordingWorker(controlPlane, gateway, browsers, 5_000); + await worker.start(); + + gateway.addSession(activeId); + await delay(20); + assert.equal(controlPlane.recycles.length, 0); + + gateway.addSession(secondId); + await waitFor(() => controlPlane.recycles.length === 1); + await waitFor(() => worker.state === 'running'); + assert.equal(browsers.length, 2); + assert.deepEqual(gateway.activeSessionIds, [activeId, secondId]); + assert.deepEqual(gateway.closedSessions, []); await worker.shutdown(0); }); +test('a spent budget while the old browser has open sessions keeps the current browser', async () => { + const controlPlane = new FakeControlPlane(1); + const gateway = new FakeGateway(); + const browsers: FakeBrowser[] = []; + const worker = createRecordingWorker(controlPlane, gateway, browsers, 5_000); + await worker.start(); + + gateway.addSession(activeId); + await waitFor(() => controlPlane.recycles.length === 1 && worker.state === 'running'); + assert.equal(browsers.length, 2); + + gateway.addSession(secondId); + await waitFor(() => controlPlane.recycles.length === 2 && worker.state === 'running'); + assert.equal(browsers.length, 2); + assert.equal(browsers[0]!.closed, false); + assert.deepEqual(gateway.closedSessions, []); + + gateway.closeSession(activeId); + await waitFor(() => browsers[0]!.closed); + gateway.addSession(thirdId); + await waitFor(() => controlPlane.recycles.length === 3 && worker.state === 'running'); + assert.equal(browsers.length, 3); + assert.equal(browsers[1]!.closed, false); + assert.equal(browsers.filter(browser => !browser.closed).length, 2); + + gateway.closeSession(secondId); + await delay(20); + assert.equal(browsers[1]!.closed, false); + gateway.closeSession(thirdId); + await waitFor(() => browsers[1]!.closed); + assert.equal(browsers[2]!.closed, false); + await worker.shutdown(0); +}); + +test('heartbeats report the sessions on both browsers during a recycle', async () => { + const reports: string[][] = []; + const controlPlane: ControlPlaneLike = { + register: async () => ({ id: 'worker-1' }), + recycle: async workerId => ({ id: workerId }), + heartbeat: async (_workerId, activeSessionIds) => { + reports.push(activeSessionIds); + return { status: 'available', stale_session_ids: [] }; + }, + setStatus: async () => undefined, + }; + const gateway = new FakeGateway([activeId]); + const worker = createInjectedWorker(controlPlane, gateway, 5_000); + await worker.start(); + + await worker.requestDrain('control plane'); + gateway.addSession(secondId); + const reportsBefore = reports.length; + await waitFor(() => reports.length > reportsBefore); + + assert.deepEqual(reports.at(-1), [activeId, secondId]); + await worker.shutdown(0); +}); + +test('a shutdown during a recycle closes both browsers', async () => { + const controlPlane = new FakeControlPlane(); + const gateway = new FakeGateway([activeId]); + const browsers: FakeBrowser[] = []; + const worker = createRecordingWorker(controlPlane, gateway, browsers, 5_000); + await worker.start(); + + await worker.requestDrain('control plane'); + assert.equal(browsers[0]!.closed, false); + await worker.shutdown(0); + + assert.equal(await worker.waitForExit(), 0); + assert.equal(browsers[0]!.closed, true); + assert.equal(browsers[1]!.closed, true); +}); + test('the worker keeps its ID and refreshes its session budget after recycle', async () => { const controlPlane = new FakeControlPlane(2, 3); - const gateway = new FakeGateway([activeId, secondId]); + const gateway = new FakeGateway(); const worker = createInjectedWorker(controlPlane, gateway, 5_000); await worker.start(); - gateway.closeSession(activeId); + gateway.addSession(activeId); assert.equal(worker.state, 'running'); - gateway.closeSession(secondId); - await waitFor(() => controlPlane.recycles.length === 1); - assert.equal(worker.state, 'running'); + gateway.addSession(secondId); + await waitFor(() => controlPlane.recycles.length === 1 && worker.state === 'running'); assert.equal(worker.workerId, 'worker-1'); assert.equal(controlPlane.registrations.length, 1); + gateway.closeSession(activeId); + gateway.closeSession(secondId); gateway.addSession(thirdId); gateway.addSession(fourthId); - gateway.addSession(fifthId); - gateway.closeSession(thirdId); - gateway.closeSession(fourthId); + await delay(20); assert.equal(controlPlane.recycles.length, 1); - gateway.closeSession(fifthId); + gateway.addSession(fifthId); await waitFor(() => controlPlane.recycles.length === 2); await worker.shutdown(0); }); @@ -838,6 +961,7 @@ test('the replacement endpoint is swapped before the old browser closes', async await worker.start(); await worker.requestDrain('control plane'); + await waitFor(() => events.length === 2); assert.deepEqual(events, ['swap:ws://127.0.0.1:12345/browser-2', 'close:browser-1']); await worker.shutdown(0); }); @@ -864,15 +988,15 @@ test('heartbeat-404 re-registration keeps the spent session budget', async () => }, setStatus: async () => undefined, }; - const gateway = new FakeGateway([activeId, secondId]); + const gateway = new FakeGateway(); const worker = createInjectedWorker(controlPlane, gateway, 5_000); await worker.start(); - gateway.closeSession(activeId); // budget 1 of 2 spent + gateway.addSession(activeId); // budget 1 of 2 spent await waitFor(() => registrations === 2); // 404 re-registers, same browser - gateway.closeSession(secondId); // 2 of 2 — must drain despite re-registration - await waitFor(() => recycles === 1); + gateway.addSession(secondId); // 2 of 2 — must recycle despite re-registration + await waitFor(() => recycles === 1 && worker.state === 'running'); assert.equal(registrations, 2); assert.equal(worker.state, 'running'); await worker.shutdown(0); @@ -920,7 +1044,47 @@ test('a failed browser relaunch exits for a supervisor restart', async () => { assert.equal(await worker.waitForExit(), 1); }); -test('drain timeout swaps the browser under a live session on the real gateway', async t => { +test('a recycle on the real gateway serves old and new sessions until the old one closes', async t => { + const { worker, browsers, gateway: getGateway } = createNetworkWorker(t, 5_000); + await worker.start(); + const gateway = getGateway(); + const oldClient = await connectWebSocket(`ws://127.0.0.1:${gateway.listeningPort}`, activeId); + await waitFor(() => gateway.activeConnectionCount === 1); + + await worker.requestDrain('control plane'); + assert.equal(worker.state, 'running'); + assert.equal(browsers.length, 2); + + const newClient = await connectWebSocket(`ws://127.0.0.1:${gateway.listeningPort}`, secondId); + assert.equal(await bounded(echo(newClient, 'new'), 1_000), 'browser-2:new'); + assert.equal(await bounded(echo(oldClient, 'old'), 1_000), 'browser-1:old'); + assert.equal(browsers[0]!.closed, false); + + oldClient.close(); + await waitFor(() => browsers[0]!.closed); + assert.equal(await bounded(echo(newClient, 'still'), 1_000), 'browser-2:still'); + newClient.terminate(); +}); + +test('drain timeout closes an old session on the real gateway with code 1001', async t => { + const { worker, browsers, gateway: getGateway } = createNetworkWorker(t, 50); + await worker.start(); + const gateway = getGateway(); + const oldClient = await connectWebSocket(`ws://127.0.0.1:${gateway.listeningPort}`, activeId); + await waitFor(() => gateway.activeConnectionCount === 1); + const oldClosed = new Promise(resolve => oldClient.once('close', code => resolve(code))); + + await worker.requestDrain('control plane'); + + assert.equal(await bounded(oldClosed, 1_000), 1001); + await waitFor(() => browsers[0]!.closed); + assert.equal(browsers[1]!.closed, false); +}); + +function createNetworkWorker( + t: { after: (fn: () => Promise) => void }, + drainTimeoutMs: number, +): { worker: BrowserWorker; browsers: NetworkBrowser[]; gateway: () => SessionGateway } { const browsers: NetworkBrowser[] = []; const launchNetworkBrowser = async (): Promise => { const server = new WebSocketServer({ host: '127.0.0.1', port: 0 }); @@ -929,15 +1093,15 @@ test('drain timeout swaps the browser under a live session on the real gateway', if (!address || typeof address === 'string') { throw new Error('browser fixture did not use a TCP port'); } - server.on('connection', socket => socket.on('message', data => socket.send(data))); + const name = `browser-${browsers.length + 1}`; + server.on('connection', socket => socket.on('message', data => socket.send(`${name}:${String(data)}`))); const browser = new NetworkBrowser(server, `ws://127.0.0.1:${address.port}`); browsers.push(browser); return browser; }; - const controlPlane = new FakeControlPlane(); let gateway!: SessionGateway; - const worker = new BrowserWorker(testConfig(30, 0), { - controlPlane, + const worker = new BrowserWorker(testConfig(drainTimeoutMs, 0), { + controlPlane: new FakeControlPlane(), createGateway: (endpoint, port) => { gateway = new SessionGateway(endpoint, port, '127.0.0.1'); return gateway; @@ -951,25 +1115,14 @@ test('drain timeout swaps the browser under a live session on the real gateway', await browser.kill(); } }); - await worker.start(); - const client = await connectWebSocket(`ws://${'127.0.0.1'}:${gateway.listeningPort}`, activeId); - await waitFor(() => gateway?.activeConnectionCount === 1); - - await worker.requestDrain('control plane'); - await waitFor(() => controlPlane.recycles.length === 1); + return { worker, browsers, gateway: () => gateway }; +} - assert.equal(worker.state, 'running'); - assert.equal(browsers.length, 2); - assert.equal(gateway.activeConnectionCount, 0); - client.terminate(); - - // A post-swap session must reach the replacement browser end to end. - const postSwap = await connectWebSocket(`ws://127.0.0.1:${gateway.listeningPort}`, secondId); - const echoed = new Promise(resolve => postSwap.once('message', data => resolve(String(data)))); - postSwap.send('after-swap'); - assert.equal(await bounded(echoed, 1_000), 'after-swap'); - postSwap.terminate(); -}); +function echo(socket: WebSocket, message: string): Promise { + const reply = new Promise(resolve => socket.once('message', data => resolve(String(data)))); + socket.send(message); + return reply; +} test('startup does not resume after shutdown wins a registration race', async () => { let resolveRegistration!: (worker: { id: string }) => void; @@ -1059,19 +1212,43 @@ class FakeGateway extends EventEmitter implements GatewayLike { readonly closedSessions: Array<{ sessionId: string; code: number }> = []; readonly browserEndpoints: string[] = []; shutdownCalled = false; - private readonly sessions: string[]; + private endpoint: string | undefined; + private readonly sessions: Array<{ sessionId: string; endpoint: string | undefined }>; + private readonly idleWaiters: Array<{ endpoint: string; resolve: () => void }> = []; constructor(sessionIds: string[] = [], readonly listeningPort = 3131) { super(); - this.sessions = [...sessionIds]; + this.sessions = sessionIds.map(sessionId => ({ sessionId, endpoint: undefined })); + } + + bind(endpoint: string): this { + this.endpoint = endpoint; + for (const session of this.sessions) { + session.endpoint ??= endpoint; + } + return this; } setBrowserEndpoint(endpoint: string): void { this.browserEndpoints.push(endpoint); + this.endpoint = endpoint; + } + + whenEndpointIdle(endpoint: string): Promise { + if (!this.sessions.some(session => session.endpoint === endpoint)) { + return Promise.resolve(); + } + return new Promise(resolve => this.idleWaiters.push({ endpoint, resolve })); + } + + closeEndpointSessions(endpoint: string, code: number): void { + for (const session of this.sessions.filter(candidate => candidate.endpoint === endpoint)) { + this.closeSession(session.sessionId, code); + } } get activeSessionIds(): string[] { - return [...this.sessions]; + return this.sessions.map(session => session.sessionId); } get activeConnectionCount(): number { @@ -1081,14 +1258,16 @@ class FakeGateway extends EventEmitter implements GatewayLike { async start(): Promise {} addSession(sessionId: string): void { - this.sessions.push(sessionId); + this.sessions.push({ sessionId, endpoint: this.endpoint }); + this.emit('sessionOpen', sessionId); } closeSession(sessionId: string, code = 1001): void { - const index = this.sessions.indexOf(sessionId); + const index = this.sessions.findIndex(session => session.sessionId === sessionId); if (index >= 0) { this.sessions.splice(index, 1); this.closedSessions.push({ sessionId, code }); + this.releaseIdleWaiters(); this.emit('sessionClose', sessionId); } } @@ -1096,6 +1275,16 @@ class FakeGateway extends EventEmitter implements GatewayLike { async shutdown(): Promise { this.shutdownCalled = true; this.sessions.splice(0); + this.releaseIdleWaiters(); + } + + private releaseIdleWaiters(): void { + for (const waiter of [...this.idleWaiters]) { + if (!this.sessions.some(session => session.endpoint === waiter.endpoint)) { + this.idleWaiters.splice(this.idleWaiters.indexOf(waiter), 1); + waiter.resolve(); + } + } } } @@ -1105,26 +1294,32 @@ class HangingShutdownGateway extends FakeGateway { } } +let fakeBrowserCount = 0; + class FakeBrowser extends EventEmitter implements BrowserServerLike { killCalled = false; + closed = false; + private readonly endpoint = `ws://127.0.0.1:12345/fake-browser-${++fakeBrowserCount}`; constructor(private readonly closeHangs = false) { super(); } wsEndpoint(): string { - return 'ws://127.0.0.1:12345/fake-browser'; + return this.endpoint; } async close(): Promise { if (this.closeHangs) { await new Promise(() => undefined); } + this.closed = true; this.emit('close'); } async kill(): Promise { this.killCalled = true; + this.closed = true; this.emit('close'); } } @@ -1145,7 +1340,7 @@ class WedgedBrowser extends EventEmitter implements BrowserServerLike { } class NetworkBrowser extends EventEmitter implements BrowserServerLike { - private closed = false; + closed = false; constructor( private readonly server: WebSocketServer, @@ -1212,7 +1407,7 @@ function createWorker(serverUrl: string, gateway: FakeGateway, heartbeatTimeoutM undefined, heartbeatTimeoutMs ?? config.lifecycle.heartbeatIntervalMs, ), - createGateway: () => gateway, + createGateway: endpoint => gateway.bind(endpoint), launchBrowser: async () => new FakeBrowser(), exit: () => undefined, }); @@ -1222,11 +1417,24 @@ function createInjectedWorker( controlPlane: ControlPlaneLike, gateway: FakeGateway, drainTimeoutMs: number, +): BrowserWorker { + return createRecordingWorker(controlPlane, gateway, [], drainTimeoutMs); +} + +function createRecordingWorker( + controlPlane: ControlPlaneLike, + gateway: FakeGateway, + browsers: FakeBrowser[], + drainTimeoutMs: number, ): BrowserWorker { return new BrowserWorker(testConfig(drainTimeoutMs), { controlPlane, - createGateway: () => gateway, - launchBrowser: async () => new FakeBrowser(), + createGateway: endpoint => gateway.bind(endpoint), + launchBrowser: async () => { + const browser = new FakeBrowser(); + browsers.push(browser); + return browser; + }, exit: () => undefined, }); } diff --git a/worker/src/worker.ts b/worker/src/worker.ts index c0d44f9..910c513 100644 --- a/worker/src/worker.ts +++ b/worker/src/worker.ts @@ -22,9 +22,11 @@ export interface GatewayLike { readonly listeningPort: number | null; start(): Promise; setBrowserEndpoint(endpoint: string): void; + whenEndpointIdle(endpoint: string): Promise; + closeEndpointSessions(endpoint: string, code: number, reason: string): void; closeSession(sessionId: string, code?: number, reason?: string): void; shutdown(): Promise; - on(event: 'sessionClose', listener: (sessionId: string) => void): this; + on(event: 'sessionOpen' | 'sessionClose', listener: (sessionId: string) => void): this; } export interface ControlPlaneLike { @@ -74,8 +76,9 @@ export class BrowserWorker { private gateway: GatewayLike | null = null; private heartbeatTimer: NodeJS.Timeout | null = null; private drainTimer: NodeJS.Timeout | null = null; - private recycleOnDrained = false; private recyclePromise: Promise | null = null; + private retiringBrowser: BrowserServerLike | null = null; + private retirePromise: Promise | null = null; private recycleGeneration = 0; private registerPromise: Promise | null = null; private sessionBudget = 0; @@ -140,19 +143,8 @@ export class BrowserWorker { this.attachBrowserServer(browserServer); this.gateway = this.createGateway(browserServer.wsEndpoint(), this.config.gateway.port); - this.gateway.on('sessionClose', () => { + this.gateway.on('sessionOpen', () => { this.servedSessions += 1; - if (this.currentState === 'draining' && this.gateway?.activeConnectionCount === 0) { - void this.drainCompleted(); - return; - } - // Drain without waiting for the heartbeat to say so: until the - // recycled worker becomes available, the scheduler has already - // stopped selecting this one, so every waiting second is dead - // capacity on small deployments. This counts gateway closes - // while the server counts claims, so a claim that never - // reaches the gateway leaves this counter short — the - // heartbeat's draining status stays as the backstop. if ( this.currentState === 'running' && this.sessionBudget > 0 && @@ -161,6 +153,11 @@ export class BrowserWorker { void this.requestDrain('session budget spent'); } }); + this.gateway.on('sessionClose', () => { + if (this.currentState === 'draining' && this.gateway?.activeConnectionCount === 0) { + void this.drainCompleted(); + } + }); await this.gateway.start(); if (this.shutdownStarted) { return; @@ -184,7 +181,17 @@ export class BrowserWorker { } } + /** + * A signal drain stops new sessions and exits after open sessions end. + * Any other drain request recycles the browser in place. + */ async requestDrain(reason: string, fromSignal = false): Promise { + if (!fromSignal) { + if (this.currentState === 'running') { + await this.recycle(reason); + } + return; + } if (this.currentState === 'shutting_down') { return; } @@ -199,16 +206,8 @@ export class BrowserWorker { } this.currentState = 'draining'; - // A control-plane drain means "give me a fresh browser" (the shipped - // compose files restart exited workers anyway). Recycling in place - // avoids the supervisor's crash-loop backoff, which grows into tens - // of seconds of downtime when sessions turn over quickly. A signal - // drain is a real stop request and still exits. - this.recycleOnDrained = !fromSignal; - this.logger.info('Worker is draining', { reason, recycleOnDrained: this.recycleOnDrained }); - if (fromSignal && this.workerIdValue) { - await this.bestEffortStatus('draining'); - } + this.logger.info('Worker is draining', { reason }); + await this.bestEffortStatus('draining'); if (!this.gateway || this.gateway.activeConnectionCount === 0) { await this.drainCompleted(); return; @@ -227,21 +226,17 @@ export class BrowserWorker { clearTimeout(this.drainTimer); this.drainTimer = null; } - if (this.recycleOnDrained) { - await this.recycle(); - } else { - await this.shutdown(0); - } + await this.shutdown(0); } - private recycle(): Promise { - this.recyclePromise ??= this.runRecycle().finally(() => { + private recycle(reason: string): Promise { + this.recyclePromise ??= this.runRecycle(reason).finally(() => { this.recyclePromise = null; }); return this.recyclePromise; } - private async runRecycle(): Promise { + private async runRecycle(reason: string): Promise { if (this.shutdownStarted) { return; } @@ -249,41 +244,32 @@ export class BrowserWorker { // dispatched before this point must not act after the swap. this.recycleGeneration += 1; this.currentState = 'recycling'; - this.logger.info('Recycling the browser in place'); - // A forced recycle (drain timeout) can still have live sessions; - // close them first so the old browser shuts down promptly. - for (const sessionId of this.gateway?.activeSessionIds ?? []) { - this.gateway?.closeSession(sessionId, 1001, 'worker is recycling'); - } - const oldServer = this.browserServer; - this.browserServer = null; + this.logger.info('Recycling the browser in place', { reason }); try { const workerId = this.workerIdValue; if (!workerId) { throw new Error('worker has no registration to recycle'); } - // Launch the replacement before closing the old browser: sessions - // that dial the gateway mid-recycle reach a live endpoint instead - // of a dead one, at the cost of two browsers briefly coexisting. - const browserServer = await this.launchBrowser(this.config); - if (this.shutdownStarted) { - await this.closeBrowserServer(browserServer); + const keepBrowser = this.retirePromise !== null; + if (keepBrowser) { + this.logger.info('The previous browser still has open sessions; keeping the current browser'); + } else { + const browserServer = await this.launchBrowser(this.config); + if (this.shutdownStarted) { + await this.closeBrowserServer(browserServer); + return; + } + const oldServer = this.browserServer; + this.attachBrowserServer(browserServer); + this.gateway?.setBrowserEndpoint(browserServer.wsEndpoint()); if (oldServer) { - await this.closeBrowserServer(oldServer); + this.retireBrowser(oldServer); } - return; } - this.attachBrowserServer(browserServer); - this.gateway?.setBrowserEndpoint(browserServer.wsEndpoint()); // Reset at the swap, not after register(): sessions served by the // replacement browser while the old one is still closing must // count toward the new budget, not be erased by a later reset. this.servedSessions = 0; - if (oldServer && !(await this.closeBrowserServer(oldServer))) { - // A wedged browser we cannot stop would leak a process per - // recycle; exit instead so the supervisor reclaims them all. - throw new Error('old browser could not be stopped'); - } try { const recycled = await this.controlPlane.recycle( workerId, @@ -321,7 +307,10 @@ export class BrowserWorker { // their late responses distinguishable from post-swap ones. this.recycleGeneration += 1; this.currentState = 'running'; - this.logger.info('Worker recycled with a fresh browser', { workerId: this.workerIdValue }); + this.logger.info( + keepBrowser ? 'Worker renewed its session budget' : 'Worker recycled with a fresh browser', + { workerId: this.workerIdValue }, + ); } catch (error) { if (this.shutdownStarted) { return; @@ -331,6 +320,43 @@ export class BrowserWorker { } } + private retireBrowser(browserServer: BrowserServerLike): void { + this.retiringBrowser = browserServer; + this.retirePromise = this.runRetire(browserServer).finally(() => { + this.retirePromise = null; + }); + } + + private async runRetire(browserServer: BrowserServerLike): Promise { + const endpoint = browserServer.wsEndpoint(); + let timer: NodeJS.Timeout | null = null; + const timedOut = await Promise.race([ + this.gateway ? this.gateway.whenEndpointIdle(endpoint).then(() => false) : Promise.resolve(false), + new Promise(resolve => { + timer = setTimeout(() => resolve(true), this.config.lifecycle.drainTimeoutMs); + timer.unref(); + }), + ]); + if (timer) { + clearTimeout(timer); + } + if (this.shutdownStarted || this.retiringBrowser !== browserServer) { + return; + } + if (timedOut) { + this.logger.warn('Drain timeout expired; closing sessions on the old browser'); + this.gateway?.closeEndpointSessions(endpoint, 1001, 'worker is recycling'); + } + this.retiringBrowser = null; + if (!(await this.closeBrowserServer(browserServer))) { + // A wedged browser we cannot stop would leak a process per + // recycle; exit instead so the supervisor reclaims them all. + await this.fatal(new Error('old browser could not be stopped')); + return; + } + this.logger.info('Old browser closed'); + } + private attachBrowserServer(browserServer: BrowserServerLike): void { this.browserServer = browserServer; browserServer.once('close', () => { @@ -475,6 +501,8 @@ export class BrowserWorker { this.drainTimer = null; } this.logger.info('Worker is shutting down', { exitCode }); + const retiringBrowser = this.retiringBrowser; + this.retiringBrowser = null; try { await this.bestEffortStatus('shutting_down'); @@ -488,6 +516,9 @@ export class BrowserWorker { if (this.browserServer) { await this.closeBrowserServer(this.browserServer); } + if (retiringBrowser) { + await this.closeBrowserServer(retiringBrowser); + } } finally { this.resolveExit(exitCode); this.exit(exitCode); From 55a3f1742650a44ff04e28ea3c17d0f7575c37b0 Mon Sep 17 00:00:00 2001 From: mbroton Date: Fri, 13 Nov 2026 01:55:38 +0100 Subject: [PATCH 2/2] fix(worker): drain open sessions when a signal arrives during a recycle --- worker/src/worker.test.ts | 48 +++++++++++++++++++++++++++++++++++++++ worker/src/worker.ts | 6 ++++- 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/worker/src/worker.test.ts b/worker/src/worker.test.ts index f25ce25..1f3002d 100644 --- a/worker/src/worker.test.ts +++ b/worker/src/worker.test.ts @@ -737,6 +737,54 @@ test('a recycle 409 exits cleanly without registering again', async () => { assert.equal(worker.state, 'shutting_down'); }); +test('a signal during a recycle drains open sessions before the worker exits', async () => { + let releaseRecycle!: () => void; + const recycleGate = new Promise(resolve => { + releaseRecycle = resolve; + }); + let recycleStarted = false; + const statuses: string[] = []; + const controlPlane: ControlPlaneLike = { + register: async () => ({ id: 'worker-1' }), + recycle: async workerId => { + recycleStarted = true; + await recycleGate; + return { id: workerId }; + }, + heartbeat: async () => ({ status: 'available', stale_session_ids: [] }), + setStatus: async (_workerId, status) => { + statuses.push(status); + }, + }; + const gateway = new FakeGateway([activeId]); + const browsers: FakeBrowser[] = []; + const worker = createRecordingWorker(controlPlane, gateway, browsers, 5_000); + await worker.start(); + let exited = false; + void worker.waitForExit().then(() => { + exited = true; + }); + + const recyclePromise = worker.requestDrain('control plane'); + await waitFor(() => recycleStarted); + gateway.addSession(secondId); + await worker.requestDrain('SIGTERM', true); + + assert.equal(worker.state, 'draining'); + assert.deepEqual(statuses, ['draining']); + releaseRecycle(); + await recyclePromise; + assert.equal(worker.state, 'draining'); + assert.deepEqual(statuses, ['draining', 'draining']); + assert.equal(exited, false); + assert.deepEqual(gateway.closedSessions, []); + + gateway.closeSession(activeId); + gateway.closeSession(secondId); + assert.equal(await bounded(worker.waitForExit(), 500), 0); + assert.equal(browsers.every(browser => browser.closed), true); +}); + test('a worker reports recycling during the swap and a drain request then shuts it down', async () => { let releaseRecycle!: () => void; const recycleGate = new Promise(resolve => { diff --git a/worker/src/worker.ts b/worker/src/worker.ts index 910c513..ab38736 100644 --- a/worker/src/worker.ts +++ b/worker/src/worker.ts @@ -200,7 +200,7 @@ export class BrowserWorker { await this.shutdown(0); return; } - if (this.currentState === 'starting' || this.currentState === 'recycling') { + if (this.currentState === 'starting') { await this.shutdown(0); return; } @@ -306,6 +306,10 @@ export class BrowserWorker { // already-bumped generation, so only a bump at completion makes // their late responses distinguishable from post-swap ones. this.recycleGeneration += 1; + if (this.state === 'draining') { + await this.bestEffortStatus('draining'); + return; + } this.currentState = 'running'; this.logger.info( keepBrowser ? 'Worker renewed its session budget' : 'Worker recycled with a fresh browser',