perf(vite): speed up inspect payload allocation

packages/vite/src/node/__tests__/inspect-context.test.ts

@@ -8,6 +8,7 @@ import { join } from 'node:path'
 import { afterEach, describe, expect, it, vi } from 'vitest'
 import { ViteInspectContext } from '../inspect/context'
 import { hijackPlugin } from '../inspect/hijack'
+import { createInspectPayloadStore } from '../inspect/store/payload'
 
 const contexts: ViteInspectContext[] = []
 
@@ -1410,3 +1411,79 @@ describe('vite inspect invalidation', () => {
     expect(await ctx.store.getModuleIds(scope)).toEqual(['/abs/e.js'])
   })
 })
+
+describe('vite inspect payload archive free space', () => {
+  async function createPayloads() {
+    const dir = mkdtempSync(join(tmpdir(), 'vite-inspect-free-'))
+    const payloads = await createInspectPayloadStore(join(dir, 'payloads.bin'))
+    return {
+      payloads,
+      cleanup: async () => {
+        await payloads.close()
+        rmSync(dir, { recursive: true, force: true })
+      },
+    }
+  }
+
+  it('merges a freed range with the free ranges before and after it', async () => {
+    const { payloads, cleanup } = await createPayloads()
+    try {
+      const [a, b, c, d] = await payloads.write(['a'.repeat(10), 'b'.repeat(10), 'c'.repeat(10), 'd'.repeat(10)])
+      // b joins the free range before it.
+      payloads.reclaim([a!])
+      payloads.reclaim([b!])
+      const [ab] = await payloads.write(['m'.repeat(20)])
+      expect(ab).toEqual({ offset: 0, length: 20 })
+
+      // ab joins the free range after it.
+      payloads.reclaim([c!])
+      payloads.reclaim([ab!])
+      const [abc] = await payloads.write(['n'.repeat(30)])
+      expect(abc).toEqual({ offset: 0, length: 30 })
+
+      expect(await payloads.read(abc!)).toBe('n'.repeat(30))
+      expect(await payloads.read(d!)).toBe('d'.repeat(10))
+    }
+    finally {
+      await cleanup()
+    }
+  })
+
+  it('picks the smallest free range that fits, then the lowest offset', async () => {
+    const { payloads, cleanup } = await createPayloads()
+    try {
+      const [large, , smallFirst, , smallSecond, last] = await payloads.write([
+        'l'.repeat(50),
+        'x',
+        's'.repeat(20),
+        'y',
+        't'.repeat(20),
+        'z',
+      ])
+      payloads.reclaim([smallSecond!, large!, smallFirst!])
+
+      const [fit] = await payloads.write(['f'.repeat(15)])
+      expect(fit).toEqual({ offset: smallFirst!.offset, length: 15 })
+      // The rest of that range is still the best fit for a payload that matches its size.
+      const [rest] = await payloads.write(['r'.repeat(5)])
+      expect(rest).toEqual({ offset: smallFirst!.offset + 15, length: 5 })
+      const [big] = await payloads.write(['b'.repeat(40)])
+      expect(big).toEqual({ offset: large!.offset, length: 40 })
+
+      // Sizes in the same size class: the smaller range wins even though it sits at a higher offset.
+      const [wide, , narrow] = await payloads.write(['w'.repeat(30), 'g', 'n'.repeat(18)])
+      payloads.reclaim([wide!, narrow!])
+      const [snug] = await payloads.write(['q'.repeat(17)])
+      expect(snug).toEqual({ offset: narrow!.offset, length: 17 })
+      expect(await payloads.read(snug!)).toBe('q'.repeat(17))
+
+      expect(await payloads.read(fit!)).toBe('f'.repeat(15))
+      expect(await payloads.read(rest!)).toBe('r'.repeat(5))
+      expect(await payloads.read(big!)).toBe('b'.repeat(40))
+      expect(await payloads.read(last!)).toBe('z')
+    }
+    finally {
+      await cleanup()
+    }
+  })
+})

packages/vite/src/node/inspect/store/payload.ts

@@ -4,6 +4,9 @@ import { Buffer } from 'node:buffer'
 import { open } from 'node:fs/promises'
 import { diagnostics } from '../../diagnostics'
 
