From 9f41edda223fabb4c7bb33840ee55d03f952bdda Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Wed, 28 Dec 2022 09:54:13 +0100
Subject: [PATCH] fix: start tracking module resolution as soon as possible for
 easier tracking (#2560)

* fix: start tracking module resolution as soon as possible for easier tracking

* chore: reduce normalization overhead
---
 packages/vite-node/src/client.ts    | 41 ++++++++++++++++++++++-------
 packages/vite-node/src/types.ts     |  1 +
 packages/vitest/src/utils/import.ts | 20 +++++++-------
 3 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/packages/vite-node/src/client.ts b/packages/vite-node/src/client.ts
index cee1de50f15f..7d7fa70bfc74 100644
--- a/packages/vite-node/src/client.ts
+++ b/packages/vite-node/src/client.ts
@@ -64,21 +64,30 @@ export class ModuleCacheMap extends Map<string, ModuleCache> {
     return this
   }
 
+  setByModuleId(modulePath: string, mod: ModuleCache) {
+    return super.set(modulePath, mod)
+  }
+
   set(fsPath: string, mod: ModuleCache) {
-    fsPath = this.normalizePath(fsPath)
-    return super.set(fsPath, mod)
+    return this.setByModuleId(this.normalizePath(fsPath), mod)
+  }
+
+  getByModuleId(modulePath: string): ModuleCache {
+    if (!super.has(modulePath))
+      super.set(modulePath, {})
+    return super.get(modulePath)!
   }
 
   get(fsPath: string): ModuleCache {
-    fsPath = this.normalizePath(fsPath)
-    if (!super.has(fsPath))
-      super.set(fsPath, {})
-    return super.get(fsPath)!
+    return this.getByModuleId(this.normalizePath(fsPath))
+  }
+
+  deleteByModuleId(modulePath: string): boolean {
+    return super.delete(modulePath)
   }
 
   delete(fsPath: string) {
-    fsPath = this.normalizePath(fsPath)
-    return super.delete(fsPath)
+    return this.deleteByModuleId(this.normalizePath(fsPath))
   }
 
   /**
@@ -197,7 +206,7 @@ export class ViteNodeRunner {
     return !isInternalRequest(id) && !isNodeBuiltin(id)
   }
 
-  async resolveUrl(id: string, importee?: string): Promise<[url: string, fsPath: string]> {
+  private async _resolveUrl(id: string, importee?: string): Promise<[url: string, fsPath: string]> {
     if (!this.shouldResolveId(id))
       return [id, id]
     // we don't pass down importee here, because otherwise Vite doesn't resolve it correctly
@@ -215,6 +224,18 @@ export class ViteNodeRunner {
     return [resolvedId, fsPath]
   }
 
+  async resolveUrl(id: string, importee?: string) {
+    const resolveKey = `resolve:${id}`
+    // put info about new import as soon as possible, so we can start tracking it
+    this.moduleCache.setByModuleId(resolveKey, { resolving: true })
+    try {
+      return await this._resolveUrl(id, importee)
+    }
+    finally {
+      this.moduleCache.deleteByModuleId(resolveKey)
+    }
+  }
+
   /** @internal */
   async dependencyRequest(id: string, fsPath: string, callstack: string[]) {
     const getStack = () => {
@@ -246,7 +267,7 @@ export class ViteNodeRunner {
     const moduleId = normalizeModuleId(fsPath)
     const callstack = [..._callstack, moduleId]
 
-    const mod = this.moduleCache.get(fsPath)
+    const mod = this.moduleCache.getByModuleId(moduleId)
 
     const request = async (dep: string) => {
       const [id, depFsPath] = await this.resolveUrl(dep, fsPath)
diff --git a/packages/vite-node/src/types.ts b/packages/vite-node/src/types.ts
index 606c0f00666a..58d327e673e3 100644
--- a/packages/vite-node/src/types.ts
+++ b/packages/vite-node/src/types.ts
@@ -45,6 +45,7 @@ export interface ModuleCache {
   promise?: Promise<any>
   exports?: any
   evaluated?: boolean
+  resolving?: boolean
   code?: string
   map?: RawSourceMap
   /**
diff --git a/packages/vitest/src/utils/import.ts b/packages/vitest/src/utils/import.ts
index 2c964054d032..16dab1312edf 100644
--- a/packages/vitest/src/utils/import.ts
+++ b/packages/vitest/src/utils/import.ts
@@ -1,21 +1,23 @@
 import { getWorkerState } from './global'
 import { setTimeout } from './timers'
 
-export async function waitForImportsToResolve(tries = 0) {
-  await new Promise(resolve => setTimeout(resolve, 0))
+function waitNextTick() {
+  return new Promise(resolve => setTimeout(resolve, 0))
+}
+
+export async function waitForImportsToResolve() {
+  await waitNextTick()
   const state = getWorkerState()
   const promises: Promise<unknown>[] = []
+  let resolvingCount = 0
   for (const mod of state.moduleCache.values()) {
     if (mod.promise && !mod.evaluated)
       promises.push(mod.promise)
+    if (mod.resolving)
+      resolvingCount++
   }
-  if (!promises.length && tries >= 3)
+  if (!promises.length && !resolvingCount)
     return
   await Promise.allSettled(promises)
-  // wait until the end of the loop, so `.then` on modules is called,
-  // like in import('./example').then(...)
-  // also call dynamicImportSettled again in case new imports were added
-  await new Promise(resolve => setTimeout(resolve, 1))
-    .then(() => Promise.resolve())
-    .then(() => waitForImportsToResolve(tries + 1))
+  await waitForImportsToResolve()
 }