perf(bundled-dev): serve an already-built lazy route instead of recompiling it on reload

packages/vite/src/node/server/bundledDev.ts

@@ -58,6 +58,8 @@ export class MemoryFiles {
   }
 }
 
+const LAZY_PROXY_QUERY = '?rolldown-lazy=1'
+
 export class BundledDev {
   private _devEngine!: DevEngine
   private viteRuntime?: string
@@ -106,6 +108,11 @@ export class BundledDev {
 
   private pendingPayloadFilenames = new Set<string>()
 
+  private builtLazyChunks = new Map<string, string>()
+  private buildGeneration = 0
+  private clientBuildGeneration = new Map<string, number>()
+  private clientsWithLazyPayloads = new Set<string>()
+
   get hasBuildOutput(): boolean {
     return (
       this.memoryFiles.size > 1 ||
@@ -135,6 +142,7 @@ export class BundledDev {
       'vite:client-connected',
       async (payload, client) => {
         this.clients.setupIfNeeded(client, payload.clientId)
+        this.clientBuildGeneration.set(payload.clientId, this.buildGeneration)
         this.devEngine.registerClient(payload.clientId)
       },
     )
@@ -157,6 +165,8 @@ export class BundledDev {
       if (clientId) {
         this.devEngine.removeClient(clientId)
         this.reloadNeededClientIds.delete(clientId)
+        this.clientBuildGeneration.delete(clientId)
+        this.clientsWithLazyPayloads.delete(clientId)
       }
     })
     this.environment.hot.on(
@@ -235,6 +245,19 @@ export class BundledDev {
         this.lastBuildError = null
 
         this.storeOutputFiles(result.output)
+        this.buildGeneration++
+        this.builtLazyChunks.clear()
+        for (const file of result.output) {
+          if (
+            file.type === 'chunk' &&
+            file.isDynamicEntry &&
+            file.facadeModuleId &&
+            // the stub of a never-fetched route is a dynamic entry too
+            !file.facadeModuleId.endsWith(LAZY_PROXY_QUERY)
+          ) {
+            this.builtLazyChunks.set(file.facadeModuleId, file.fileName)
+          }
+        }
 
         // Trigger a full reload if there's no error in the result and a reload is pending from HMR.
         if (this.fullReloadPending) {
@@ -344,6 +367,7 @@ export class BundledDev {
     debug?.(
       `TRIGGER-LAZY: trigger lazy bundling for module ${moduleId} for client ${clientId}`,
     )
+    this.clientsWithLazyPayloads.add(clientId)
     const result = await this.devEngine.compileEntry(moduleId, clientId)
     this.pendingPayloadFilenames.add(result.filename)
     return {
@@ -352,6 +376,32 @@ export class BundledDev {
     }
   }
 
+  /**
+   * A full-build chunk runs and re-registers every module in it, unlike a
+   * lazy payload whose factories skip modules the client already holds. So a
+   * client gets a chunk only while it took no lazy payload and saw no full
+   * build since it connected.
+   *
+   * The build state is read after the running build finished, because an
+   * edit marks the output stale only when its build task ends.
+   */
+  async builtLazyChunk(
+    proxyModuleId: string,
+    clientId: string,
+  ): Promise<string | undefined> {
+    if (this._closed || !this.initialBuildCompleted) return
+    if (!proxyModuleId.endsWith(LAZY_PROXY_QUERY)) return
+    if (this.clientsWithLazyPayloads.has(clientId)) return
+    await this.devEngine.ensureCurrentBuildFinish()
+    const state = await this.devEngine.getBundleState()
+    if (state.hasStaleOutput || state.lastBuildErrored) return
+    if (this.clientBuildGeneration.get(clientId) !== this.buildGeneration)
+      return
+    return this.builtLazyChunks.get(
+      proxyModuleId.slice(0, -LAZY_PROXY_QUERY.length),
+    )
+  }
+
   /**
    * Called when the client reports that it evaluated a payload (the line appended
    * by `payloadDeliveredAck`). Only then is the payload recorded on the server's
@@ -369,6 +419,7 @@ export class BundledDev {
   async close(): Promise<void> {
     this._closed = true
     this.memoryFiles.clear()
+    this.builtLazyChunks.clear()
     await this._devEngine?.close()
     this.initialBuildCompleted = false
   }

packages/vite/src/node/server/middlewares/triggerLazyBundling.ts

@@ -1,3 +1,4 @@
+import path from 'node:path'
 import type { Connect } from '#dep-types/connect'
 import type { ViteDevServer } from '..'
 
@@ -28,6 +29,18 @@ export function triggerLazyBundlingMiddleware(
     const clientId = params.get('clientId')
     let result: { code: string; filename: string } | undefined
     try {
+      const builtChunk =
+        moduleId &&
+        clientId &&
+        (await bundledDev.builtLazyChunk(moduleId, clientId))
+      if (builtChunk) {
+        res.statusCode = 302
+        res.setHeader(
+          'Location',
+          path.posix.join(server.config.base, builtChunk),
+        )
+        return res.end()
+      }
       result = await bundledDev.triggerLazyBundling(moduleId, clientId)
     } catch (e) {
       server.config.logger.error(

playground/lazy-compilation/__tests__/lazy-compilation.spec.ts

@@ -1,6 +1,16 @@
+import path from 'node:path'
 import type { Response, Route } from 'playwright-chromium'
+import { normalizePath } from 'vite'
 import { beforeAll, describe, expect, test } from 'vitest'
-import { isServe, page, promiseWithResolvers } from '~utils'
+import {
+  browser,
+  editFile,
+  isServe,
+  page,
+  promiseWithResolvers,
+  viteServer,
+  viteTestUrl,
+} from '~utils'
 
 // Regression test for rolldown/rolldown#10774.
 //
@@ -16,9 +26,9 @@ import { isServe, page, promiseWithResolvers } from '~utils'
 // response is fetched from the server (so the server finished writing it) but
 // held back from the page until B has been requested and evaluated.
 //
-// Both tests share one page session on purpose. A lazy entry is served through
+// Both tests share one page session on purpose. A lazy entry is compiled by
 // `/@vite/lazy` only until it is first fetched on this server; a later page
-// load gets it as an ordinary chunk. So the second test cannot reload — it
+// load is redirected to its built chunk. So the second test cannot reload — it
 // continues from the state the first test leaves behind.
 
 const factoryFor = (file: string) =>
@@ -112,3 +122,174 @@ describe.runIf(isServe)('lazy compilation', () => {
     expect(bodyC).not.toMatch(factoryFor('shared.js'))
   })
 })
+
+const bundledDev = () => viteServer.environments.client.bundledDev as any
+
+const rebuildLanded = async () => {
+  const state = await bundledDev().devEngine.getBundleState()
+  return !state.hasStaleOutput
+}
+
+// the map is keyed by rolldown's module id, which keeps backslashes on Windows
+const builtChunkFor = (file: string): string | undefined => {
+  const id = path.posix.join(viteServer.config.root, file)
+  for (const [key, chunk] of bundledDev().builtLazyChunks as Map<
+    string,
+    string
+  >) {
+    if (normalizePath(key) === id) return chunk
+  }
+}
+
+const sharedRuns = (p: typeof page) =>
+  p.evaluate(() => (globalThis as any).__sharedRuns as number)
+
+describe.runIf(isServe)('lazy compilation: reload after fetch', () => {
+  test('a reload after the first fetch loads the route from the built chunk', async () => {
+    await expect.poll(rebuildLanded).toBe(true)
+
+    const lazyStatuses: number[] = []
+    const chunkRequests: string[] = []
+    page.on('response', (res) => {
+      if (res.url().includes('/@vite/lazy?')) lazyStatuses.push(res.status())
+    })
+    page.on('request', (req) => {
+      if (/\/assets\/page-a-.*\.js$/.test(req.url()))
+        chunkRequests.push(req.url())
+    })
+
+    await page.reload()
+    await page.click('#route-a-btn')
+    await expect
+      .poll(() => page.textContent('#route-a-content'))
+      .toBe('A:shared-value')
+
+    expect(lazyStatuses).toEqual([302])
+    expect(chunkRequests).toHaveLength(1)
+  })
+
+  test('a page that loaded before an edit gets the edited route', async () => {
+    const context = await browser.newContext()
+    const other = await context.newPage()
+    await other.goto(viteTestUrl)
+    await expect
+      .poll(() => other.textContent('#route-b-content'))
+      .toBe('pending')
+
+    // Wait until the server has seen the edit, so the click is not racing
+    // the watcher. Nothing has run B since the reload, so no rebuild follows;
+    // the output just turns stale and the click compiles from source.
+    await expect.poll(rebuildLanded).toBe(true)
+    const before = builtChunkFor('page-b.js')
+    expect(before).toBeDefined()
+    editFile('page-b.js', (code) => code.replace('B:', 'B-edited:'))
+    await expect
+      .poll(
+        async () =>
+          !(await rebuildLanded()) || builtChunkFor('page-b.js') !== before,
+      )
+      .toBe(true)
+
+    await other.click('#route-b-btn')
+    await expect
+      .poll(() => other.textContent('#route-b-content'))
+      .toBe('B-edited:shared-value')
+    await context.close()
+  })
+
+  test('a page that took a lazy payload keeps taking them, so shared.js runs once', async () => {
+    // With no client that ran A, an edit to A leaves the output stale
+    // instead of rebuilding it.
+    await page.reload()
+    const context = await browser.newContext()
+    const other = await context.newPage()
+    try {
+      await other.goto(viteTestUrl)
+      await expect
+        .poll(() => other.textContent('#route-a-content'))
+        .toBe('pending')
+
+      await expect.poll(rebuildLanded).toBe(true)
+      editFile('page-a.js', (code) => code.replace('A:', 'A-edited:'))
+      await expect.poll(async () => !(await rebuildLanded())).toBe(true)
+      await other.click('#route-a-btn')
+      await expect
+        .poll(() => other.textContent('#route-a-content'))
+        .toBe('A-edited:shared-value')
+      await expect.poll(rebuildLanded).toBe(true)
+      expect(builtChunkFor('page-b.js')).toBeDefined()
+
+      // A built chunk for B would run shared.js a second time
+      await other.click('#route-b-btn')
+      await expect
+        .poll(() => other.textContent('#route-b-content'))
+        .toBe('B-edited:shared-value')
+      expect(await sharedRuns(other)).toBe(1)
+    } finally {
+      await context.close()
+    }
+  })
+
+  test('a click while the edit is still being processed gets the edited route', async () => {
+    // The output turns stale only when the HMR task for the edit ends. The
+    // `hold-page-b` plugin keeps that task inside `transform`.
+    const holdPageB = viteServer.config.plugins.find(
+      (p) => p.name === 'hold-page-b',
+    )!.api as { hold: Promise<void> | null; holding: boolean }
+    const taskHold = promiseWithResolvers<void>()
+    // A redirected chunk is held until the page has seen the update, so the
+    // update cannot repair the page.
+    const chunkHold = promiseWithResolvers<void>()
+    // The page has not run B, so it never fetches the patch; the update
+    // message is the only trace.
+    const updateReceived = promiseWithResolvers<void>()
+
+    await expect.poll(rebuildLanded).toBe(true)
+    const context = await browser.newContext()
+    const held = await context.newPage()
+    try {
+      held.on('websocket', (ws) => {
+        ws.on('framereceived', (frame) => {
+          if (String(frame.payload).includes('bundled-dev-update')) {
+            updateReceived.resolve()
+          }
+        })
+      })
+      await held.goto(viteTestUrl)
+      await expect
+        .poll(() => held.textContent('#route-b-content'))
+        .toBe('pending')
+      await expect.poll(rebuildLanded).toBe(true)
+      expect(builtChunkFor('page-b.js')).toBeDefined()
+      await held.route('**/assets/page-b-*.js', async (route) => {
+        await chunkHold.promise
+        await route.continue()
+      })
+
+      holdPageB.hold = taskHold.promise
+      editFile('page-b.js', (code) =>
+        code.replace('B-edited:', 'B-edited-twice:'),
+      )
+      await expect.poll(() => holdPageB.holding).toBe(true)
+
+      const lazyRequest = held.waitForRequest((req) =>
+        req.url().includes('/@vite/lazy?'),
+      )
+      await held.click('#route-b-btn')
+      await lazyRequest
+
+      taskHold.resolve()
+      await updateReceived.promise
+      chunkHold.resolve()
+
+      await expect
+        .poll(() => held.textContent('#route-b-content'))
+        .toBe('B-edited-twice:shared-value')
+    } finally {
+      holdPageB.hold = null
+      taskHold.resolve()
+      chunkHold.resolve()
+      await context.close()
+    }
+  })
+})

playground/lazy-compilation/shared.js

@@ -1 +1,4 @@
+// how many times this body ran in the page; the specs read it back
+globalThis.__sharedRuns = (globalThis.__sharedRuns ?? 0) + 1
+
 export const shared = 'shared-value'

playground/lazy-compilation/vite.config.ts

@@ -1,8 +1,23 @@
-import { defineConfig } from 'vite'
+import { defineConfig, type Plugin } from 'vite'
+
+// Lets a spec hold the HMR task for page-b.js inside `transform`.
+export const holdPageB = { hold: null as Promise<void> | null, holding: false }
+
+const holdPageBPlugin: Plugin = {
+  name: 'hold-page-b',
+  api: holdPageB,
+  async transform(_code, id) {
+    if (!id.endsWith('/page-b.js') || !holdPageB.hold) return
+    holdPageB.holding = true
+    await holdPageB.hold
+    holdPageB.holding = false
+  },
+}
 
 // Lazy compilation only exists in bundled dev, so this playground always runs
 // with it on, also under plain `pnpm test-serve`.
 export default defineConfig({
+  plugins: [holdPageBPlugin],
   experimental: {
     bundledDev: true,
   },