From fb940125be4bb0785578c4e92e3b046d96483589 Mon Sep 17 00:00:00 2001 From: yuranich Date: Tue, 28 Jul 2026 13:49:33 +0600 Subject: [PATCH] fix: drain in-flight blob blocks before cleanup --- packages/registry-server/client/lib/client.js | 27 ++- .../tests/unit/client.download-retry.test.js | 84 +++++++ .../unit/client.download-storage.test.js | 226 ++++++++++++++++++ 3 files changed, 336 insertions(+), 1 deletion(-) create mode 100644 packages/registry-server/client/tests/unit/client.download-storage.test.js diff --git a/packages/registry-server/client/lib/client.js b/packages/registry-server/client/lib/client.js index 9189095315..a8430bdb6d 100644 --- a/packages/registry-server/client/lib/client.js +++ b/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 }) diff --git a/packages/registry-server/client/tests/unit/client.download-retry.test.js b/packages/registry-server/client/tests/unit/client.download-retry.test.js index dc6837f310..8267502ef3 100644 --- a/packages/registry-server/client/tests/unit/client.download-retry.test.js +++ b/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 = [] diff --git a/packages/registry-server/client/tests/unit/client.download-storage.test.js b/packages/registry-server/client/tests/unit/client.download-storage.test.js new file mode 100644 index 0000000000..f1ef9276fa --- /dev/null +++ b/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') +})