fix(worker): keep taking sessions while a recycled browser finishes slow ones

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
 

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 {

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).

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`. |
 

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
 

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<IncomingMessage>;
     nextUpstreamSocket: () => Promise<WebSocket>;
     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<void> | 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));

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<GatewayEvents> {
     private readonly webSocketServer = new WebSocketServer({ noServer: true });
     private readonly sessions = new Map<string, SessionSockets>();
     private readonly pendingSessions = new Map<string, PendingSession>();
+    private readonly idleWaiters: Array<{ endpoint: string; resolve: () => void }> = [];
     private shuttingDown = false;
     private shutdownPromise: Promise<void> | null = null;
     private listenPort: number | null = null;
@@ -54,6 +58,22 @@ export class SessionGateway extends EventEmitter<GatewayEvents> {
         this.browserEndpoint = endpoint;
     }
 
+    /** Resolves when no open or connecting session uses the browser endpoint. */
+    whenEndpointIdle(endpoint: string): Promise<void> {
+        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<GatewayEvents> {
                 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<GatewayEvents> {
             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<GatewayEvents> {
                 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<GatewayEvents> {
             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<GatewayEvents> {
             }
             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<GatewayEvents> {
                     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<GatewayEvents> {
             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 {