fix(hmr): call dispose before prune (#15782)

Author: bluwy

packages/vite/src/client/client.ts

@@ -268,7 +268,7 @@ async function handleMessage(payload: HMRPayload) {
       break
     case 'prune':
       notifyListeners('vite:beforePrune', payload)
-      hmrClient.prunePaths(payload.paths)
+      await hmrClient.prunePaths(payload.paths)
       break
     case 'error': {
       notifyListeners('vite:error', payload)

packages/vite/src/runtime/hmrHandler.ts

@@ -67,7 +67,7 @@ export async function handleHMRPayload(
     }
     case 'prune':
       await hmrClient.notifyListeners('vite:beforePrune', payload)
-      hmrClient.prunePaths(payload.paths)
+      await hmrClient.prunePaths(payload.paths)
       break
     case 'error': {
       await hmrClient.notifyListeners('vite:error', payload)

packages/vite/src/shared/hmr.ts

@@ -232,8 +232,13 @@ export class HMRClient {
   // After an HMR update, some modules are no longer imported on the page
   // but they may have left behind side effects that need to be cleaned up
   // (.e.g style injections)
-  // TODO Trigger their dispose callbacks.
-  public prunePaths(paths: string[]): void {
+  public async prunePaths(paths: string[]): Promise<void> {
+    await Promise.all(
+      paths.map((path) => {
+        const disposer = this.disposeMap.get(path)
+        if (disposer) return disposer(this.dataMap.get(path))
+      }),
+    )
     paths.forEach((path) => {
       const fn = this.pruneMap.get(path)
       if (fn) {

playground/hmr-ssr/__tests__/hmr.spec.ts renamed to playground/hmr-ssr/__tests__/hmr-ssr.spec.ts

@@ -7,7 +7,14 @@ import type { InlineConfig, Logger, ViteDevServer } from 'vite'
 import { createServer, createViteRuntime } from 'vite'
 import type { ViteRuntime } from 'vite/runtime'
 import type { RollupError } from 'rollup'
-import { page, promiseWithResolvers, slash, untilUpdated } from '~utils'
+import {
+  addFile,
+  page,
+  promiseWithResolvers,
+  readFile,
+  slash,
+  untilUpdated,
+} from '~utils'
 
 let server: ViteDevServer
 const clientLogs: string[] = []
@@ -737,31 +744,19 @@ test.todo('should hmr when file is deleted and restored', async () => {
   )
   await untilUpdated(() => hmr('.file-delete-restore'), 'parent:child1')
 
+  // delete the file
   editFile(parentFile, (code) =>
     code.replace(
       "export { value as childValue } from './child'",
       "export const childValue = 'not-child'",
     ),
   )
+  const originalChildFileCode = readFile(childFile)
   removeFile(childFile)
   await untilUpdated(() => hmr('.file-delete-restore'), 'parent:not-child')
 
-  createFile(
-    childFile,
-    `
-import { rerender } from './runtime'
-
-export const value = 'child'
-
-if (import.meta.hot) {
-  import.meta.hot.accept((newMod) => {
-    if (!newMod) return
-
-    rerender({ child: newMod.value })
-  })
-}
-`,
-  )
+  // restore the file
+  createFile(childFile, originalChildFileCode)
   editFile(parentFile, (code) =>
     code.replace(
       "export const childValue = 'not-child'",
@@ -822,6 +817,45 @@ test.todo('delete file should not break hmr', async () => {
   )
 })
 
+test.todo(
+  'deleted file should trigger dispose and prune callbacks',
+  async () => {
+    await setupViteRuntime('/hmr.ts')
+
+    const parentFile = 'file-delete-restore/parent.js'
+    const childFile = 'file-delete-restore/child.js'
+
+    // delete the file
+    editFile(parentFile, (code) =>
+      code.replace(
+        "export { value as childValue } from './child'",
+        "export const childValue = 'not-child'",
+      ),
+    )
+    const originalChildFileCode = readFile(childFile)
+    removeFile(childFile)
+    await untilUpdated(
+      () => page.textContent('.file-delete-restore'),
+      'parent:not-child',
+    )
+    expect(clientLogs).to.include('file-delete-restore/child.js is disposed')
+    expect(clientLogs).to.include('file-delete-restore/child.js is pruned')
+
+    // restore the file
+    addFile(childFile, originalChildFileCode)
+    editFile(parentFile, (code) =>
+      code.replace(
+        "export const childValue = 'not-child'",
+        "export { value as childValue } from './child'",
+      ),
+    )
+    await untilUpdated(
+      () => page.textContent('.file-delete-restore'),
+      'parent:child',
+    )
+  },
+)
+
 test('import.meta.hot?.accept', async () => {
   await setupViteRuntime('/hmr.ts')
   await untilConsoleLogAfter(

playground/hmr/__tests__/hmr.spec.ts

@@ -8,6 +8,7 @@ import {
   getColor,
   isBuild,
   page,
+  readFile,
   removeFile,
   serverLogs,
   untilBrowserLogAfter,
@@ -784,34 +785,21 @@ if (!isBuild) {
       'parent:child1',
     )
 
+    // delete the file
     editFile(parentFile, (code) =>
       code.replace(
         "export { value as childValue } from './child'",
         "export const childValue = 'not-child'",
       ),
     )
+    const originalChildFileCode = readFile(childFile)
     removeFile(childFile)
     await untilUpdated(
       () => page.textContent('.file-delete-restore'),
       'parent:not-child',
     )
 
-    addFile(
-      childFile,
-      `
-import { rerender } from './runtime'
-
-export const value = 'child'
-
-if (import.meta.hot) {
-  import.meta.hot.accept((newMod) => {
-    if (!newMod) return
-
-    rerender({ child: newMod.value })
-  })
-}
-`,
-    )
+    addFile(childFile, originalChildFileCode)
     editFile(parentFile, (code) =>
       code.replace(
         "export const childValue = 'not-child'",
@@ -875,6 +863,42 @@ if (import.meta.hot) {
     )
   })
 
+  test('deleted file should trigger dispose and prune callbacks', async () => {
+    await page.goto(viteTestUrl)
+
+    const parentFile = 'file-delete-restore/parent.js'
+    const childFile = 'file-delete-restore/child.js'
+
+    // delete the file
+    editFile(parentFile, (code) =>
+      code.replace(
+        "export { value as childValue } from './child'",
+        "export const childValue = 'not-child'",
+      ),
+    )
+    const originalChildFileCode = readFile(childFile)
+    removeFile(childFile)
+    await untilUpdated(
+      () => page.textContent('.file-delete-restore'),
+      'parent:not-child',
+    )
+    expect(browserLogs).to.include('file-delete-restore/child.js is disposed')
+    expect(browserLogs).to.include('file-delete-restore/child.js is pruned')
+
+    // restore the file
+    addFile(childFile, originalChildFileCode)
+    editFile(parentFile, (code) =>
+      code.replace(
+        "export const childValue = 'not-child'",
+        "export { value as childValue } from './child'",
+      ),
+    )
+    await untilUpdated(
+      () => page.textContent('.file-delete-restore'),
+      'parent:child',
+    )
+  })
+
   test('import.meta.hot?.accept', async () => {
     const el = await page.$('.optional-chaining')
     await untilBrowserLogAfter(

playground/hmr/file-delete-restore/child.js

@@ -8,4 +8,12 @@ if (import.meta.hot) {
 
     rerender({ child: newMod.value })
   })
+
+  import.meta.hot.dispose(() => {
+    console.log('file-delete-restore/child.js is disposed')
+  })
+
+  import.meta.hot.prune(() => {
+    console.log('file-delete-restore/child.js is pruned')
+  })
 }