diff --git a/lib/storage.ts b/lib/storage.ts index d40ce61..afe6826 100644 --- a/lib/storage.ts +++ b/lib/storage.ts @@ -1,11 +1,11 @@ /* eslint-disable ts/method-signature-style */ import type { ContainerClient } from '@azure/storage-blob' import type { Kysely } from 'kysely' +import type { Writable } from 'node:stream' import type { Database, StorageLocation } from './db' import type { Reservation } from './helpers' import type { Env } from './schemas' import { randomUUID } from 'node:crypto' -import { once } from 'node:events' import { createReadStream, createWriteStream } from 'node:fs' import fs from 'node:fs/promises' import { Agent } from 'node:https' @@ -66,6 +66,23 @@ function escapeLikePattern(value: string) { .replaceAll('_', String.raw`\_`) } +/** + * Resolves once `stream` wants more data or is closed. A destroyed stream emits + * no 'drain', and no 'error' when destroyed without one, so waiting on 'drain' + * alone hangs forever once its reader is gone. + */ +function drainOrClose(stream: Writable) { + const { promise, resolve } = Promise.withResolvers() + const done = () => { + stream.off('drain', done) + stream.off('close', done) + resolve() + } + stream.on('drain', done) + stream.on('close', done) + return promise +} + export class ObjectNotFoundError extends Error { constructor(objectName: string) { super(`Object not found in storage: ${objectName}`) @@ -198,16 +215,25 @@ export class Storage { ) { if (location.partsDeletedAt) throw new Error('No parts to feed') + // Either sink can be destroyed mid-download: the route destroys the response + // when the client hangs up, and a failed merge upload destroys the merger. + // Waiting on a destroyed sink's 'drain' would park the pump forever, pinning + // the current part's backend read and renewing the merge lease. Keep feeding + // whichever sink is still open, and stop reading once both are gone. + const sinks = [responseStream, mergerStream] for await (const chunk of this.streamParts(location)) { - // attach both drain listeners simultaneously to prevent a race condition where the second stream being faster hangs forever + // attach every drain wait before awaiting any, so a faster sink's 'drain' is never missed const drains = [] - if (!responseStream.write(chunk)) drains.push(once(responseStream, 'drain')) - if (!mergerStream.write(chunk)) drains.push(once(mergerStream, 'drain')) + for (const sink of sinks) { + if (sink.destroyed) continue + if (!sink.write(chunk) && !sink.destroyed) drains.push(drainOrClose(sink)) + } await Promise.all(drains) + // Leaving the loop makes `streamParts` destroy the current part stream. + if (sinks.every((sink) => sink.destroyed)) return } - responseStream.end() - mergerStream.end() + for (const sink of sinks) if (!sink.destroyed) sink.end() await globalThis.gc?.() } diff --git a/tests/download-merge-sink-gone.test.ts b/tests/download-merge-sink-gone.test.ts new file mode 100644 index 0000000..151fbd7 --- /dev/null +++ b/tests/download-merge-sink-gone.test.ts @@ -0,0 +1,158 @@ +import { Buffer } from 'node:buffer' +import { randomUUID } from 'node:crypto' +import { once } from 'node:events' +import { Readable } from 'node:stream' +import { setTimeout as sleep } from 'node:timers/promises' +import { afterEach, describe, expect, test, vi } from 'vitest' +import { getDatabase } from '~/lib/db' +import { getStorage } from '~/lib/storage' + +const PART_COUNT = 3 +const PART_BYTES = 1024 * 1024 + +async function createUnmergedEntry() { + const storage = await getStorage() + const db = await getDatabase() + const key = `merge-sink-gone-${randomUUID()}` + const entry = { key, version: 'v1', scope: 'refs/heads/main', repoId: '123' } + const upload = await storage.createUpload(entry) + if (!upload) throw new Error('createUpload returned nothing') + for (let i = 0; i < PART_COUNT; i++) + await storage.uploadPart(upload.id, i, Readable.from([Buffer.alloc(PART_BYTES, i)])) + await storage.completeUpload(entry) + return db + .selectFrom('cache_entries') + .where('key', '=', key) + .select('locationId') + .executeTakeFirstOrThrow() +} + +/** Records every part stream the download opens, so the test can see each one released. */ +async function trackPartStreams() { + const storage = await getStorage() + const original = storage.adapter.createDownloadStream.bind(storage.adapter) + const parts: Readable[] = [] + vi.spyOn(storage.adapter, 'createDownloadStream').mockImplementation(async (name, range) => { + const download = await original(name, range) + if (name.includes('/parts/')) parts.push(download.stream) + return download + }) + return parts +} + +/** Fails the background merge upload without reading it, as a broken S3 upload would. */ +async function failMergeUpload() { + const storage = await getStorage() + vi.spyOn(storage.adapter, 'uploadStream').mockRejectedValue(new Error('merge upload failed')) +} + +/** Reads one chunk, then hangs up the way the route does when the client leaves. */ +async function leaveAfterFirstChunk(stream: Readable) { + await once(stream, 'readable') + expect(stream.read()).not.toBeNull() + stream.destroy() +} + +async function settles(promise: Promise, timeoutMs = 5000) { + return Promise.race([promise.then(() => 'settled'), sleep(timeoutMs, 'still pending')]) +} + +async function locationState(locationId: string) { + const db = await getDatabase() + const { mergedAt } = await db + .selectFrom('storage_locations') + .where('id', '=', locationId) + .select('mergedAt') + .executeTakeFirstOrThrow() + const mergeLeases = await db + .selectFrom('merge_leases') + .where('storageLocationId', '=', locationId) + .select('token') + .execute() + return { mergedAt, mergeLeases } +} + +afterEach(() => { + vi.restoreAllMocks() +}) + +/** + * The first download of an unmerged entry tees its parts into the client + * response and the background merge upload. When one of them is destroyed + * (client hung up, merge upload failed), a pump waiting for 'drain' parks + * forever: a destroyed stream emits neither 'drain' nor 'error'. That holds the + * part's backend read (an S3 socket), starves the merge, and keeps renewing + * the merge lease for the life of the process. + */ +describe('first download of an unmerged entry', () => { + test('client leaving mid-stream still completes the merge and releases the parts', async () => { + const storage = await getStorage() + const { locationId } = await createUnmergedEntry() + const parts = await trackPartStreams() + + const download = await storage.download(locationId) + if (!download) throw new Error('download returned nothing') + const { stream } = download + await leaveAfterFirstChunk(stream) + + expect(await settles(storage.waitForOngoingMerges()), 'merge settled').toBe('settled') + const { mergedAt, mergeLeases } = await locationState(locationId) + expect(mergedAt, 'merge completed').not.toBeNull() + expect(mergeLeases, 'merge lease released').toHaveLength(0) + expect(parts).toHaveLength(PART_COUNT) + expect( + parts.filter((part) => !part.destroyed), + 'part streams left open', + ).toHaveLength(0) + }, 20_000) + + test('a failed merge upload still serves the whole entry to the client', async () => { + const storage = await getStorage() + const { locationId } = await createUnmergedEntry() + await failMergeUpload() + + const download = await storage.download(locationId) + if (!download) throw new Error('download returned nothing') + const { stream } = download + let received = 0 + const consumed = (async () => { + for await (const chunk of stream) received += (chunk as Buffer).length + })() + + expect(await settles(consumed), 'client download finished').toBe('settled') + expect(received).toBe(PART_COUNT * PART_BYTES) + expect(await settles(storage.waitForOngoingMerges()), 'merge settled').toBe('settled') + const { mergedAt, mergeLeases } = await locationState(locationId) + expect(mergedAt, 'failed merge not marked complete').toBeNull() + expect(mergeLeases, 'merge lease released').toHaveLength(0) + }, 20_000) + + test('client and merge both gone stops reading and destroys the part stream', async () => { + const storage = await getStorage() + const { locationId } = await createUnmergedEntry() + await failMergeUpload() + // An unending part: only being destroyed lets it go, as with a GetObject body. + const part = Readable.from( + (async function* () { + for (;;) { + await sleep(1) + yield Buffer.alloc(64 * 1024) + } + })(), + ) + vi.spyOn(storage.adapter, 'createDownloadStream').mockResolvedValue({ stream: part }) + + const download = await storage.download(locationId) + if (!download) throw new Error('download returned nothing') + const { stream } = download + await leaveAfterFirstChunk(stream) + + await vi.waitFor(() => expect(part.destroyed, 'part stream destroyed').toBe(true), { + timeout: 5000, + interval: 25, + }) + expect(await settles(storage.waitForOngoingMerges()), 'merge settled').toBe('settled') + const { mergeLeases } = await locationState(locationId) + expect(mergeLeases, 'merge lease released').toHaveLength(0) + }, 20_000) +})