feat: batch package info fetching (#67)

9romise · web-flow · commit 7a75949c1110 · 2026-03-07T13:25:40.000+08:00
* refactor: use `getPackageInfo` directly in document-link

* feat: add batch runner

* chore: improve log

* feat: support `maxSize`
diff --git a/src/providers/document-link/npmx.ts b/src/providers/document-link/npmx.ts
@@ -1,49 +1,19 @@
 import type { Extractor } from '#types/extractor'
 import type { DocumentLink, DocumentLinkProvider, TextDocument } from 'vscode'
-import { config, logger } from '#state'
+import { config } from '#state'
 import { getPackageInfo } from '#utils/api/package'
 import { npmxPackageUrl } from '#utils/links'
 import { resolveExactVersion } from '#utils/package'
 import { isSupportedProtocol, parseVersion } from '#utils/version'
 import { Uri, DocumentLink as VscodeDocumentLink } from 'vscode'
 
-// Limit concurrent lookups to avoid overwhelming the registry and hitting rate limits
-const RESOLVED_LOOKUP_CONCURRENCY = 6
-
-type PackageInfoResult = Awaited<ReturnType<typeof getPackageInfo>>
-
 export class NpmxDocumentLinkProvider<T extends Extractor> implements DocumentLinkProvider {
   extractor: T
 
   constructor(extractor: T) {
     this.extractor = extractor
   }
 
-  private async fetchResolvedPackageInfoMap(names: string[]): Promise<Map<string, PackageInfoResult>> {
-    const packageInfoMap = new Map<string, PackageInfoResult>()
-
-    for (let i = 0; i < names.length; i += RESOLVED_LOOKUP_CONCURRENCY) {
-      const batch = names.slice(i, i + RESOLVED_LOOKUP_CONCURRENCY)
-      const results = await Promise.allSettled(batch.map(async (name) => [name, await getPackageInfo(name)] as const))
-
-      for (const result of results) {
-        if (result.status === 'fulfilled') {
-          const [name, pkg] = result.value
-          packageInfoMap.set(name, pkg)
-        }
-      }
-
-      for (const [index, result] of results.entries()) {
-        if (result.status === 'rejected') {
-          const name = batch[index]
-          logger.warn(`[package-link] failed to fetch package info for ${name}: ${String(result.reason)}`)
-        }
-      }
-    }
-
-    return packageInfoMap
-  }
-
   async provideDocumentLinks(document: TextDocument): Promise<DocumentLink[]> {
     const root = this.extractor.parse(document)
     if (!root)
@@ -67,13 +37,6 @@ export class NpmxDocumentLinkProvider<T extends Extractor> implements DocumentLi
       parsedDeps.push({ dep, parsed })
     }
 
-    let packageInfoMap = new Map<string, PackageInfoResult>()
-
-    if (linkMode === 'resolved') {
-      const names = [...new Set(parsedDeps.map(({ dep }) => dep.name))]
-      packageInfoMap = await this.fetchResolvedPackageInfoMap(names)
-    }
-
     for (const { dep, parsed } of parsedDeps) {
       const { name, nameNode } = dep
 
@@ -82,7 +45,7 @@ export class NpmxDocumentLinkProvider<T extends Extractor> implements DocumentLi
       if (linkMode === 'declared') {
         targetVersion = parsed.version
       } else if (linkMode === 'resolved') {
-        const pkg = packageInfoMap.get(name)
+        const pkg = await getPackageInfo(name)
         const exactVersion = pkg ? resolveExactVersion(pkg, parsed.version) : null
         targetVersion = exactVersion ?? parsed.version
       }
diff --git a/src/utils/api/package.ts b/src/utils/api/package.ts
@@ -1,27 +1,18 @@
-import type { PackageVersionsInfoWithMetadata } from 'fast-npm-meta'
+import type { MaybeError, PackageVersionsInfoWithMetadata } from 'fast-npm-meta'
 import { logger } from '#state'
-import { getVersions } from 'fast-npm-meta'
+import { createBatchRunner } from '#utils/batch'
+import { getVersionsBatch } from 'fast-npm-meta'
 import { memoize } from '../memoize'
 
+const BATCH_SIZE = 20
+
 export interface PackageInfo extends PackageVersionsInfoWithMetadata {
   versionToTag: Map<string, string>
 }
 
-/**
- * Fetch npm package versions and build a version-to-tag lookup map.
- *
- * @see https://github.com/antfu/fast-npm-meta
- */
-export const getPackageInfo = memoize<string, Promise<PackageInfo | null>>(async (name) => {
-  logger.info(`Fetching package info for ${name}`)
-
-  const pkg = await getVersions(name, {
-    metadata: true,
-    throw: false,
-  })
-
+function parsePackageInfo(name: string, pkg: MaybeError<PackageVersionsInfoWithMetadata>) {
   if ('error' in pkg) {
-    logger.warn(`Fetching package info for ${name} error: ${JSON.stringify(pkg)}`)
+    logger.warn(`[package] Fetching error(${name}): ${JSON.stringify(pkg)}`)
 
     // Return null to trigger a cache hit
     if (pkg.status === 404)
@@ -30,8 +21,6 @@ export const getPackageInfo = memoize<string, Promise<PackageInfo | null>>(async
     throw pkg
   }
 
-  logger.info(`Fetched package info for ${name}`)
-
   const versionToTag = new Map<string, string>()
   if (pkg.distTags) {
     for (const [tag, ver] of Object.entries(pkg.distTags)) {
@@ -40,4 +29,45 @@ export const getPackageInfo = memoize<string, Promise<PackageInfo | null>>(async
   }
 
   return { ...pkg, versionToTag }
+}
+
+const getPackageInfoBatch = createBatchRunner<string, PackageInfo | null>({
+  maxSize: BATCH_SIZE,
+  runBatch: async (names) => {
+    const logName = names.join(', ')
+    logger.info(`[package] Fetching ${logName}`)
+
+    const list = await getVersionsBatch(names, {
+      metadata: true,
+      throw: false,
+    })
+
+    logger.info(`[package] Fetched ${logName}`)
+
+    const values = new Map<string, PackageInfo | null>()
+    const errors = new Map<string, unknown>()
+
+    names.forEach((name, index) => {
+      const item = list[index]
+      if (!item) {
+        errors.set(name, new Error(`Missing package info response for ${name}`))
+        return
+      }
+
+      try {
+        values.set(name, parsePackageInfo(name, item))
+      } catch (error) {
+        errors.set(name, error)
+      }
+    })
+
+    return { values, errors }
+  },
 })
+
+/**
+ * Fetch npm package versions and build a version-to-tag lookup map.
+ *
+ * @see https://github.com/antfu/fast-npm-meta
+ */
+export const getPackageInfo = memoize<string, Promise<PackageInfo | null>>(async (name) => getPackageInfoBatch(name))
diff --git a/src/utils/batch.ts b/src/utils/batch.ts
@@ -0,0 +1,97 @@
+export interface BatchRunResult<TResult> {
+  values: Map<string, TResult>
+  errors: Map<string, unknown>
+}
+
+export interface BatchRunnerOptions<TTask, TResult> {
+  getKey?: (task: TTask) => string
+  maxSize?: number
+  runBatch: (tasks: TTask[]) => Promise<BatchRunResult<TResult>>
+}
+
+interface PendingRequest<TResult> {
+  resolve: (value: TResult) => void
+  reject: (reason?: unknown) => void
+}
+
+interface PendingTaskGroup<TTask, TResult> {
+  task: TTask
+  waiters: PendingRequest<TResult>[]
+}
+
+export function createBatchRunner<TTask, TResult>(options: BatchRunnerOptions<TTask, TResult>): (task: TTask) => Promise<TResult> {
+  const {
+    getKey = String,
+    maxSize,
+    runBatch,
+  } = options
+  const pendingTasksByKey = new Map<string, PendingTaskGroup<TTask, TResult>>()
+
+  let isFlushScheduled = false
+
+  const resolveTaskGroup = (group: PendingTaskGroup<TTask, TResult>, value: TResult) => {
+    group.waiters.forEach(({ resolve }) => resolve(value))
+  }
+
+  const rejectTaskGroup = (group: PendingTaskGroup<TTask, TResult>, error: unknown) => {
+    group.waiters.forEach(({ reject }) => reject(error))
+  }
+
+  const flushPendingTasks = async () => {
+    isFlushScheduled = false
+    if (pendingTasksByKey.size === 0)
+      return
+
+    const pendingTaskGroups = [...pendingTasksByKey.values()]
+    pendingTasksByKey.clear()
+
+    let batchResult: BatchRunResult<TResult>
+
+    try {
+      batchResult = await runBatch(pendingTaskGroups.map((group) => group.task))
+    } catch (batchError) {
+      pendingTaskGroups.forEach((group) => {
+        rejectTaskGroup(group, batchError)
+      })
+      return
+    }
+
+    pendingTaskGroups.forEach((group) => {
+      const taskKey = getKey(group.task)
+      if (batchResult.values.has(taskKey)) {
+        resolveTaskGroup(group, batchResult.values.get(taskKey)!)
+        return
+      }
+
+      if (batchResult.errors.has(taskKey)) {
+        rejectTaskGroup(group, batchResult.errors.get(taskKey))
+        return
+      }
+
+      rejectTaskGroup(group, new Error(`Missing batch outcome for key "${taskKey}"`))
+    })
+  }
+
+  return (task: TTask) => new Promise<TResult>((resolve, reject) => {
+    const taskKey = getKey(task)
+    const existingGroup = pendingTasksByKey.get(taskKey)
+
+    if (existingGroup) {
+      existingGroup.waiters.push({ resolve, reject })
+    } else {
+      pendingTasksByKey.set(taskKey, {
+        task,
+        waiters: [{ resolve, reject }],
+      })
+    }
+
+    if (maxSize && pendingTasksByKey.size >= maxSize) {
+      void flushPendingTasks()
+    } else if (!isFlushScheduled) {
+      isFlushScheduled = true
+      queueMicrotask(() => {
+        void flushPendingTasks()
+      })
+    }
+  })
+}
diff --git a/tests/utils/batch.test.ts b/tests/utils/batch.test.ts
@@ -0,0 +1,128 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { createBatchRunner } from '../../src/utils/batch'
+
+describe('createBatchRunner', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('should batch different tasks in the same tick', async () => {
+    const runBatch = vi.fn(async (tasks: string[]) => {
+      const values = new Map<string, string>()
+      tasks.forEach((task) => {
+        values.set(task, task.toUpperCase())
+      })
+      return { values, errors: new Map<string, unknown>() }
+    })
+
+    const run = createBatchRunner<string, string>({
+      runBatch,
+    })
+
+    const [a, b] = await Promise.all([
+      run('a'),
+      run('b'),
+    ])
+
+    expect(a).toBe('A')
+    expect(b).toBe('B')
+    expect(runBatch).toHaveBeenCalledTimes(1)
+    expect(runBatch).toHaveBeenCalledWith(['a', 'b'])
+  })
+
+  it('should deduplicate same-key requests in one batch', async () => {
+    const runBatch = vi.fn(async (tasks: string[]) => {
+      const values = new Map<string, string>()
+      tasks.forEach((task) => {
+        values.set(task, `${task}-resolved`)
+      })
+      return { values, errors: new Map<string, unknown>() }
+    })
+
+    const run = createBatchRunner<string, string>({
+      runBatch,
+    })
+
+    const [first, second] = await Promise.all([
+      run('pkg'),
+      run('pkg'),
+    ])
+
+    expect(first).toBe('pkg-resolved')
+    expect(second).toBe('pkg-resolved')
+    expect(runBatch).toHaveBeenCalledTimes(1)
+    expect(runBatch).toHaveBeenCalledWith(['pkg'])
+  })
+
+  it('should reject all tasks when runBatch throws', async () => {
+    const runBatch = vi.fn(async () => {
+      throw new Error('batch failed')
+    })
+
+    const run = createBatchRunner<string, string>({
+      runBatch,
+    })
+
+    const [a, b] = await Promise.allSettled([
+      run('a'),
+      run('b'),
+    ])
+
+    expect(a.status).toBe('rejected')
+    expect(b.status).toBe('rejected')
+    expect(runBatch).toHaveBeenCalledTimes(1)
+  })
+
+  it('should flush immediately when maxSize is reached and split into multiple batches', async () => {
+    const runBatch = vi.fn(async (tasks: string[]) => {
+      const values = new Map<string, string>()
+      tasks.forEach((task) => {
+        values.set(task, task.toUpperCase())
+      })
+      return { values, errors: new Map<string, unknown>() }
+    })
+
+    const run = createBatchRunner<string, string>({
+      maxSize: 2,
+      runBatch,
+    })
+
+    const a = run('a')
+    expect(runBatch).not.toHaveBeenCalled()
+
+    const [, b, c] = await Promise.all([
+      a,
+      run('b'),
+      run('c'),
+    ])
+
+    expect(await a).toBe('A')
+    expect(b).toBe('B')
+    expect(c).toBe('C')
+    expect(runBatch).toHaveBeenCalledTimes(2)
+    expect(runBatch).toHaveBeenNthCalledWith(1, ['a', 'b'])
+    expect(runBatch).toHaveBeenNthCalledWith(2, ['c'])
+  })
+
+  it('should reject errored or missing outcomes', async () => {
+    const runBatch = vi.fn(async () => ({
+      values: new Map<string, string>([['a', 'A']]),
+      errors: new Map<string, unknown>([['b', new Error('b failed')]]),
+    }))
+
+    const run = createBatchRunner<string, string>({
+      runBatch,
+    })
+
+    const [a, b, c] = await Promise.allSettled([
+      run('a'),
+      run('b'),
+      run('c'),
+    ])
+
+    expect(a).toEqual({ status: 'fulfilled', value: 'A' })
+    expect(b.status).toBe('rejected')
+    expect(c.status).toBe('rejected')
+    expect(runBatch).toHaveBeenCalledTimes(1)
+  })
+})
