From 78ee37c2896b08030892a8f80c5539441e67c03f Mon Sep 17 00:00:00 2001 From: James Date: Fri, 10 Jul 2026 12:25:28 -0700 Subject: [PATCH] refactor(engine): add fingerprinted browser leases --- packages/cli/src/server/studioServer.ts | 31 +- packages/engine/src/index.ts | 4 + .../src/services/browserLeasePool.test.ts | 161 ++++++++++ .../engine/src/services/browserLeasePool.ts | 231 ++++++++++++++ .../src/services/browserManager.test.ts | 26 +- .../engine/src/services/browserManager.ts | 296 +++++++----------- ...frameCapture-constructionOwnership.test.ts | 60 ++++ packages/engine/src/services/frameCapture.ts | 68 +++- 8 files changed, 656 insertions(+), 221 deletions(-) create mode 100644 packages/engine/src/services/browserLeasePool.test.ts create mode 100644 packages/engine/src/services/browserLeasePool.ts create mode 100644 packages/engine/src/services/frameCapture-constructionOwnership.test.ts diff --git a/packages/cli/src/server/studioServer.ts b/packages/cli/src/server/studioServer.ts index 344f2acef0..37eb6b79ff 100644 --- a/packages/cli/src/server/studioServer.ts +++ b/packages/cli/src/server/studioServer.ts @@ -167,12 +167,16 @@ async function downloadRemoteGifImageSources( // Uses the engine's browser pool so the thumbnail browser and render workers // share a single Chrome process instead of running two independent ones. -let _thumbnailBrowser: import("puppeteer-core").Browser | null = null; -let _thumbnailBrowserInitializing: Promise | null = null; +let _thumbnailBrowserLease: import("@hyperframes/engine").BrowserLease | null = null; +let _thumbnailBrowserInitializing: Promise< + import("@hyperframes/engine").BrowserLease | null +> | null = null; async function getThumbnailBrowser(): Promise { - if (_thumbnailBrowser?.connected) return _thumbnailBrowser; - if (_thumbnailBrowserInitializing) return _thumbnailBrowserInitializing; + if (_thumbnailBrowserLease?.browser.connected) return _thumbnailBrowserLease.browser; + if (_thumbnailBrowserInitializing) { + return (await _thumbnailBrowserInitializing)?.browser ?? null; + } _thumbnailBrowserInitializing = (async () => { try { @@ -192,12 +196,12 @@ async function getThumbnailBrowser(): Promise { - _thumbnailBrowser = null; + _thumbnailBrowserLease = acquired; + acquired.browser.on("disconnected", () => { + if (_thumbnailBrowserLease === acquired) _thumbnailBrowserLease = null; _thumbnailBrowserInitializing = null; }); - return _thumbnailBrowser; + return acquired; } catch (err) { console.warn( "[Studio] Failed to launch thumbnail browser:", @@ -208,16 +212,15 @@ async function getThumbnailBrowser(): Promise { - if (!_thumbnailBrowser) return; - const browser = _thumbnailBrowser; - _thumbnailBrowser = null; + if (!_thumbnailBrowserLease) return; + const lease = _thumbnailBrowserLease; + _thumbnailBrowserLease = null; _thumbnailBrowserInitializing = null; - const { releaseBrowser } = await import("@hyperframes/engine"); - await releaseBrowser(browser).catch(() => {}); + await lease.release().catch(() => {}); } // ── Server factory ────────────────────────────────────────────────────────── diff --git a/packages/engine/src/index.ts b/packages/engine/src/index.ts index 1848dea4cd..b00b991637 100644 --- a/packages/engine/src/index.ts +++ b/packages/engine/src/index.ts @@ -73,7 +73,11 @@ export { resolveBrowserGpuMode, buildChromeArgs, ENABLE_BROWSER_POOL, + BrowserLeasePool, type BuildChromeArgsOptions, + type BrowserLaunchFingerprint, + type BrowserLease, + type BrowserPoolState, type CaptureMode, type AcquiredBrowser, } from "./services/browserManager.js"; diff --git a/packages/engine/src/services/browserLeasePool.test.ts b/packages/engine/src/services/browserLeasePool.test.ts new file mode 100644 index 0000000000..692ffe1c1f --- /dev/null +++ b/packages/engine/src/services/browserLeasePool.test.ts @@ -0,0 +1,161 @@ +import { describe, expect, it, vi } from "vitest"; +import type { Browser } from "puppeteer-core"; +import { BrowserLeasePool, type BrowserLaunchFingerprint } from "./browserLeasePool.js"; + +function browser(name: string): Browser { + return { + connected: true, + close: vi.fn().mockResolvedValue(undefined), + disconnect: vi.fn(), + process: () => null, + version: vi.fn().mockResolvedValue(name), + } as unknown as Browser; +} + +function fingerprint(args: string[] = ["--one"]): BrowserLaunchFingerprint { + return { + args, + executablePath: "/chrome", + browserTimeoutMs: 1_000, + protocolTimeoutMs: 2_000, + requestedCaptureMode: "screenshot", + }; +} + +describe("BrowserLeasePool", () => { + it("shares only an exact immutable launch fingerprint", async () => { + const firstBrowser = browser("first"); + const secondBrowser = browser("second"); + const launch = vi + .fn() + .mockResolvedValueOnce({ browser: firstBrowser, captureMode: "screenshot" as const }) + .mockResolvedValueOnce({ browser: secondBrowser, captureMode: "screenshot" as const }); + const pool = new BrowserLeasePool({ + launch, + close: async (value) => value.close(), + forceClose: vi.fn(), + }); + const mutableArgs = ["--one"]; + const first = await pool.acquire(fingerprint(mutableArgs), true); + mutableArgs.push("--mutated-after-acquire"); + const shared = await pool.acquire(fingerprint(), true); + const isolated = await pool.acquire(fingerprint(["--two"]), true); + + expect(shared.browser).toBe(first.browser); + expect(isolated.browser).toBe(secondBrowser); + expect(first.fingerprint.args).toEqual(["--one"]); + expect(Object.isFrozen(first.fingerprint)).toBe(true); + expect(Object.isFrozen(first.fingerprint.args)).toBe(true); + + await Promise.all([first.release(), shared.release(), isolated.release()]); + }); + + it("removes a final lease from availability before awaiting close", async () => { + const firstBrowser = browser("first"); + const secondBrowser = browser("second"); + let finishFirstClose!: () => void; + const firstClose = new Promise((resolve) => { + finishFirstClose = resolve; + }); + const launch = vi + .fn() + .mockResolvedValueOnce({ browser: firstBrowser, captureMode: "screenshot" as const }) + .mockResolvedValueOnce({ browser: secondBrowser, captureMode: "screenshot" as const }); + const close = vi.fn((value: Browser) => + value === firstBrowser ? firstClose : Promise.resolve(), + ); + const pool = new BrowserLeasePool({ launch, close, forceClose: vi.fn() }); + const first = await pool.acquire(fingerprint(), true); + + const releasing = first.release(); + const replacement = await pool.acquire(fingerprint(), true); + + expect(replacement.browser).toBe(secondBrowser); + expect(launch).toHaveBeenCalledTimes(2); + finishFirstClose(); + await releasing; + await replacement.release(); + }); + + it("makes lease release idempotent", async () => { + const value = browser("only"); + const close = vi.fn().mockResolvedValue(undefined); + const pool = new BrowserLeasePool({ + launch: vi.fn().mockResolvedValue({ browser: value, captureMode: "screenshot" }), + close, + forceClose: vi.fn(), + }); + const lease = await pool.acquire(fingerprint(), true); + + await Promise.all([lease.release(), lease.release(), lease.release()]); + + expect(close).toHaveBeenCalledTimes(1); + }); + + it("never reclaims an arbitrary owner through the shared browser handle", async () => { + const value = browser("shared"); + const close = vi.fn().mockResolvedValue(undefined); + const forceClose = vi.fn(); + const pool = new BrowserLeasePool({ + launch: vi.fn().mockResolvedValue({ browser: value, captureMode: "screenshot" }), + close, + forceClose, + }); + const renderLease = await pool.acquire(fingerprint(), true); + const thumbnailLease = await pool.acquire(fingerprint(), true); + + expect(await pool.releaseUnleased(value)).toBe(false); + expect(pool.forceReleaseUnleased(value)).toBe(false); + expect(close).not.toHaveBeenCalled(); + expect(forceClose).not.toHaveBeenCalled(); + + await thumbnailLease.release(); + expect(close).not.toHaveBeenCalled(); + await renderLease.release(); + expect(close).toHaveBeenCalledTimes(1); + }); + + it("evicts a failed launch so the next acquire can recover", async () => { + const recovered = browser("recovered"); + const launch = vi + .fn() + .mockRejectedValueOnce(new Error("launch failed")) + .mockResolvedValueOnce({ browser: recovered, captureMode: "screenshot" as const }); + const pool = new BrowserLeasePool({ + launch, + close: async (value) => value.close(), + forceClose: vi.fn(), + }); + + await expect(pool.acquire(fingerprint(), true)).rejects.toThrow("launch failed"); + const lease = await pool.acquire(fingerprint(), true); + + expect(lease.browser).toBe(recovered); + expect(launch).toHaveBeenCalledTimes(2); + await lease.release(); + }); + + it("drains an in-flight launch without returning the closing browser", async () => { + const value = browser("pending"); + let finishLaunch!: () => void; + const pendingLaunch = new Promise<{ browser: Browser; captureMode: "screenshot" }>( + (resolve) => { + finishLaunch = () => resolve({ browser: value, captureMode: "screenshot" }); + }, + ); + const close = vi.fn().mockResolvedValue(undefined); + const pool = new BrowserLeasePool({ + launch: vi.fn().mockReturnValue(pendingLaunch), + close, + forceClose: vi.fn(), + }); + + const acquiring = pool.acquire(fingerprint(), true); + const draining = pool.drain(); + finishLaunch(); + + await expect(acquiring).rejects.toThrow("drained during acquisition"); + await draining; + expect(close).toHaveBeenCalledWith(value); + }); +}); diff --git a/packages/engine/src/services/browserLeasePool.ts b/packages/engine/src/services/browserLeasePool.ts new file mode 100644 index 0000000000..de11987179 --- /dev/null +++ b/packages/engine/src/services/browserLeasePool.ts @@ -0,0 +1,231 @@ +import type { Browser } from "puppeteer-core"; + +export type CaptureMode = "beginframe" | "screenshot" | "drawelement"; +export type BrowserPoolState = "launching" | "ready" | "closing"; + +export interface BrowserLaunchFingerprint { + readonly args: readonly string[]; + readonly executablePath?: string; + readonly browserTimeoutMs: number; + readonly protocolTimeoutMs: number; + readonly requestedCaptureMode: CaptureMode; +} + +export interface BrowserLaunchResult { + browser: Browser; + captureMode: CaptureMode; +} + +export interface BrowserLease extends BrowserLaunchResult { + readonly fingerprint: Readonly; + release(): Promise; + forceRelease(): void; +} + +interface BrowserPoolEntry { + readonly key: string; + readonly fingerprint: Readonly; + readonly pooled: boolean; + state: BrowserPoolState; + refCount: number; + closeRequested: boolean; + result?: BrowserLaunchResult; + launchPromise: Promise; + closePromise?: Promise; +} + +export interface BrowserLeasePoolOptions { + launch(fingerprint: Readonly): Promise; + close(browser: Browser): Promise; + forceClose(browser: Browser): void; +} + +function freezeFingerprint( + fingerprint: BrowserLaunchFingerprint, +): Readonly { + return Object.freeze({ + ...fingerprint, + args: Object.freeze([...fingerprint.args]), + }); +} + +function fingerprintKey(fingerprint: Readonly): string { + return JSON.stringify([ + fingerprint.args, + fingerprint.executablePath ?? null, + fingerprint.browserTimeoutMs, + fingerprint.protocolTimeoutMs, + fingerprint.requestedCaptureMode, + ]); +} + +/** Owns pooled browser generations and hands callers exactly-once leases. */ +export class BrowserLeasePool { + private readonly available = new Map(); + private readonly entries = new Set(); + private readonly leasesByBrowser = new Map>(); + private drainPromise: Promise | null = null; + + constructor(private readonly options: BrowserLeasePoolOptions) {} + + async acquire( + fingerprintInput: BrowserLaunchFingerprint, + pooled: boolean, + ): Promise { + if (this.drainPromise) await this.drainPromise; + + const fingerprint = freezeFingerprint(fingerprintInput); + const entry = this.reserveEntry(fingerprint, pooled); + return this.issueLease(entry); + } + + private async issueLease(entry: BrowserPoolEntry): Promise { + try { + const result = await entry.launchPromise; + if (entry.state !== "ready") { + throw new Error("Browser pool drained during acquisition"); + } + return this.createLease(entry, result); + } catch (error) { + entry.refCount = Math.max(0, entry.refCount - 1); + throw error; + } + } + + private reserveEntry( + fingerprint: Readonly, + pooled: boolean, + ): BrowserPoolEntry { + const key = fingerprintKey(fingerprint); + let entry = pooled ? this.available.get(key) : undefined; + if (entry?.state === "ready" && !entry.result?.browser.connected) { + this.requestClose(entry, true); + entry = undefined; + } + if (!entry || entry.state === "closing") return this.createEntry(key, fingerprint, pooled); + entry.refCount += 1; + return entry; + } + + /** Close only a browser that is not owned by this pool's active leases. */ + async releaseUnleased(browser: Browser): Promise { + if (this.leasesByBrowser.has(browser)) return false; + await this.options.close(browser).catch(() => {}); + return true; + } + + /** Force-close only a browser that is not owned by an active lease. */ + forceReleaseUnleased(browser: Browser): boolean { + if (this.leasesByBrowser.has(browser)) return false; + this.options.forceClose(browser); + return true; + } + + drain(): Promise { + if (this.drainPromise) return this.drainPromise; + const entries = [...this.entries]; + for (const entry of entries) this.requestClose(entry, false); + this.drainPromise = Promise.all(entries.map((entry) => entry.closePromise)).then(() => { + this.drainPromise = null; + }); + return this.drainPromise; + } + + /** Test-only state reset. Call `drain()` first when entries own real browsers. */ + reset(): void { + this.available.clear(); + this.entries.clear(); + this.leasesByBrowser.clear(); + this.drainPromise = null; + } + + private createEntry( + key: string, + fingerprint: Readonly, + pooled: boolean, + ): BrowserPoolEntry { + const entry: BrowserPoolEntry = { + key, + fingerprint, + pooled, + state: "launching", + refCount: 1, + closeRequested: false, + launchPromise: undefined as unknown as Promise, + }; + entry.launchPromise = this.options.launch(fingerprint).then( + (result) => { + entry.result = result; + if (entry.closeRequested) { + entry.state = "closing"; + } else { + entry.state = "ready"; + } + return result; + }, + (error: unknown) => { + if (this.available.get(key) === entry) this.available.delete(key); + this.entries.delete(entry); + throw error; + }, + ); + this.entries.add(entry); + if (pooled) this.available.set(key, entry); + return entry; + } + + private createLease(entry: BrowserPoolEntry, result: BrowserLaunchResult): BrowserLease { + let active = true; + let lease: BrowserLease; + const deactivate = (force: boolean): boolean => { + if (!active) return false; + active = false; + this.removeLease(result.browser, lease); + entry.refCount = Math.max(0, entry.refCount - 1); + if (entry.refCount === 0) this.requestClose(entry, force); + return true; + }; + lease = { + ...result, + fingerprint: entry.fingerprint, + release: async () => { + if (!deactivate(false)) return; + await entry.closePromise; + }, + forceRelease: () => { + deactivate(true); + }, + }; + const leases = this.leasesByBrowser.get(result.browser) ?? new Set(); + leases.add(lease); + this.leasesByBrowser.set(result.browser, leases); + return lease; + } + + private requestClose(entry: BrowserPoolEntry, force: boolean): void { + if (entry.closePromise) return; + entry.closeRequested = true; + if (this.available.get(entry.key) === entry) this.available.delete(entry.key); + entry.closePromise = entry.launchPromise + .then(async (result) => { + entry.state = "closing"; + if (force) { + this.options.forceClose(result.browser); + } else { + await this.options.close(result.browser).catch(() => {}); + } + }) + .catch(() => {}) + .finally(() => { + this.entries.delete(entry); + if (entry.result) this.leasesByBrowser.delete(entry.result.browser); + }); + } + + private removeLease(browser: Browser, lease: BrowserLease): void { + const leases = this.leasesByBrowser.get(browser); + if (!leases) return; + leases.delete(lease); + if (leases.size === 0) this.leasesByBrowser.delete(browser); + } +} diff --git a/packages/engine/src/services/browserManager.test.ts b/packages/engine/src/services/browserManager.test.ts index cbabfbaf94..86ff57b0f8 100644 --- a/packages/engine/src/services/browserManager.test.ts +++ b/packages/engine/src/services/browserManager.test.ts @@ -10,7 +10,6 @@ import { buildChromeArgs, drainBrowserPool, forceReleaseBrowser, - releaseBrowser, resolveHeadlessShellPath, resolveBrowserGpuMode, } from "./browserManager.js"; @@ -305,8 +304,8 @@ describe("browser pool", () => { expect(first.browser).toBe(second.browser); expect(launchFn).toHaveBeenCalledTimes(1); - await releaseBrowser(first.browser, poolCfg); - await releaseBrowser(second.browser, poolCfg); + await first.release(); + await second.release(); }); it("concurrent acquires via Promise.all trigger exactly one launch", async () => { @@ -320,14 +319,14 @@ describe("browser pool", () => { expect(a.browser).toBe(b.browser); expect(b.browser).toBe(c.browser); - await releaseBrowser(a.browser, poolCfg); - await releaseBrowser(b.browser, poolCfg); - await releaseBrowser(c.browser, poolCfg); + await a.release(); + await b.release(); + await c.release(); }); it("pool recovers from a disconnected browser", async () => { const first = await acquireBrowser(["--no-sandbox"], poolCfg); - await releaseBrowser(first.browser, poolCfg); + await first.release(); // Simulate Chrome crash (first.browser as unknown as { connected: boolean }).connected = false; @@ -340,14 +339,14 @@ describe("browser pool", () => { expect(second.browser).not.toBe(first.browser); expect(launchFn).toHaveBeenCalledTimes(2); - await releaseBrowser(second.browser, poolCfg); + await second.release(); }); it("release at refCount 0 closes the browser", async () => { const result = await acquireBrowser(["--no-sandbox"], poolCfg); const closeFn = result.browser.close as ReturnType; - await releaseBrowser(result.browser, poolCfg); + await result.release(); expect(closeFn).toHaveBeenCalledTimes(1); }); @@ -360,8 +359,8 @@ describe("browser pool", () => { expect(second.browser).toBe(first.browser); expect(launchFn).toHaveBeenCalledTimes(1); - await releaseBrowser(first.browser, poolCfg); - await releaseBrowser(second.browser, poolCfg); + await first.release(); + await second.release(); }); it("forceReleaseBrowser does not kill Chrome when other sessions hold refs", async () => { @@ -375,8 +374,9 @@ describe("browser pool", () => { // Should NOT have disconnected — other session still holds a ref expect(disconnectFn).not.toHaveBeenCalled(); - // Release the remaining ref normally - await releaseBrowser(second.browser, poolCfg); + // Each owner releases its own identity; neither can consume the other. + result.forceRelease(); + await second.release(); }); it("drainBrowserPool is safe to call when no browser is pooled", async () => { diff --git a/packages/engine/src/services/browserManager.ts b/packages/engine/src/services/browserManager.ts index 942f16834c..9644f4d4e4 100644 --- a/packages/engine/src/services/browserManager.ts +++ b/packages/engine/src/services/browserManager.ts @@ -12,6 +12,20 @@ import { join } from "path"; import { homedir } from "os"; import { DEFAULT_CONFIG, type EngineConfig } from "../config.js"; import { getSystemTotalMb, LOW_MEMORY_TOTAL_MB_THRESHOLD } from "./systemMemory.js"; +import { + BrowserLeasePool, + type BrowserLaunchFingerprint, + type BrowserLease, + type CaptureMode, +} from "./browserLeasePool.js"; + +export { BrowserLeasePool } from "./browserLeasePool.js"; +export type { + BrowserLaunchFingerprint, + BrowserLease, + BrowserPoolState, + CaptureMode, +} from "./browserLeasePool.js"; let _puppeteer: PuppeteerNode | undefined; @@ -96,15 +110,7 @@ async function probeHardwareWebGlInfo( } } -// "beginframe" = atomic compositor control via HeadlessExperimental.beginFrame (Linux only) -// "screenshot" = renderSeek + Page.captureScreenshot (all platforms) -// "drawelement" = BeginFrame compositor advance + canvas.drawElementImage capture -export type CaptureMode = "beginframe" | "screenshot" | "drawelement"; - -export interface AcquiredBrowser { - browser: Browser; - captureMode: CaptureMode; -} +export type AcquiredBrowser = BrowserLease; /** * Resolve chrome-headless-shell binary for deterministic BeginFrame rendering. @@ -148,11 +154,6 @@ export function resolveHeadlessShellPath( return undefined; } -let pooledBrowser: Browser | null = null; -let pooledBrowserRefCount = 0; -let pooledCaptureMode: CaptureMode = "screenshot"; -let _pooledBrowserLaunchPromise: Promise | null = null; - // Preserve the producer-era export so re-export shims keep the same public API. export const ENABLE_BROWSER_POOL = DEFAULT_CONFIG.enableBrowserPool; @@ -338,20 +339,30 @@ function logResolvedBrowserGpuMode(resolved: "hardware" | "software", reason: st console.error(`[hyperframes] browserGpuMode auto → ${resolved} (${reason})`); } -/** - * Resolve the capture mode the caller expects, WITHOUT launching a browser. - * Used to validate pool compatibility before returning a cached instance. - */ -function resolveRequestedCaptureMode( - config?: Partial>, -): CaptureMode { - const headlessShell = resolveHeadlessShellPath(config); - // BeginFrame requires chrome-headless-shell AND Linux — crashes on - // macOS/Windows (crbug.com/40656275). - const isLinux = process.platform === "linux"; - const forceScreenshot = config?.forceScreenshot ?? DEFAULT_CONFIG.forceScreenshot; - if (headlessShell && isLinux && !forceScreenshot) return "beginframe"; - return "screenshot"; +function createBrowserLaunchFingerprint( + chromeArgs: string[], + config?: Partial< + Pick + >, +): BrowserLaunchFingerprint { + const launchConfig = { + browserTimeout: DEFAULT_CONFIG.browserTimeout, + protocolTimeout: DEFAULT_CONFIG.protocolTimeout, + forceScreenshot: DEFAULT_CONFIG.forceScreenshot, + ...config, + }; + const headlessShell = resolveHeadlessShellPath(launchConfig); + const requestedCaptureMode: CaptureMode = + headlessShell && process.platform === "linux" && !launchConfig.forceScreenshot + ? "beginframe" + : "screenshot"; + return { + args: chromeArgs, + executablePath: headlessShell, + browserTimeoutMs: launchConfig.browserTimeout, + protocolTimeoutMs: launchConfig.protocolTimeout, + requestedCaptureMode, + }; } export async function acquireBrowser( @@ -364,162 +375,62 @@ export async function acquireBrowser( >, ): Promise { const enablePool = config?.enableBrowserPool ?? DEFAULT_CONFIG.enableBrowserPool; - - if (enablePool && pooledBrowser) { - if (!pooledBrowser.connected) { - pooledBrowser = null; - pooledBrowserRefCount = 0; - _pooledBrowserLaunchPromise = null; - } else { - // Validate mode compatibility: a caller that needs screenshot mode - // (forceScreenshot, alpha output, BeginFrame timeout retry) must not - // receive a beginframe browser — the BeginFrame-only flags make the - // compositor wait for frames the screenshot path never sends. - const requestedMode = resolveRequestedCaptureMode(config); - if (pooledCaptureMode === requestedMode) { - pooledBrowserRefCount += 1; - return { browser: pooledBrowser, captureMode: pooledCaptureMode }; - } - // Mode mismatch — skip pool, launch a dedicated browser for this caller. - // Don't evict the pooled browser: other sessions may still hold refs. - } - } - - // Dedup concurrent launches: when the pool is enabled and multiple callers - // (e.g. parallel workers via Promise.all) race into acquireBrowser before - // the first launch completes, they would all see pooledBrowser === null and - // each spawn a separate Chrome. Cache the in-flight launch Promise so the - // second+ callers await the same one instead of launching again. - if (enablePool && _pooledBrowserLaunchPromise) { - const result = await _pooledBrowserLaunchPromise; - const requestedMode = resolveRequestedCaptureMode(config); - if (result.captureMode === requestedMode) { - pooledBrowserRefCount += 1; - return result; - } - // Mode mismatch with pending launch — launch a dedicated browser. - } - - const launchPromise = launchBrowser(chromeArgs, config); - - if (enablePool && !pooledBrowser && !_pooledBrowserLaunchPromise) { - _pooledBrowserLaunchPromise = launchPromise; - try { - const result = await launchPromise; - pooledBrowser = result.browser; - pooledBrowserRefCount = 1; - pooledCaptureMode = result.captureMode; - return result; - } finally { - _pooledBrowserLaunchPromise = null; - } - } - - return launchPromise; + return browserLeasePool.acquire(createBrowserLaunchFingerprint(chromeArgs, config), enablePool); } // fallow-ignore-next-line complexity async function launchBrowser( - chromeArgs: string[], - config?: Partial< - Pick - >, -): Promise { - // Config chromePath overrides env var / auto-detection. - const headlessShell = resolveHeadlessShellPath(config); - - // BeginFrame requires chrome-headless-shell AND Linux (crashes on - // macOS/Windows — crbug.com/40656275). - const isLinux = process.platform === "linux"; - const forceScreenshot = config?.forceScreenshot ?? DEFAULT_CONFIG.forceScreenshot; - let captureMode: CaptureMode; - let executablePath: string | undefined; - - if (headlessShell && isLinux && !forceScreenshot) { - captureMode = "beginframe"; - executablePath = headlessShell; - } else { - // Screenshot mode with renderSeek: works on all platforms. - captureMode = "screenshot"; - executablePath = headlessShell ?? undefined; - } - + fingerprint: Readonly, +): Promise<{ browser: Browser; captureMode: CaptureMode }> { const ppt = await getPuppeteer(); - const browserTimeout = config?.browserTimeout ?? DEFAULT_CONFIG.browserTimeout; - const protocolTimeout = config?.protocolTimeout ?? DEFAULT_CONFIG.protocolTimeout; - let browser = await ppt.launch({ - headless: true, - args: chromeArgs, - defaultViewport: null, - executablePath, - timeout: browserTimeout, - protocolTimeout, - }); - - const browserVersion = await browser.version().catch(() => "unknown"); - const gpuFlags = chromeArgs.filter( - (a) => a.startsWith("--use-gl=") || a.startsWith("--use-angle="), - ); - console.log( - `[BrowserManager] Browser launched (${browserVersion}, ${captureMode}, gl=${gpuFlags.join(" ") || "default"}, headlessShell=${!!headlessShell}, platform=${process.platform})`, - ); - - if (captureMode === "beginframe") { - const supported = await probeBeginFrameSupport(browser).catch(() => true); - if (!supported) { - await browser.close().catch(() => {}); - console.warn( - "[BrowserManager] HeadlessExperimental.beginFrame unavailable in this Chromium build; falling back to screenshot mode.", - ); - captureMode = "screenshot"; - browser = await ppt.launch({ - headless: true, - args: stripBeginFrameFlags(chromeArgs), - defaultViewport: null, - executablePath, - timeout: browserTimeout, - protocolTimeout, - }); - } - } + let captureMode = fingerprint.requestedCaptureMode; + let browser: Browser | undefined; + try { + browser = await ppt.launch({ + headless: true, + args: [...fingerprint.args], + defaultViewport: null, + executablePath: fingerprint.executablePath, + timeout: fingerprint.browserTimeoutMs, + protocolTimeout: fingerprint.protocolTimeoutMs, + }); - return { browser, captureMode }; -} + const browserVersion = await browser.version().catch(() => "unknown"); + const gpuFlags = fingerprint.args.filter( + (a) => a.startsWith("--use-gl=") || a.startsWith("--use-angle="), + ); + console.log( + `[BrowserManager] Browser launched (${browserVersion}, ${captureMode}, gl=${gpuFlags.join(" ") || "default"}, headlessShell=${!!fingerprint.executablePath}, platform=${process.platform})`, + ); -export async function releaseBrowser( - browser: Browser, - config?: Partial>, -): Promise { - const enablePool = config?.enableBrowserPool ?? DEFAULT_CONFIG.enableBrowserPool; - if (!enablePool) { - await browser.close().catch(() => {}); - return; - } - if (pooledBrowser && pooledBrowser === browser) { - pooledBrowserRefCount = Math.max(0, pooledBrowserRefCount - 1); - if (pooledBrowserRefCount === 0) { - await browser.close().catch(() => {}); - pooledBrowser = null; - _pooledBrowserLaunchPromise = null; + if (captureMode === "beginframe") { + const supported = await probeBeginFrameSupport(browser).catch(() => true); + if (!supported) { + await browser.close().catch(() => {}); + browser = undefined; + console.warn( + "[BrowserManager] HeadlessExperimental.beginFrame unavailable in this Chromium build; falling back to screenshot mode.", + ); + captureMode = "screenshot"; + browser = await ppt.launch({ + headless: true, + args: stripBeginFrameFlags([...fingerprint.args]), + defaultViewport: null, + executablePath: fingerprint.executablePath, + timeout: fingerprint.browserTimeoutMs, + protocolTimeout: fingerprint.protocolTimeoutMs, + }); + } } - return; + + return { browser, captureMode }; + } catch (error) { + await browser?.close().catch(() => {}); + throw error; } - await browser.close().catch(() => {}); } -export function forceReleaseBrowser(browser: Browser): void { - if (pooledBrowser && pooledBrowser === browser) { - // If other sessions still hold refs, just drop ours — don't kill the - // shared Chrome out from under them. The browser will be cleaned up when - // the last session releases or drainBrowserPool is called. - if (pooledBrowserRefCount > 1) { - pooledBrowserRefCount -= 1; - return; - } - pooledBrowserRefCount = 0; - pooledBrowser = null; - _pooledBrowserLaunchPromise = null; - } +function forceCloseBrowserProcess(browser: Browser): void { const proc = ( browser as unknown as { process?: () => { kill: (signal?: NodeJS.Signals) => boolean; killed?: boolean } | null; @@ -539,33 +450,38 @@ export function forceReleaseBrowser(browser: Browser): void { } } +const browserLeasePool = new BrowserLeasePool({ + launch: launchBrowser, + close: async (browser) => browser.close(), + forceClose: forceCloseBrowserProcess, +}); + +export async function releaseBrowser( + browser: Browser, + _config?: Partial>, +): Promise { + // Backward-compatible escape hatch for browsers created outside this pool. + // A pooled browser must be released through its BrowserLease; reclaiming an + // arbitrary lease by shared Browser identity can release another owner. + await browserLeasePool.releaseUnleased(browser); +} + +export function forceReleaseBrowser(browser: Browser): void { + browserLeasePool.forceReleaseUnleased(browser); +} + /** * Forcefully close the pooled browser if one exists, regardless of refCount. * Used for explicit cleanup at process exit or between independent render jobs * that should not share browser state. */ export async function drainBrowserPool(): Promise { - // Await any in-flight launch first — otherwise the launch resolves after we - // drain and produces a browser that nobody references (orphan). - const pending = _pooledBrowserLaunchPromise; - _pooledBrowserLaunchPromise = null; - if (pending) { - await pending.then((r) => r.browser.close()).catch(() => {}); - } - if (pooledBrowser) { - const browser = pooledBrowser; - pooledBrowser = null; - pooledBrowserRefCount = 0; - await browser.close().catch(() => {}); - } + await browserLeasePool.drain(); } /** Test-only: reset all pool state. */ export function _resetBrowserPoolForTests(): void { - pooledBrowser = null; - pooledBrowserRefCount = 0; - pooledCaptureMode = "screenshot"; - _pooledBrowserLaunchPromise = null; + browserLeasePool.reset(); } /** Test-only: inject a mock PuppeteerNode so tests bypass the dynamic import. */ diff --git a/packages/engine/src/services/frameCapture-constructionOwnership.test.ts b/packages/engine/src/services/frameCapture-constructionOwnership.test.ts new file mode 100644 index 0000000000..ccd3fbecf1 --- /dev/null +++ b/packages/engine/src/services/frameCapture-constructionOwnership.test.ts @@ -0,0 +1,60 @@ +// @vitest-environment node +import { afterEach, describe, expect, it, vi } from "vitest"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { Browser, Page, PuppeteerNode } from "puppeteer-core"; +import { + _resetBrowserPoolForTests, + _setPuppeteerForTests, + drainBrowserPool, +} from "./browserManager.js"; +import { createCaptureSession } from "./frameCapture.js"; + +describe("createCaptureSession construction ownership", () => { + afterEach(async () => { + await drainBrowserPool(); + _resetBrowserPoolForTests(); + _setPuppeteerForTests(undefined); + }); + + it("closes the page and releases its exact browser lease when bootstrap fails", async () => { + const outputDir = mkdtempSync(join(tmpdir(), "hf-session-owner-")); + const page = { + evaluateOnNewDocument: vi.fn().mockRejectedValue(new Error("bootstrap failed")), + close: vi.fn().mockResolvedValue(undefined), + } as unknown as Page; + const browser = { + connected: true, + newPage: vi.fn().mockResolvedValue(page), + version: vi.fn().mockResolvedValue("HeadlessChrome/150.0.0.0"), + close: vi.fn().mockResolvedValue(undefined), + disconnect: vi.fn(), + process: () => null, + } as unknown as Browser; + _setPuppeteerForTests({ + launch: vi.fn().mockResolvedValue(browser), + } as unknown as PuppeteerNode); + + try { + await expect( + createCaptureSession( + "http://127.0.0.1:3000", + outputDir, + { width: 320, height: 180, fps: { num: 30, den: 1 }, format: "jpeg" }, + null, + { + browserGpuMode: "software", + enableBrowserPool: true, + forceScreenshot: true, + }, + ), + ).rejects.toThrow("bootstrap failed"); + + expect(page.close).toHaveBeenCalledTimes(1); + expect(browser.close).toHaveBeenCalledTimes(1); + } finally { + rmSync(outputDir, { recursive: true, force: true }); + } + }); +}); diff --git a/packages/engine/src/services/frameCapture.ts b/packages/engine/src/services/frameCapture.ts index 0bb503fb40..044e6250b7 100644 --- a/packages/engine/src/services/frameCapture.ts +++ b/packages/engine/src/services/frameCapture.ts @@ -21,6 +21,7 @@ import { buildChromeArgs, resolveBrowserGpuMode, resolveHeadlessShellPath, + type BrowserLease, type CaptureMode, } from "./browserManager.js"; import { @@ -61,6 +62,8 @@ export type BeforeCaptureHook = (page: Page, time: number) => Promise; export interface CaptureSession { browser: Browser; + /** Exact ownership token for this browser acquisition. */ + browserLease?: BrowserLease; page: Page; options: CaptureOptions; serverUrl: string; @@ -829,9 +832,63 @@ export async function createCaptureSession( { ...config, browserGpuMode: resolvedGpuMode }, ); - const { browser, captureMode } = await acquireBrowser(chromeArgs, config); + const browserLease = await acquireBrowser(chromeArgs, config); + return constructCaptureSessionWithRollback({ + browserLease, + serverUrl, + outputDir, + options, + onBeforeCapture, + config, + useDrawElement, + }); +} + +interface CaptureSessionConstructionInput { + browserLease: BrowserLease; + serverUrl: string; + outputDir: string; + options: CaptureOptions; + onBeforeCapture: BeforeCaptureHook | null; + config?: Partial; + useDrawElement: boolean; +} + +async function constructCaptureSessionWithRollback( + input: CaptureSessionConstructionInput, +): Promise { + let page: Page | undefined; + try { + return await constructCaptureSession({ + ...input, + onPageCreated: (createdPage) => { + page = createdPage; + }, + }); + } catch (error) { + await page?.close().catch(() => {}); + await input.browserLease.release(); + throw error; + } +} + +async function constructCaptureSession( + input: CaptureSessionConstructionInput & { onPageCreated(page: Page): void }, +): Promise { + const { + browserLease, + serverUrl, + outputDir, + options, + onBeforeCapture, + config, + useDrawElement, + onPageCreated, + } = input; + const { browser, captureMode } = browserLease; const page = await browser.newPage(); + onPageCreated(page); // Polyfill esbuild's keepNames helper inside the page. // // The engine is published as raw TypeScript (`packages/engine/package.json` @@ -941,6 +998,7 @@ export async function createCaptureSession( return { browser, + browserLease, page, options: sessionOptions, serverUrl, @@ -3148,18 +3206,20 @@ export async function closeCaptureSession(session: CaptureSession): Promise