+// Lengths are below 2^53, so their log2 bucket never exceeds this.
+const MAX_SIZE_BUCKET = 53
+
 export interface InspectPayloadStore {
   write: (values: Array<string | null | undefined>) => Promise<Array<ViteInspectPayloadRange | undefined>>
   read: (range: ViteInspectPayloadRange) => Promise<string>
@@ -55,7 +58,9 @@ class MemoryInspectPayloadStore implements InspectPayloadStore {
 
 class FileInspectPayloadStore implements InspectPayloadStore {
   private nextOffset = 0
+  // Free ranges sorted by offset, plus the same ranges bucketed by size so allocation doesn't scan them all.
   private freeRanges: ViteInspectPayloadRange[] = []
+  private readonly freeRangesBySize = new Map<number, Set<ViteInspectPayloadRange>>()
   // Payloads keyed by offset once the archive file can no longer be written.
   private fallback?: Map<number, Buffer>
 
@@ -125,39 +130,48 @@ class FileInspectPayloadStore implements InspectPayloadStore {
 
   async close(): Promise<void> {
     this.freeRanges.length = 0
+    this.freeRangesBySize.clear()
     this.fallback?.clear()
     await this.file.close()
   }
 
   reclaim(ranges: ViteInspectPayloadRange[]): void {
     for (const range of ranges)
       this.fallback?.delete(range.offset)
-    this.freeRanges.push(...ranges.filter(range => range.length > 0))
-    this.mergeFreeRanges()
+    for (const range of ranges) {
+      if (range.length > 0)
+        this.insertFreeRange(range.offset, range.length)
+    }
   }
 
   private allocate(length: number): ViteInspectPayloadRange {
     if (length === 0)
       return { offset: this.nextOffset, length }
 
-    let bestIndex = -1
-    for (let index = 0; index < this.freeRanges.length; index++) {
-      const range = this.freeRanges[index]!
-      if (range.length < length)
-        continue
-      if (bestIndex < 0 || range.length < this.freeRanges[bestIndex]!.length)
-        bestIndex = index
+    // Best fit: the smallest free range that is large enough, the lowest offset among equal sizes.
+    let best: ViteInspectPayloadRange | undefined
+    for (let bucket = sizeBucket(length); bucket <= MAX_SIZE_BUCKET; bucket++) {
+      for (const range of this.freeRangesBySize.get(bucket) ?? []) {
+        if (range.length < length)
+          continue
+        if (!best || range.length < best.length || (range.length === best.length && range.offset < best.offset))
+          best = range
+      }
+      // Every range in a higher bucket is larger than any range in this one.
+      if (best)
+        break
     }
 
-    if (bestIndex >= 0) {
-      const range = this.freeRanges[bestIndex]!
-      const allocated = { offset: range.offset, length }
-      if (range.length === length) {
-        this.freeRanges.splice(bestIndex, 1)
+    if (best) {
+      const allocated = { offset: best.offset, length }
+      this.removeFromSizeBucket(best)
+      if (best.length === length) {
+        this.freeRanges.splice(this.findFreeRangeIndex(best.offset), 1)
       }
       else {
-        range.offset += length
-        range.length -= length
+        best.offset += length
+        best.length -= length
+        this.addToSizeBucket(best)
       }
       return allocated
     }
@@ -167,6 +181,63 @@ class FileInspectPayloadStore implements InspectPayloadStore {
     return allocated
   }
 
+  private insertFreeRange(offset: number, length: number): void {
+    let end = offset + length
+    let index = this.findFreeRangeIndex(offset)
+    let start = offset
+    // Merge with the previous range when it touches or overlaps this one.
+    const previous = this.freeRanges[index - 1]
+    if (previous && previous.offset + previous.length >= start) {
+      index -= 1
+      start = previous.offset
+      end = Math.max(end, previous.offset + previous.length)
+    }
+    // Merge every following range that starts before this one ends.
+    let removeCount = 0
+    while (index + removeCount < this.freeRanges.length && this.freeRanges[index + removeCount]!.offset <= end) {
+      const next = this.freeRanges[index + removeCount]!
+      end = Math.max(end, next.offset + next.length)
+      removeCount += 1
+    }
+    for (let i = index; i < index + removeCount; i++)
+      this.removeFromSizeBucket(this.freeRanges[i]!)
+    const merged = { offset: start, length: end - start }
+    this.freeRanges.splice(index, removeCount, merged)
+    this.addToSizeBucket(merged)
+  }
+
+  // Index of the first free range whose offset is not below `offset`.
+  private findFreeRangeIndex(offset: number): number {
+    let low = 0
+    let high = this.freeRanges.length
+    while (low < high) {
+      const middle = (low + high) >>> 1
+      if (this.freeRanges[middle]!.offset < offset)
+        low = middle + 1
+      else
+        high = middle
+    }
+    return low
+  }
+
+  private addToSizeBucket(range: ViteInspectPayloadRange): void {
+    const bucket = sizeBucket(range.length)
+    let ranges = this.freeRangesBySize.get(bucket)
+    if (!ranges) {
+      ranges = new Set()
+      this.freeRangesBySize.set(bucket, ranges)
+    }
+    ranges.add(range)
+  }
+
+  private removeFromSizeBucket(range: ViteInspectPayloadRange): void {
+    const bucket = sizeBucket(range.length)
+    const ranges = this.freeRangesBySize.get(bucket)
+    ranges?.delete(range)
+    if (ranges?.size === 0)
+      this.freeRangesBySize.delete(bucket)
+  }
+
   private async writeAllocatedPayloads(
     writes: Array<{ buffer: Buffer, range: ViteInspectPayloadRange }>,
   ): Promise<void> {
@@ -203,23 +274,10 @@ class FileInspectPayloadStore implements InspectPayloadStore {
       offset += bytesWritten
     }
   }
+}
 
-  private mergeFreeRanges(): void {
-    this.freeRanges.sort((a, b) => a.offset - b.offset)
-    const merged: ViteInspectPayloadRange[] = []
-    for (const range of this.freeRanges) {
-      const previous = merged.at(-1)
-      if (!previous || range.offset > previous.offset + previous.length) {
-        merged.push({ ...range })
-        continue
-      }
-      previous.length = Math.max(
-        previous.length,
-        range.offset + range.length - previous.offset,
-      )
-    }
-    this.freeRanges = merged
-  }
+function sizeBucket(length: number): number {
+  return Math.floor(Math.log2(length))
 }
 
 function consumeBuffers(input: Buffer[], bytes: number): Buffer[] {