fix(browser): close the pages of a project that has no tests left

Browser pages were only closed by the provider's `close()`, which runs once the
whole workspace run is over. Projects run concurrently and each opens up to
`maxWorkers` pages, so a run held `projects * maxWorkers` pages at its peak and
the projects that finished early kept their page alive until the slowest one
was done.

Add an optional `closePage` to `BrowserProvider`, implement it for Playwright,
and call it from the pool once a session has no test files left. The results of
the last test can still be on their way when that happens, so the session waits
for them to be handled before its page goes away.

In watch mode the pages are kept so that reruns still reuse the session.

Author: JCQuintas

packages/browser-playwright/src/playwright.ts

@@ -653,6 +653,19 @@ export class PlaywrightBrowserProvider implements BrowserProvider {
     await this._throwIfClosing(browserPage)
   }
 
+  async closePage(sessionId: string): Promise<void> {
+    const page = this.pages.get(sessionId)
+
+    if (!page) {
+      return
+    }
+
+    debug?.('[%s][%s] closing the page', sessionId, this.browserName)
+    this.pages.delete(sessionId)
+    await page.close()
+    // the context is closed with the provider, it can still hold pending tracing chunks
+  }
+
   private async _throwIfClosing(disposable?: { close: () => Promise<void> }) {
     if (this.closing) {
       debug?.('[%s] provider was closed, cannot perform the action on %s', this.browserName, String(disposable))

packages/browser/src/node/rpc.ts

@@ -291,11 +291,18 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject, defaultMocke
           return project.benchmark.writeResult(relativePath, data)
         },
         async onTaskUpdate(method, packs, events) {
-          if (method === 'collect') {
-            vitest.state.updateTasks(packs)
+          const sessions = vitest._browserSessions
+          sessions.startUpdate(options.sessionId)
+          try {
+            if (method === 'collect') {
+              vitest.state.updateTasks(packs)
+            }
+            else {
+              await vitest._testRun.updated(packs, events)
+            }
           }
-          else {
-            await vitest._testRun.updated(packs, events)
+          finally {
+            sessions.finishUpdate(options.sessionId)
           }
         },
         onAfterSuiteRun(meta) {

packages/vitest/src/node/browser/sessions.ts

@@ -5,13 +5,48 @@ import { createDefer } from '@vitest/utils/helpers'
 
 export class BrowserSessions {
   private sessions = new Map<string, BrowserServerStateSession>()
+  private pendingUpdates = new Map<string, { count: number; waiters: (() => void)[] }>()
 
   public sessionIds: Set<string> = new Set()
 
   getSession(sessionId: string): BrowserServerStateSession | undefined {
     return this.sessions.get(sessionId)
   }
 
+  startUpdate(sessionId: string): void {
+    const pending = this.pendingUpdates.get(sessionId)
+    if (pending) {
+      pending.count++
+    }
+    else {
+      this.pendingUpdates.set(sessionId, { count: 1, waiters: [] })
+    }
+  }
+
+  finishUpdate(sessionId: string): void {
+    const pending = this.pendingUpdates.get(sessionId)
+    if (!pending) {
+      return
+    }
+    pending.count--
+    if (pending.count > 0) {
+      return
+    }
+    this.pendingUpdates.delete(sessionId)
+    pending.waiters.forEach(resolve => resolve())
+  }
+
+  // the results of the last test can still be arriving when the session has no tests left
+  waitForUpdates(sessionId: string): Promise<void> {
+    const pending = this.pendingUpdates.get(sessionId)
+    if (!pending) {
+      return Promise.resolve()
+    }
+    return new Promise<void>((resolve) => {
+      pending.waiters.push(resolve)
+    })
+  }
+
   destroySession(sessionId: string): void {
     this.sessions.delete(sessionId)
   }

packages/vitest/src/node/pools/browser.ts

@@ -37,6 +37,7 @@ export function createBrowserPool(vitest: Vitest): ProcessPool {
     : Math.max(maxThreadsCount, 1)
 
   const projectPools = new WeakMap<TestProject, BrowserPool>()
+  const pools = new Set<BrowserPool>()
 
   const ensurePool = (project: TestProject) => {
     if (projectPools.has(project)) {
@@ -49,6 +50,7 @@ export function createBrowserPool(vitest: Vitest): ProcessPool {
       maxWorkers: getThreadsCount(project),
     })
     projectPools.set(project, pool)
+    pools.add(pool)
     vitest.onCancel(() => {
       pool.cancel()
     })
@@ -170,21 +172,31 @@ export function createBrowserPool(vitest: Vitest): ProcessPool {
       // a frozen or crashed browser never answers the close message;
       // don't wait for it forever, the browser process is killed
       // when this process exits anyway
-      await Promise.all(Array.from(providers, (provider) => {
+      const withTimeout = (promise: Promise<void>, message: string) => {
         let timer: ReturnType<typeof setTimeout>
         return Promise.race([
-          Promise.resolve(provider.close()).finally(() => clearTimeout(timer)),
+          promise.finally(() => clearTimeout(timer)),
           new Promise<void>((resolve) => {
             timer = setTimeout(() => {
-              vitest.logger.warn(`The browser did not close within ${PROVIDER_CLOSE_TIMEOUT}ms. The browser process will be killed when the process exits.`)
+              vitest.logger.warn(message)
               resolve()
             }, PROVIDER_CLOSE_TIMEOUT)
             timer.unref()
           }),
         ])
-      }))
+      }
+
+      await Promise.all(Array.from(pools, pool => withTimeout(
+        pool.waitForClosedPages(),
+        `The browser did not close its pages within ${PROVIDER_CLOSE_TIMEOUT}ms. The browser process will be killed when the process exits.`,
+      )))
+      await Promise.all(Array.from(providers, provider => withTimeout(
+        Promise.resolve(provider.close()),
+        `The browser did not close within ${PROVIDER_CLOSE_TIMEOUT}ms. The browser process will be killed when the process exits.`,
+      )))
       vitest._browserSessions.sessionIds.clear()
       providers.clear()
+      pools.clear()
       vitest.projects.forEach((project) => {
         project.browser?.state.orchestrators.forEach((orchestrator) => {
           orchestrator.$close()
@@ -207,6 +219,7 @@ class BrowserPool {
   private _providedContext: string | undefined
 
   private readySessions: Set<string>
+  private _closingPages: Promise<void>[] = []
 
   private _traces: Traces
   private _otel: {
@@ -344,11 +357,44 @@ class BrowserPool {
     return orchestrator
   }
 
+  public async waitForClosedPages(): Promise<void> {
+    await Promise.all(this._closingPages)
+  }
+
+  // in watch mode the pages are reused by the next run
+  private closePages(): void {
+    const provider = this.project.browser!.provider
+
+    if (this.project.vitest.config.watch || !provider.closePage) {
+      return
+    }
+
+    const sessionIds = [...this.orchestrators.keys()]
+    sessionIds.forEach((sessionId) => {
+      this.readySessions.delete(sessionId)
+      this.orchestrators.delete(sessionId)
+    })
+
+    // not awaited here: a browser that stopped answering would hold up the whole run,
+    // `close` waits for it instead
+    this._closingPages.push(Promise.all(sessionIds.map(async (sessionId) => {
+      try {
+        // the page cannot go away while its results are still being handled
+        await this.project.vitest._browserSessions.waitForUpdates(sessionId)
+        await provider.closePage!(sessionId)
+      }
+      catch (error) {
+        debug?.('[%s] failed to close the page: %s', sessionId, error)
+      }
+    })).then(() => undefined))
+  }
+
   private finishSession(sessionId: string): void {
     this.readySessions.add(sessionId)
 
     // the last worker finished running tests
     if (this.readySessions.size === this.orchestrators.size) {
+      this.closePages()
       this._otel.span.end()
       this._promise?.resolve()
       this._promise = undefined

packages/vitest/src/node/types/browser.ts

@@ -52,6 +52,7 @@ export interface BrowserProvider {
   supportsParallelism: boolean
   getCommandsContext: (sessionId: string) => Record<string, unknown>
   openPage: (sessionId: string, url: string, options: { parallel: boolean }) => Promise<void>
+  closePage?: (sessionId: string) => Awaitable<void>
   getCDPSession?: (sessionId: string) => Promise<CDPSession>
   close: () => Awaitable<void>
 }