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..1f3002d 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); }); @@ -614,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 => { @@ -838,6 +1009,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 +1036,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 +1092,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 +1141,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 +1163,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 +1260,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 +1306,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 +1323,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 +1342,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 +1388,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 +1455,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 +1465,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..ab38736 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; } @@ -193,22 +200,14 @@ export class BrowserWorker { await this.shutdown(0); return; } - if (this.currentState === 'starting' || this.currentState === 'recycling') { + if (this.currentState === 'starting') { await this.shutdown(0); return; } 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, @@ -320,8 +306,15 @@ 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('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 +324,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 +505,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 +520,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);