refactor: make HMR agnostic to environment (#15179)

Author: sheremet-va

packages/vite/src/client/client.ts

@@ -1,6 +1,7 @@
-import type { ErrorPayload, HMRPayload, Update } from 'types/hmrPayload'
-import type { ModuleNamespace, ViteHotContext } from 'types/hot'
+import type { ErrorPayload, HMRPayload } from 'types/hmrPayload'
+import type { ViteHotContext } from 'types/hot'
 import type { InferCustomEventPayload } from 'types/customEvent'
+import { HMRClient, HMRContext } from '../shared/hmr'
 import { ErrorOverlay, overlayId } from './overlay'
 import '@vite/env'
 
@@ -110,17 +111,6 @@ function setupWebSocket(
   return socket
 }
 
-function warnFailedFetch(err: Error, path: string | string[]) {
-  if (!err.message.match('fetch')) {
-    console.error(err)
-  }
-  console.error(
-    `[hmr] Failed to reload ${path}. ` +
-      `This could be due to syntax errors or importing non-existent ` +
-      `modules. (see errors above)`,
-  )
-}
-
 function cleanUrl(pathname: string): string {
   const url = new URL(pathname, location.toString())
   url.searchParams.delete('direct')
@@ -144,6 +134,22 @@ const debounceReload = (time: number) => {
 }
 const pageReload = debounceReload(50)
 
+const hmrClient = new HMRClient(console, async function importUpdatedModule({
+  acceptedPath,
+  timestamp,
+  explicitImportRequired,
+}) {
+  const [acceptedPathWithoutQuery, query] = acceptedPath.split(`?`)
+  return await import(
+    /* @vite-ignore */
+    base +
+      acceptedPathWithoutQuery.slice(1) +
+      `?${explicitImportRequired ? 'import&' : ''}t=${timestamp}${
+        query ? `&${query}` : ''
+      }`
+  )
+})
+
 async function handleMessage(payload: HMRPayload) {
   switch (payload.type) {
     case 'connected':
@@ -173,7 +179,7 @@ async function handleMessage(payload: HMRPayload) {
       await Promise.all(
         payload.updates.map(async (update): Promise<void> => {
           if (update.type === 'js-update') {
-            return queueUpdate(fetchUpdate(update))
+            return queueUpdate(hmrClient.fetchUpdate(update))
           }
 
           // css-update
@@ -245,16 +251,7 @@ async function handleMessage(payload: HMRPayload) {
       break
     case 'prune':
       notifyListeners('vite:beforePrune', payload)
-      // 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.
-      payload.paths.forEach((path) => {
-        const fn = pruneMap.get(path)
-        if (fn) {
-          fn(dataMap.get(path))
-        }
-      })
+      hmrClient.prunePaths(payload.paths)
       break
     case 'error': {
       notifyListeners('vite:error', payload)
@@ -280,10 +277,7 @@ function notifyListeners<T extends string>(
   data: InferCustomEventPayload<T>,
 ): void
 function notifyListeners(event: string, data: any): void {
-  const cbs = customListenersMap.get(event)
-  if (cbs) {
-    cbs.forEach((cb) => cb(data))
-  }
+  hmrClient.notifyListeners(event, data)
 }
 
 const enableOverlay = __HMR_ENABLE_OVERLAY__
@@ -430,206 +424,22 @@ export function removeStyle(id: string): void {
   }
 }
 
-async function fetchUpdate({
-  path,
-  acceptedPath,
-  timestamp,
-  explicitImportRequired,
-}: Update) {
-  const mod = hotModulesMap.get(path)
-  if (!mod) {
-    // In a code-splitting project,
-    // it is common that the hot-updating module is not loaded yet.
-    // https://github.com/vitejs/vite/issues/721
-    return
-  }
-
-  let fetchedModule: ModuleNamespace | undefined
-  const isSelfUpdate = path === acceptedPath
-
-  // determine the qualified callbacks before we re-import the modules
-  const qualifiedCallbacks = mod.callbacks.filter(({ deps }) =>
-    deps.includes(acceptedPath),
-  )
-
-  if (isSelfUpdate || qualifiedCallbacks.length > 0) {
-    const disposer = disposeMap.get(acceptedPath)
-    if (disposer) await disposer(dataMap.get(acceptedPath))
-    const [acceptedPathWithoutQuery, query] = acceptedPath.split(`?`)
-    try {
-      fetchedModule = await import(
-        /* @vite-ignore */
-        base +
-          acceptedPathWithoutQuery.slice(1) +
-          `?${explicitImportRequired ? 'import&' : ''}t=${timestamp}${
-            query ? `&${query}` : ''
-          }`
-      )
-    } catch (e) {
-      warnFailedFetch(e, acceptedPath)
-    }
-  }
-
-  return () => {
-    for (const { deps, fn } of qualifiedCallbacks) {
-      fn(deps.map((dep) => (dep === acceptedPath ? fetchedModule : undefined)))
-    }
-    const loggedPath = isSelfUpdate ? path : `${acceptedPath} via ${path}`
-    console.debug(`[vite] hot updated: ${loggedPath}`)
-  }
-}
-
 function sendMessageBuffer() {
   if (socket.readyState === 1) {
     messageBuffer.forEach((msg) => socket.send(msg))
     messageBuffer.length = 0
   }
 }
 
-interface HotModule {
-  id: string
-  callbacks: HotCallback[]
-}
-
-interface HotCallback {
-  // the dependencies must be fetchable paths
-  deps: string[]
-  fn: (modules: Array<ModuleNamespace | undefined>) => void
-}
-
-type CustomListenersMap = Map<string, ((data: any) => void)[]>
-
-const hotModulesMap = new Map<string, HotModule>()
-const disposeMap = new Map<string, (data: any) => void | Promise<void>>()
-const pruneMap = new Map<string, (data: any) => void | Promise<void>>()
-const dataMap = new Map<string, any>()
-const customListenersMap: CustomListenersMap = new Map()
-const ctxToListenersMap = new Map<string, CustomListenersMap>()
-
 export function createHotContext(ownerPath: string): ViteHotContext {
-  if (!dataMap.has(ownerPath)) {
-    dataMap.set(ownerPath, {})
-  }
-
-  // when a file is hot updated, a new context is created
-  // clear its stale callbacks
-  const mod = hotModulesMap.get(ownerPath)
-  if (mod) {
-    mod.callbacks = []
-  }
-
-  // clear stale custom event listeners
-  const staleListeners = ctxToListenersMap.get(ownerPath)
-  if (staleListeners) {
-    for (const [event, staleFns] of staleListeners) {
-      const listeners = customListenersMap.get(event)
-      if (listeners) {
-        customListenersMap.set(
-          event,
-          listeners.filter((l) => !staleFns.includes(l)),
-        )
-      }
-    }
-  }
-
-  const newListeners: CustomListenersMap = new Map()
-  ctxToListenersMap.set(ownerPath, newListeners)
-
-  function acceptDeps(deps: string[], callback: HotCallback['fn'] = () => {}) {
-    const mod: HotModule = hotModulesMap.get(ownerPath) || {
-      id: ownerPath,
-      callbacks: [],
-    }
-    mod.callbacks.push({
-      deps,
-      fn: callback,
-    })
-    hotModulesMap.set(ownerPath, mod)
-  }
-
-  const hot: ViteHotContext = {
-    get data() {
-      return dataMap.get(ownerPath)
-    },
-
-    accept(deps?: any, callback?: any) {
-      if (typeof deps === 'function' || !deps) {
-        // self-accept: hot.accept(() => {})
-        acceptDeps([ownerPath], ([mod]) => deps?.(mod))
-      } else if (typeof deps === 'string') {
-        // explicit deps
-        acceptDeps([deps], ([mod]) => callback?.(mod))
-      } else if (Array.isArray(deps)) {
-        acceptDeps(deps, callback)
-      } else {
-        throw new Error(`invalid hot.accept() usage.`)
-      }
-    },
-
-    // export names (first arg) are irrelevant on the client side, they're
-    // extracted in the server for propagation
-    acceptExports(_, callback) {
-      acceptDeps([ownerPath], ([mod]) => callback?.(mod))
+  return new HMRContext(ownerPath, hmrClient, {
+    addBuffer(message) {
+      messageBuffer.push(message)
     },
-
-    dispose(cb) {
-      disposeMap.set(ownerPath, cb)
-    },
-
-    prune(cb) {
-      pruneMap.set(ownerPath, cb)
-    },
-
-    // Kept for backward compatibility (#11036)
-    // @ts-expect-error untyped
-    // eslint-disable-next-line @typescript-eslint/no-empty-function
-    decline() {},
-
-    // tell the server to re-perform hmr propagation from this module as root
-    invalidate(message) {
-      notifyListeners('vite:invalidate', { path: ownerPath, message })
-      this.send('vite:invalidate', { path: ownerPath, message })
-      console.debug(
-        `[vite] invalidate ${ownerPath}${message ? `: ${message}` : ''}`,
-      )
-    },
-
-    // custom events
-    on(event, cb) {
-      const addToMap = (map: Map<string, any[]>) => {
-        const existing = map.get(event) || []
-        existing.push(cb)
-        map.set(event, existing)
-      }
-      addToMap(customListenersMap)
-      addToMap(newListeners)
-    },
-
-    // remove a custom event
-    off(event, cb) {
-      const removeFromMap = (map: Map<string, any[]>) => {
-        const existing = map.get(event)
-        if (existing === undefined) {
-          return
-        }
-        const pruned = existing.filter((l) => l !== cb)
-        if (pruned.length === 0) {
-          map.delete(event)
-          return
-        }
-        map.set(event, pruned)
-      }
-      removeFromMap(customListenersMap)
-      removeFromMap(newListeners)
-    },
-
-    send(event, data) {
-      messageBuffer.push(JSON.stringify({ type: 'custom', event, data }))
+    send() {
       sendMessageBuffer()
     },
-  }
-
-  return hot
+  })
 }
 
 /**