fix[notask]: drain in-flight blob blocks before cleanup

packages/registry-server/client/lib/client.js

@@ -15,6 +15,8 @@ const fs = require('#fs')
 
 const DEFAULT_DOWNLOAD_MAX_RETRIES = 3
 const RETRIABLE_DOWNLOAD_CODES = ['REQUEST_TIMEOUT']
+const INFLIGHT_DRAIN_TIMEOUT_MS = 5000
+const INFLIGHT_DRAIN_POLL_MS = 10
 
 // While the app is backgrounded the swarm is suspended; a retry must wait for
 // resume rather than burn its (small) retry budget timing out against a dead
@@ -515,10 +517,11 @@ class QVACRegistryClient extends ReadyResource {
   }
 
   async _releaseDownload(core, blobs, rangeDownload, blockStart, blockEnd) {
-    // Stop replication before clearing to prevent blocks from being refetched.
+    // Hypercore can commit in-flight responses after a range is destroyed.
     if (rangeDownload) rangeDownload.destroy()
 
     if (core && blockStart !== undefined) {
+      if (rangeDownload) await this._waitForInflightBlocks(core)
       await this._clearBlobBlocks(core, blockStart, blockEnd)
     }
     if (blobs) {
@@ -539,6 +542,28 @@ class QVACRegistryClient extends ReadyResource {
     this.logger.debug('Blob resources released')
   }
 
+  async _waitForInflightBlocks(core) {
+    if (!Array.isArray(core.peers)) return
+
+    const timeoutMs = this._inflightDrainTimeoutMs ?? INFLIGHT_DRAIN_TIMEOUT_MS
+    const pollMs = this._inflightDrainPollMs ?? INFLIGHT_DRAIN_POLL_MS
+    const deadline = Date.now() + timeoutMs
+    while (core.peers.some((peer) => peer.inflight > 0 || peer.dataProcessing > 0)) {
+      if (Date.now() >= deadline) {
+        this.logger.warn('Timed out waiting for in-flight blob blocks before clearing', {
+          timeoutMs,
+          peers: core.peers.map((peer, index) => ({
+            peer: index,
+            inflight: peer.inflight,
+            dataProcessing: peer.dataProcessing
+          }))
+        })
+        return
+      }
+      await new Promise((resolve) => setTimeout(resolve, pollMs))
+    }
+  }
+
   async _clearBlobBlocks(core, start, end) {
     try {
       const cleared = await core.clear(start, end, { diff: true })

packages/registry-server/client/tests/unit/client.download-retry.test.js

@@ -367,6 +367,90 @@ test('downloadModel clears cached blocks when the download fails', async (t) =>
   )
 })
 
+test('downloadModel waits for in-flight blocks before clearing after failure', async (t) => {
+  const dir = await tmp(t)
+  const outputFile = path.join(dir, 'model.gguf')
+
+  const client = makeClient(t)
+  const peer = { inflight: 1, dataProcessing: 0 }
+  client._core.peers = [peer]
+  client._core.download = () => ({
+    destroy() {
+      client._events.push('destroy')
+      setTimeout(() => {
+        peer.inflight = 0
+        client._events.push('drained')
+      }, 20)
+    }
+  })
+  // lunte-disable-next-line require-await
+  client._clearBlobBlocks = async () => {
+    client._events.push('clear')
+  }
+  // lunte-disable-next-line require-await
+  client._streamBlobToFile = async () => {
+    throw new Error('Download cancelled')
+  }
+
+  await t.exception(
+    () => client.downloadModel('models/tiny.gguf', 's3', { outputFile }),
+    /Download cancelled/,
+    'the failed download rejects'
+  )
+
+  const drainedAt = client._events.indexOf('drained')
+  t.ok(
+    drainedAt !== -1 && drainedAt < client._events.indexOf('clear'),
+    'in-flight blocks drain before the cached range is cleared'
+  )
+})
+
+test('downloadModel logs remaining peer counters when drain times out', async (t) => {
+  const dir = await tmp(t)
+  const outputFile = path.join(dir, 'model.gguf')
+
+  const client = makeClient(t)
+  const warnings = []
+  const baseWarn = client.logger.warn.bind(client.logger)
+  client.logger.warn = (msg, data) => {
+    warnings.push({ msg, data })
+    baseWarn(msg, data)
+  }
+  client._inflightDrainTimeoutMs = 30
+  client._inflightDrainPollMs = 5
+  client._core.peers = [{ inflight: 2, dataProcessing: 1 }]
+  client._core.download = () => ({
+    destroy() {
+      client._events.push('destroy')
+    }
+  })
+  // lunte-disable-next-line require-await
+  client._clearBlobBlocks = async () => {
+    client._events.push('clear')
+  }
+  // lunte-disable-next-line require-await
+  client._streamBlobToFile = async () => {
+    throw new Error('Download cancelled')
+  }
+
+  await t.exception(
+    () => client.downloadModel('models/tiny.gguf', 's3', { outputFile }),
+    /Download cancelled/,
+    'the failed download rejects'
+  )
+
+  const timeoutWarning = warnings.find((entry) =>
+    entry.msg.includes('Timed out waiting for in-flight blob blocks')
+  )
+  t.ok(timeoutWarning, 'drain timeout is logged')
+  t.alike(
+    timeoutWarning.data.peers,
+    [{ peer: 0, inflight: 2, dataProcessing: 1 }],
+    'timeout warning includes remaining inflight and dataProcessing counters'
+  )
+  t.ok(client._events.includes('clear'), 'best-effort clear still runs after the timeout')
+})
+
 test('downloadModel clears cached blocks when the returned stream is destroyed', async (t) => {
   const client = makeClient(t)
   const clears = []

packages/registry-server/client/tests/unit/client.download-storage.test.js

@@ -0,0 +1,226 @@
+'use strict'
+
+const test = require('brittle')
+const tmp = require('test-tmp')
+const path = require('#path')
+const Corestore = require('corestore')
+const Hyperblobs = require('hyperblobs')
+const QVACRegistryClient = require('../../lib/client')
+
+const PAYLOAD_BYTES = 8 * 1024 * 1024
+const CANCEL_AFTER_BLOCKS = 8
+
+function logger() {
+  return {
+    info() {},
+    debug() {},
+    warn() {},
+    error() {}
+  }
+}
+
+function cancellation() {
+  let onAbort = null
+  const signal = {
+    aborted: false,
+    addEventListener(event, listener) {
+      if (event === 'abort') onAbort = listener
+    }
+  }
+
+  return {
+    signal,
+    abort() {
+      signal.aborted = true
+      if (onAbort) onAbort()
+    }
+  }
+}
+
+async function countBlocks(storage, coreKey, pointer) {
+  const store = new Corestore(storage)
+  const core = store.get({ key: coreKey })
+  await core.ready()
+
+  let blocks = 0
+  for (
+    let index = pointer.blockOffset;
+    index < pointer.blockOffset + pointer.blockLength;
+    index++
+  ) {
+    if (await core.has(index)) blocks++
+  }
+
+  await core.close()
+  await store.close()
+  return blocks
+}
+
+test('cancelled download leaves no durable blocks after reopen', async (t) => {
+  const root = await tmp(t)
+  const writerStorage = path.join(root, 'writer')
+  const readerStorage = path.join(root, 'reader')
+  const outputFile = path.join(root, 'partial.gguf')
+
+  const writerStore = new Corestore(writerStorage)
+  const writerCore = writerStore.get({ name: 'blobs' })
+  await writerCore.ready()
+  const coreKey = writerCore.key
+  const writerBlobs = new Hyperblobs(writerCore)
+  await writerBlobs.ready()
+  const pointer = await writerBlobs.put(Buffer.alloc(PAYLOAD_BYTES, 1))
+
+  const readerStore = new Corestore(readerStorage)
+  const readerCore = readerStore.get({ key: coreKey })
+  await readerCore.ready()
+  const readerBlobs = new Hyperblobs(readerCore)
+  await readerBlobs.ready()
+
+  const writerReplication = writerStore.replicate(true)
+  const readerReplication = readerStore.replicate(false)
+  writerReplication.pipe(readerReplication).pipe(writerReplication)
+
+  const client = Object.create(QVACRegistryClient.prototype)
+  client.logger = logger()
+  client.hyperswarm = {
+    join() {},
+    flush() {
+      return Promise.resolve()
+    }
+  }
+  client._ensureMetadata = () => Promise.resolve()
+  client.getModel = () =>
+    Promise.resolve({
+      name: 'storage-test',
+      blobBinding: { coreKey, ...pointer }
+    })
+  client._getBlobsCore = () => Promise.resolve({ core: readerCore, blobs: readerBlobs })
+
+  const controller = cancellation()
+  let downloadedBlocks = 0
+  readerCore.on('download', () => {
+    downloadedBlocks++
+    if (downloadedBlocks === CANCEL_AFTER_BLOCKS) controller.abort()
+  })
+
+  await t.exception(
+    () =>
+      client.downloadModel('models/storage-test.gguf', 's3', {
+        outputFile,
+        signal: controller.signal,
+        maxRetries: 1
+      }),
+    /Download cancelled/,
+    'the download is cancelled while blocks are in flight'
+  )
+
+  readerReplication.destroy()
+  writerReplication.destroy()
+  await readerStore.close()
+  await writerBlobs.close()
+  await writerCore.close()
+  await writerStore.close()
+
+  t.ok(downloadedBlocks > CANCEL_AFTER_BLOCKS, 'in-flight blocks arrived after cancellation')
+  t.is(
+    await countBlocks(readerStorage, coreKey, pointer),
+    0,
+    'no downloaded blocks remain after reopening the corestore'
+  )
+})
+
+test('drain timeout logs remaining peer counters with real Hypercore peers', async (t) => {
+  const root = await tmp(t)
+  const writerStorage = path.join(root, 'writer')
+  const readerStorage = path.join(root, 'reader')
+  const outputFile = path.join(root, 'partial.gguf')
+
+  const writerStore = new Corestore(writerStorage)
+  const writerCore = writerStore.get({ name: 'blobs' })
+  await writerCore.ready()
+  const coreKey = writerCore.key
+  const writerBlobs = new Hyperblobs(writerCore)
+  await writerBlobs.ready()
+  const pointer = await writerBlobs.put(Buffer.alloc(PAYLOAD_BYTES, 1))
+
+  const readerStore = new Corestore(readerStorage)
+  const readerCore = readerStore.get({ key: coreKey })
+  await readerCore.ready()
+  const readerBlobs = new Hyperblobs(readerCore)
+  await readerBlobs.ready()
+
+  const writerReplication = writerStore.replicate(true)
+  const readerReplication = readerStore.replicate(false)
+  writerReplication.pipe(readerReplication).pipe(writerReplication)
+
+  // Hold verify long enough that drain hits the best-effort timeout while
+  // peer.dataProcessing is still nonzero.
+  const VERIFY_HOLD_MS = 250
+  const originalVerify = readerCore.core.verify.bind(readerCore.core)
+  readerCore.core.verify = async function (...args) {
+    await new Promise((resolve) => setTimeout(resolve, VERIFY_HOLD_MS))
+    return originalVerify(...args)
+  }
+
+  const warnings = []
+  const client = Object.create(QVACRegistryClient.prototype)
+  client.logger = {
+    info() {},
+    debug() {},
+    warn(msg, data) {
+      warnings.push({ msg, data })
+    },
+    error() {}
+  }
+  client._inflightDrainTimeoutMs = 40
+  client._inflightDrainPollMs = 5
+  client.hyperswarm = {
+    join() {},
+    flush() {
+      return Promise.resolve()
+    }
+  }
+  client._ensureMetadata = () => Promise.resolve()
+  client.getModel = () =>
+    Promise.resolve({
+      name: 'storage-timeout-test',
+      blobBinding: { coreKey, ...pointer }
+    })
+  client._getBlobsCore = () => Promise.resolve({ core: readerCore, blobs: readerBlobs })
+
+  const controller = cancellation()
+  let downloadedBlocks = 0
+  readerCore.on('download', () => {
+    downloadedBlocks++
+    if (downloadedBlocks === CANCEL_AFTER_BLOCKS) controller.abort()
+  })
+
+  await t.exception(
+    () =>
+      client.downloadModel('models/storage-timeout-test.gguf', 's3', {
+        outputFile,
+        signal: controller.signal,
+        maxRetries: 1
+      }),
+    /Download cancelled/,
+    'the download is cancelled while verify is still held'
+  )
+
+  readerReplication.destroy()
+  writerReplication.destroy()
+  await readerStore.close()
+  await writerBlobs.close()
+  await writerCore.close()
+  await writerStore.close()
+
+  const timeoutWarning = warnings.find((entry) =>
+    entry.msg.includes('Timed out waiting for in-flight blob blocks')
+  )
+  t.ok(timeoutWarning, 'drain timeout is logged against real peers')
+  t.ok(Array.isArray(timeoutWarning.data.peers), 'timeout warning includes peer counters')
+  t.ok(
+    timeoutWarning.data.peers.some((peer) => peer.inflight > 0 || peer.dataProcessing > 0),
+    'logged counters show work still outstanding at timeout'
+  )
+  t.ok(downloadedBlocks >= CANCEL_AFTER_BLOCKS, 'cancellation fired after blocks started arriving')
+})