fix(nbd-client): fix long datamap read nbddisk (#10157)

Customer imported large VM from VMWare (> 3 To), import switched to Qcow2 automatically but the process was stuck while building the qcow2 L1/L2 address tables, where the NbdDisk's hasBlock is called for each l2 table entry which iterated over the disk dataMap which contained ~113,528 ranges. This froze the process for more than 18 minutes after which it was either killed or died due to OOM error.

This fix improves the NbdDisk by:

    Updating processDatamap to throw on overlapping ranges and merge touching extents.
    Override default getBlockIndexesCount with optimized Nbd one.
    Improve hasBlock with a built in cursor to improve performance.

Test file added to cover the new function and updated code.

@vates/nbd-client/NbdDisk.mjs

@@ -19,6 +19,12 @@ export class NbdDisk extends RandomAccessDisk {
   /** @type {number} */
   #blockSize
 
+  /** @type {number | undefined} */
+  #hasBlockCursor
+
+  /** @type {number | undefined} */
+  #hasBlockPreviousIndex
+
   constructor(nbdInfos, blockSize, { dataMap } = {}) {
     super()
     this.#blockSize = blockSize
@@ -27,11 +33,38 @@ export class NbdDisk extends RandomAccessDisk {
   }
 
   #processDatamap(rawDataMap) {
-    return rawDataMap
-      .filter(({ type }) => type === 0)
+    const ranges = rawDataMap
+      .filter(({ type, length }) => type === 0 && length > 0)
       .map(({ offset, length }) => ({ offset, length }))
       .sort(({ offset: offset1 }, { offset: offset2 }) => offset1 - offset2)
+
+    // hasBlock()'s forward-only cursor is only correct if the extents are
+    // sorted AND disjoint: an earlier extent must never reach into a block
+    // after a later one, otherwise the cursor could skip it and wrongly report
+    // a block as empty, dropping data and corrupting the output.
+    // Touching extents (cur.offset === prevEnd) are merged
+    const merged = []
+    for (const range of ranges) {
+      const last = merged[merged.length - 1]
+      if (last !== undefined) {
+        const lastEnd = last.offset + last.length
+        if (range.offset < lastEnd) {
+          throw new Error(
+            `overlapping ranges in data map: [${last.offset}, ${lastEnd}) and [${range.offset}, ${range.offset + range.length})`
+          )
+        }
+        if (range.offset === lastEnd) {
+          // touching: extend the previous extent
+          last.length += range.length
+          continue
+        }
+      }
+      // gap (or first range): new extent
+      merged.push(range)
+    }
+    return merged
   }
+
   /**
    * @param {number} index
    * @returns {Promise<DiskBlock>}
@@ -117,6 +150,34 @@ export class NbdDisk extends RandomAccessDisk {
     return [...indexes].sort((a, b) => a - b)
   }
 
+  /**
+   * Counts the allocated blocks without materializing the index list.
+   *
+   * @returns {number}
+   */
+  getBlockIndexesCount() {
+    if (!this.#dataMap) {
+      throw new Error("can't getBlockIndexesCount before init")
+    }
+
+    const blockSize = this.getBlockSize()
+
+    let count = 0
+    let lastCountedBlock = -1
+    for (const { offset, length } of this.#dataMap) {
+      const firstBlockIndex = Math.floor(offset / blockSize)
+      const lastBlockIndex = Math.floor((offset + length - 1) / blockSize)
+      const from = Math.max(firstBlockIndex, lastCountedBlock + 1)
+
+      if (lastBlockIndex >= from) {
+        count += lastBlockIndex - from + 1
+        lastCountedBlock = lastBlockIndex
+      }
+    }
+
+    return count
+  }
+
   /**
    * @param {number} index
    * @returns {boolean}
@@ -125,8 +186,32 @@ export class NbdDisk extends RandomAccessDisk {
     if (!this.#dataMap) {
       throw new Error("can't hasBlock before init")
     }
-    const blockStart = index * this.getBlockSize()
-    const blockEnd = (index + 1) * this.getBlockSize()
-    return this.#dataMap.some(({ offset, length }) => offset + length > blockStart && offset < blockEnd)
+
+    const blockSize = this.getBlockSize()
+    const blockStart = index * blockSize
+    const blockEnd = blockStart + blockSize
+
+    let startExtentIndex = 0
+    if (this.#hasBlockCursor !== undefined && index >= this.#hasBlockPreviousIndex) {
+      startExtentIndex = this.#hasBlockCursor
+    } else {
+      this.#hasBlockCursor = undefined
+    }
+    this.#hasBlockPreviousIndex = index
+
+    const dataMap = this.#dataMap
+    const l = dataMap.length
+    for (let i = startExtentIndex; i < l; i++) {
+      const { offset, length } = dataMap[i]
+      if (offset >= blockEnd) {
+        // extents are sorted: nothing from here on can overlap this block
+        break
+      }
+      if (offset + length > blockStart) {
+        this.#hasBlockCursor = i
+        return true
+      }
+    }
+    return false
   }
 }

@vates/nbd-client/NbdDisk.test.mjs

@@ -0,0 +1,127 @@
+import { describe, it } from 'node:test'
+import assert from 'node:assert'
+
+import { NbdDisk } from './NbdDisk.mjs'
+
+const KiB = 1024
+const MiB = 1024 * KiB
+const TiB = 1024 * 1024 * MiB
+
+const SOURCE_BLOCK_SIZE = 2 * MiB // VHD_BLOCK_SIZE, the block size used when importing from VMware
+const CLUSTER_SIZE = 64 * KiB // qcow2 cluster size
+
+// hasBlock() only needs the dataMap (passed through the constructor) and the
+// block size, so we can exercise it without a live NBD connection.
+describe('NbdDisk.hasBlock', () => {
+  it('reports allocated and unallocated blocks correctly', () => {
+    const dataMap = [
+      { type: 0, offset: 0, length: SOURCE_BLOCK_SIZE }, // block 0
+      { type: 0, offset: 5 * SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE }, // block 5
+    ]
+    const disk = new NbdDisk({}, SOURCE_BLOCK_SIZE, { dataMap })
+
+    assert.strictEqual(disk.hasBlock(0), true)
+    assert.strictEqual(disk.hasBlock(1), false)
+    assert.strictEqual(disk.hasBlock(4), false)
+    assert.strictEqual(disk.hasBlock(5), true)
+    assert.strictEqual(disk.hasBlock(6), false)
+  })
+
+  it('handles repeated and out-of-order queries with the resume cursor', () => {
+    const dataMap = [
+      { type: 0, offset: 0, length: SOURCE_BLOCK_SIZE }, // block 0
+      { type: 0, offset: 5 * SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE }, // block 5
+    ]
+    const disk = new NbdDisk({}, SOURCE_BLOCK_SIZE, { dataMap })
+
+    // repeated indexes (the qcow2 path queries each 2MB block ~32 times as it
+    // walks 64KB clusters) must return a stable answer
+    assert.strictEqual(disk.hasBlock(5), true)
+    assert.strictEqual(disk.hasBlock(5), true)
+    assert.strictEqual(disk.hasBlock(6), false)
+    assert.strictEqual(disk.hasBlock(6), false)
+
+    // querying backwards must still be correct (falls back to a full scan)
+    assert.strictEqual(disk.hasBlock(0), true)
+    assert.strictEqual(disk.hasBlock(5), true)
+  })
+
+  it('invalidates the cursor on a backward miss (no stale skip)', () => {
+    const dataMap = [
+      { type: 0, offset: 3 * SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE }, // block 3
+      { type: 0, offset: 10 * SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE }, // block 10
+    ]
+    const disk = new NbdDisk({}, SOURCE_BLOCK_SIZE, { dataMap })
+
+    assert.strictEqual(disk.hasBlock(10), true) // forward hit: cursor parks on the high extent
+    assert.strictEqual(disk.hasBlock(1), false) // backward miss: must invalidate the cursor
+    assert.strictEqual(disk.hasBlock(3), true) // lower extent must not be skipped
+  })
+
+  it('throws on overlapping ranges (would break the resume cursor)', () => {
+    const dataMap = [
+      { type: 0, offset: 0, length: SOURCE_BLOCK_SIZE },
+      { type: 0, offset: SOURCE_BLOCK_SIZE / 2, length: SOURCE_BLOCK_SIZE }, // overlaps the previous range
+    ]
+    assert.throws(() => new NbdDisk({}, SOURCE_BLOCK_SIZE, { dataMap }), /overlapping ranges/)
+  })
+
+  it('merges touching ranges and drops zero-length / non type-0 entries', () => {
+    const dataMap = [
+      { type: 0, offset: 2 * SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE }, // [2, 3)
+      { type: 0, offset: 0, length: SOURCE_BLOCK_SIZE / 2 }, // [0, .5)
+      { type: 0, offset: SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE }, // [1, 2)
+      { type: 0, offset: SOURCE_BLOCK_SIZE / 2, length: SOURCE_BLOCK_SIZE / 2 }, // [.5, 1), touches [0, .5)
+      { type: 1, offset: 10 * SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE }, // wrong type: ignored
+      { type: 0, offset: 7 * SOURCE_BLOCK_SIZE, length: 0 }, // zero length: ignored
+    ]
+    const disk = new NbdDisk({}, SOURCE_BLOCK_SIZE, { dataMap })
+
+    assert.strictEqual(disk.hasBlock(0), true)
+    assert.strictEqual(disk.hasBlock(1), true)
+    assert.strictEqual(disk.hasBlock(2), true)
+    assert.strictEqual(disk.hasBlock(3), false)
+    assert.strictEqual(disk.hasBlock(7), false) // zero-length range dropped
+    assert.strictEqual(disk.hasBlock(10), false) // non type-0 range dropped
+  })
+})
+
+describe('NbdDisk.getBlockIndexesCount', () => {
+  const cases = {
+    'aligned, disjoint blocks': [
+      { type: 0, offset: 0, length: SOURCE_BLOCK_SIZE }, // block 0
+      { type: 0, offset: 5 * SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE }, // block 5
+    ],
+    'a range spanning several blocks': [
+      { type: 0, offset: SOURCE_BLOCK_SIZE / 2, length: 3 * SOURCE_BLOCK_SIZE }, // blocks 0,1,2,3
+    ],
+    'two disjoint extents sharing block 0 (needs dedup)': [
+      { type: 0, offset: 0, length: SOURCE_BLOCK_SIZE / 2 }, // [0, .5) => block 0
+      { type: 0, offset: 3 * (SOURCE_BLOCK_SIZE / 4), length: SOURCE_BLOCK_SIZE }, // [.75, 1.75) => blocks 0,1
+    ],
+    'touching sub-block fragments': [
+      { type: 0, offset: 0, length: SOURCE_BLOCK_SIZE / 2 },
+      { type: 0, offset: SOURCE_BLOCK_SIZE / 2, length: SOURCE_BLOCK_SIZE / 2 },
+      { type: 0, offset: SOURCE_BLOCK_SIZE, length: SOURCE_BLOCK_SIZE },
+    ],
+    empty: [],
+  }
+
+  for (const [name, dataMap] of Object.entries(cases)) {
+    it(`matches getBlockIndexes().length: ${name}`, () => {
+      const disk = new NbdDisk({}, SOURCE_BLOCK_SIZE, { dataMap })
+      assert.strictEqual(disk.getBlockIndexesCount(), disk.getBlockIndexes().length)
+    })
+  }
+
+  it('matches getBlockIndexes().length on the large sparse map', () => {
+    const virtualSize = Math.floor(3.5 * TiB)
+    const NB_RANGES = 113528
+    const dataMap = []
+    for (let i = 0; i < NB_RANGES; i++) {
+      dataMap.push({ type: 0, offset: Math.floor((i / NB_RANGES) * virtualSize), length: CLUSTER_SIZE })
+    }
+    const disk = new NbdDisk({}, SOURCE_BLOCK_SIZE, { dataMap })
+    assert.strictEqual(disk.getBlockIndexesCount(), disk.getBlockIndexes().length)
+  })
+})