diff --git a/packages/registry-server/client/lib/client.js b/packages/registry-server/client/lib/client.js index 655e58383c..b6a55ffd12 100644 --- a/packages/registry-server/client/lib/client.js +++ b/packages/registry-server/client/lib/client.js @@ -304,7 +304,7 @@ class QVACRegistryClient extends ReadyResource { throw new Error(`Invalid options: ${typeof options}`) } - let core, blobs + let core, blobs, blockStart, blockEnd, rangeDownload try { this.logger.info('Downloading model', { path, source }) @@ -344,13 +344,13 @@ class QVACRegistryClient extends ReadyResource { const totalSize = model.blobBinding.byteLength - const rangeDownload = core.download({ + rangeDownload = core.download({ start: model.blobBinding.blockOffset, length: model.blobBinding.blockLength }) - const blockStart = model.blobBinding.blockOffset - const blockEnd = blockStart + model.blobBinding.blockLength + blockStart = model.blobBinding.blockOffset + blockEnd = blockStart + model.blobBinding.blockLength let artifact if (options.outputFile) { @@ -369,10 +369,7 @@ class QVACRegistryClient extends ReadyResource { ) artifact = { path: options.outputFile, totalSize } - rangeDownload.destroy() - await this._clearBlobBlocks(core, blockStart, blockEnd) - if (blobs) await blobs.close() - if (core) await core.close() + await this._releaseDownload(core, blobs, rangeDownload, blockStart, blockEnd) } else { const stream = blobs.createReadStream(model.blobBinding, { wait: true, @@ -380,27 +377,7 @@ class QVACRegistryClient extends ReadyResource { }) artifact = { stream, totalSize } - const cleanup = async () => { - rangeDownload.destroy() - await this._clearBlobBlocks(core, blockStart, blockEnd) - if (blobs) { - try { - await blobs.close() - } catch (cleanupError) { - this.logger.warn('Error closing blob instance', { error: cleanupError.message }) - } - } - if (core) { - try { - await core.close() - } catch (cleanupError) { - this.logger.warn('Error closing blob core', { error: cleanupError.message }) - } - } - this.logger.debug('Blob resources closed after stream end') - } - - stream.once('end', cleanup) + this._releaseOnStreamEnd(stream, core, blobs, rangeDownload, blockStart, blockEnd) } this.logger.info('Model downloaded successfully') @@ -412,20 +389,7 @@ class QVACRegistryClient extends ReadyResource { } catch (error) { this.logger.error('Error downloading model', error) - if (blobs) { - try { - await blobs.close() - } catch (cleanupError) { - this.logger.warn('Error closing blob instance on error', { error: cleanupError.message }) - } - } - if (core) { - try { - await core.close() - } catch (cleanupError) { - this.logger.warn('Error closing blob core on error', { error: cleanupError.message }) - } - } + await this._releaseDownload(core, blobs, rangeDownload, blockStart, blockEnd) throw error } @@ -449,7 +413,7 @@ class QVACRegistryClient extends ReadyResource { throw new Error(`Invalid options: ${typeof options}`) } - let core, blobs + let core, blobs, blockStart, blockEnd, rangeDownload try { this.logger.info('Downloading blob directly', { @@ -487,10 +451,10 @@ class QVACRegistryClient extends ReadyResource { } const totalSize = blobBinding.byteLength - const blockStart = pointer.blockOffset - const blockEnd = blockStart + pointer.blockLength + blockStart = pointer.blockOffset + blockEnd = blockStart + pointer.blockLength - const rangeDownload = core.download({ + rangeDownload = core.download({ start: pointer.blockOffset, length: pointer.blockLength }) @@ -510,10 +474,7 @@ class QVACRegistryClient extends ReadyResource { ) artifact = { path: options.outputFile, totalSize } - rangeDownload.destroy() - await this._clearBlobBlocks(core, blockStart, blockEnd) - if (blobs) await blobs.close() - if (core) await core.close() + await this._releaseDownload(core, blobs, rangeDownload, blockStart, blockEnd) } else { const stream = blobs.createReadStream(pointer, { wait: true, @@ -521,27 +482,7 @@ class QVACRegistryClient extends ReadyResource { }) artifact = { stream, totalSize } - const cleanup = async () => { - rangeDownload.destroy() - await this._clearBlobBlocks(core, blockStart, blockEnd) - if (blobs) { - try { - await blobs.close() - } catch (e) { - this.logger.warn('Error closing blob instance', { error: e.message }) - } - } - if (core) { - try { - await core.close() - } catch (e) { - this.logger.warn('Error closing blob core', { error: e.message }) - } - } - this.logger.debug('Blob resources closed after stream end') - } - - stream.once('end', cleanup) + this._releaseOnStreamEnd(stream, core, blobs, rangeDownload, blockStart, blockEnd) } this.logger.info('Blob download complete (direct)') @@ -550,26 +491,50 @@ class QVACRegistryClient extends ReadyResource { } catch (error) { this.logger.error('Error downloading blob directly', error) - if (blobs) { - try { - await blobs.close() - } catch (e) { - this.logger.warn('Error closing blob instance on error', { error: e.message }) - } - } - if (core) { - try { - await core.close() - } catch (e) { - this.logger.warn('Error closing blob core on error', { error: e.message }) - } - } + await this._releaseDownload(core, blobs, rangeDownload, blockStart, blockEnd) throw error } } - async _clearBlobBlocks(core, start, end) { + _releaseOnStreamEnd (stream, core, blobs, rangeDownload, blockStart, blockEnd) { + let released = false + + // 'close' also covers a destroyed or errored stream; on 'end' alone a + // cancelled stream download would never free its blocks. + const release = () => { + if (released) return + released = true + return this._releaseDownload(core, blobs, rangeDownload, blockStart, blockEnd) + .catch(e => this.logger.warn('Error releasing blob resources', { error: e.message })) + } + + stream.once('end', release) + stream.once('close', release) + } + + async _releaseDownload (core, blobs, rangeDownload, blockStart, blockEnd) { + // Stop replication before clearing to prevent blocks from being refetched. + if (rangeDownload) rangeDownload.destroy() + + if (core && blockStart != null) { + await this._clearBlobBlocks(core, blockStart, blockEnd) + } + if (blobs) { + try { await blobs.close() } catch (e) { + this.logger.warn('Error closing blob instance', { error: e.message }) + } + } + if (core) { + try { await core.close() } catch (e) { + this.logger.warn('Error closing blob core', { error: e.message }) + } + } + + this.logger.debug('Blob resources released') + } + + async _clearBlobBlocks (core, start, end) { try { const cleared = await core.clear(start, end, { diff: true }) await core.compact() 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 52fa36d4d0..ff64a1477f 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 @@ -6,17 +6,40 @@ const fs = require('#fs') const path = require('#path') const { withRetry } = require('../../utils/retry') +// Forwards the client's own log lines into TAP output when a test passes `t`, +// so the run shows which production branches actually executed. +function makeLogger(t) { + const emit = (level, msg, data) => { + if (!t) return + let detail = '' + if (data instanceof Error) detail = ` ${data.message}` + else if (data) detail = ` ${JSON.stringify(data).slice(0, 120)}` + t.comment(`[${level}] ${msg}${detail}`) + } + + return { + info(msg, data) { emit('info', msg, data) }, + debug(msg, data) { emit('debug', msg, data) }, + warn(msg, data) { emit('warn', msg, data) }, + error(msg, data) { emit('error', msg, data) } + } +} + +function now() { + return typeof performance !== 'undefined' ? performance.now() : Date.now() +} + // Builds a QVACRegistryClient instance WITHOUT running the constructor (which // would open a real Corestore and join a real swarm). Only the collaborators // that downloadModel touches are stubbed, so the real retry path runs. -function makeClient() { +function makeClient(t) { const QVACRegistryClient = require('../../lib/client') const client = Object.create(QVACRegistryClient.prototype) const events = [] client._events = events - client.logger = { info() {}, debug() {}, warn() {}, error() {} } + client.logger = makeLogger(t) const core = { discoveryKey: Buffer.alloc(32), @@ -37,8 +60,13 @@ function makeClient() { off() {} } // lunte-disable-next-line require-await - const blobs = { async close() {} } + const blobs = { + async close() {}, + createReadStream() { return client._stream } + } client._core = core + client._blobs = blobs + client._stream = fakeStream() client.hyperswarm = { suspended: false, @@ -69,6 +97,20 @@ function makeClient() { return client } +// Stands in for the hyperblobs read stream; emit() resolves once the release +// handlers it triggered have settled. +function fakeStream() { + const handlers = {} + return { + once(event, fn) { (handlers[event] = handlers[event] || []).push(fn) }, + emit(event) { + const fns = handlers[event] || [] + handlers[event] = [] + return Promise.all(fns.map(fn => fn())) + } + } +} + function requestTimeout() { const err = new Error('request timed out waiting for peers') err.code = 'REQUEST_TIMEOUT' @@ -267,6 +309,219 @@ test('downloadModel aborts the reconnect wait when the signal is cancelled', asy t.is(attempt, 1, 'no second attempt started after the cancel') }) +test('downloadModel clears cached blocks when the download fails', async t => { + const dir = await tmp(t) + const outputFile = path.join(dir, 'model.gguf') + + const client = makeClient(t) + const clears = [] + client._core.download = () => ({ + destroy () { + client._events.push('destroy') + t.comment('step: rangeDownload.destroy() called from the catch path') + } + }) + client._clearBlobBlocks = async (core, start, end) => { + client._events.push('clear') + clears.push({ start, end }) + t.comment(`step: _clearBlobBlocks(${start}, ${end}) called`) + } + client._streamBlobToFile = async () => { + client._events.push('stream') + t.comment('step: _streamBlobToFile rejecting to force the catch path') + throw new Error('Download cancelled') + } + + const started = now() + await t.exception( + () => client.downloadModel('models/tiny.gguf', 's3', { outputFile }), + /Download cancelled/, + 'the failed download rejects' + ) + t.comment(`timing: downloadModel failure+release took ${(now() - started).toFixed(1)}ms`) + + t.is(clears.length, 1, 'partial blocks cleared exactly once on the failure path') + t.alike(clears[0], { start: 0, end: 10 }, 'cleared the model block range') + t.ok( + client._events.indexOf('stream') < client._events.indexOf('clear'), + 'blocks are cleared after the download fails (in the catch)' + ) + const destroyedAt = client._events.indexOf('destroy') + t.ok( + destroyedAt !== -1 && destroyedAt < client._events.indexOf('clear'), + 'replication is stopped before clearing so blocks are not refetched' + ) +}) + +test('downloadModel clears cached blocks when the returned stream is destroyed', async t => { + const client = makeClient(t) + const clears = [] + client._core.download = () => ({ + destroy () { + client._events.push('destroy') + t.comment('step: rangeDownload.destroy() called from the stream release') + } + }) + client._clearBlobBlocks = async (core, start, end) => { + client._events.push('clear') + clears.push({ start, end }) + t.comment(`step: _clearBlobBlocks(${start}, ${end}) called`) + } + + const { artifact } = await client.downloadModel('models/tiny.gguf', 's3') + t.comment('step: stream artifact returned, release handlers bound') + + t.is(clears.length, 0, 'nothing cleared while the stream is still live') + + // A cancelled consumer destroys the stream, which emits 'close' without 'end'. + t.comment("step: emitting 'close' without 'end' (consumer destroyed the stream)") + const started = now() + await artifact.stream.emit('close') + t.comment(`timing: stream release took ${(now() - started).toFixed(1)}ms`) + + t.is(clears.length, 1, 'blocks cleared when the stream closes without ending') + t.alike(clears[0], { start: 0, end: 10 }, 'cleared the model block range') + const destroyedAt = client._events.indexOf('destroy') + t.ok( + destroyedAt !== -1 && destroyedAt < client._events.indexOf('clear'), + 'replication is stopped before clearing so blocks are not refetched' + ) +}) + +test('downloadModel releases the stream download exactly once', async t => { + const client = makeClient(t) + let clears = 0 + client._clearBlobBlocks = async () => { + clears++ + t.comment(`step: _clearBlobBlocks call #${clears}`) + } + + const { artifact } = await client.downloadModel('models/tiny.gguf', 's3') + + t.comment("step: emitting 'end' (normal completion)") + await artifact.stream.emit('end') + t.comment("step: emitting 'close' (follows end on an autodestroyed stream)") + await artifact.stream.emit('close') + + t.is(clears, 1, 'the end-then-close sequence releases only once') +}) + +test('downloadBlob clears cached blocks when the returned stream is destroyed', async t => { + const client = makeClient(t) + client.ready = async () => {} + const clears = [] + client._clearBlobBlocks = async (core, start, end) => { + clears.push({ start, end }) + t.comment(`step: _clearBlobBlocks(${start}, ${end}) called`) + } + + const { artifact } = await client.downloadBlob({ + coreKey: Buffer.alloc(32), + blockOffset: 3, + blockLength: 7, + byteLength: 700 + }) + + t.comment("step: emitting 'close' without 'end' on the direct blob stream") + await artifact.stream.emit('close') + + t.is(clears.length, 1, 'blocks cleared when the stream closes without ending') + t.alike(clears[0], { start: 3, end: 10 }, 'cleared the blob block range') +}) + +test('downloadBlob clears cached blocks when the download fails', async t => { + const dir = await tmp(t) + const outputFile = path.join(dir, 'blob.bin') + + const client = makeClient(t) + client.ready = async () => {} + const clears = [] + client._core.download = () => ({ + destroy () { + client._events.push('destroy') + t.comment('step: rangeDownload.destroy() called from the catch path') + } + }) + client._clearBlobBlocks = async (core, start, end) => { + client._events.push('clear') + clears.push({ start, end }) + t.comment(`step: _clearBlobBlocks(${start}, ${end}) called`) + } + client._streamBlobToFile = async () => { + t.comment('step: _streamBlobToFile rejecting to force the catch path') + throw new Error('Download cancelled') + } + + const blobBinding = { + coreKey: Buffer.alloc(32), + blockOffset: 3, + blockLength: 7, + byteLength: 700 + } + + await t.exception( + () => client.downloadBlob(blobBinding, { outputFile }), + /Download cancelled/, + 'the failed direct blob download rejects' + ) + + t.is(clears.length, 1, 'partial blocks cleared exactly once on the failure path') + t.alike(clears[0], { start: 3, end: 10 }, 'cleared the blob block range') + const destroyedAt = client._events.indexOf('destroy') + t.ok( + destroyedAt !== -1 && destroyedAt < client._events.indexOf('clear'), + 'replication is stopped before clearing so blocks are not refetched' + ) +}) + +test('downloadModel releases nothing when it fails before the core is opened', async t => { + const client = makeClient(t) + let clears = 0 + client._clearBlobBlocks = async () => { + clears++ + t.comment('step: _clearBlobBlocks called') + } + client.getModel = async () => { + t.comment('step: getModel returning null, failing before _getBlobsCore') + return null + } + + await t.exception( + () => client.downloadModel('models/missing.gguf', 's3'), + /Model not found/, + 'the missing model rejects' + ) + + t.is(clears, 0, 'nothing cleared when no core or block range exists yet') +}) + +test('downloadModel closes the core without clearing when the range is unknown', async t => { + const client = makeClient(t) + let clears = 0 + let closed = 0 + client._clearBlobBlocks = async () => { + clears++ + t.comment('step: _clearBlobBlocks called') + } + client._core.close = async () => { + closed++ + t.comment('step: core.close() called') + } + client._core.update = async () => { + t.comment('step: core.update() rejecting before the block range is computed') + throw new Error('core update failed') + } + + await t.exception( + () => client.downloadModel('models/tiny.gguf', 's3'), + /core update failed/, + 'the failed core update rejects' + ) + + t.is(clears, 0, 'no clear attempted while blockStart is still unassigned') + t.is(closed, 1, 'the opened core is still closed on the way out') +}) + // Locks the generic retry contract the download path relies on. test('withRetry retries only listed codes and stays bounded', async (t) => { let calls = 0