From 9558a47ec591601ec070485271a6a25f9af454f9 Mon Sep 17 00:00:00 2001 From: Rui Chen Date: Sat, 24 Oct 2026 15:35:21 -0400 Subject: [PATCH] fix(download): stop the merge pump waiting on a destroyed stream The first download of an unmerged entry tees its parts into the client response and the background merge upload. A destroyed stream returns false from write() and never emits 'drain' (nor 'error' when destroyed without one), so the pump parks forever when: - the merge upload fails: `download()` destroys the merger stream, and the client's download stalls mid-body until the client times out; - the client hangs up while the pump waits on the merger: the next write to the destroyed response stream parks the pump, the merge never completes, and its lease keeps renewing. Either way the current part's backend read stays open (one S3 socket per occurrence), so enough of them fill the S3 socket pool. Wait for 'drain' or 'close', skip destroyed sinks, keep feeding the one that is still open, and stop reading once both are gone so `streamParts` destroys the current part stream. The regression tests fail on the previous pump with the filesystem and S3 drivers: the merge never settles after the client leaves, the client never finishes after the merge upload fails, and the part stream is never destroyed when both are gone. Signed-off-by: Rui Chen --- lib/storage.ts | 38 +++++- tests/download-merge-sink-gone.test.ts | 158 +++++++++++++++++++++++++ 2 files changed, 190 insertions(+), 6 deletions(-) create mode 100644 tests/download-merge-sink-gone.test.ts 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) +})