feat(relay): fence and drain file and git response streams on shutdown

src/relay/fs-handler-file-read.ts

@@ -2,8 +2,9 @@ import { open, readFile, stat } from 'node:fs/promises'
 import type { FileHandle } from 'node:fs/promises'
 import { extname } from 'node:path'
 import type { RelayDispatcher, RequestContext } from './dispatcher'
-import { MAX_CONCURRENT_STREAMS, STREAM_ACK_WINDOW_CHUNKS, STREAM_CHUNK_SIZE } from './protocol'
-import { TooManyStreamsError, type RelayStreamRegistry } from './fs-stream-registry'
+import { STREAM_ACK_WINDOW_CHUNKS, STREAM_CHUNK_SIZE } from './protocol'
+import type { RelayStreamRegistry } from './fs-stream-registry'
+import { reserveTerminalFrameSlot } from './fs-stream-terminal-frame-slots'
 import {
   BINARY_PROBE_BYTES,
   IMAGE_MIME_TYPES,
@@ -77,6 +78,21 @@ export async function readRelayFileStreamMetadata(
   registry: RelayStreamRegistry,
   context: RequestContext,
   pumpOptions?: StreamPumpOptions
+): Promise<StreamMetadata> {
+  const finish = registry.beginOperation()
+  try {
+    return await prepareRelayFileStream(filePath, dispatcher, registry, context, pumpOptions)
+  } finally {
+    finish()
+  }
+}
+
+async function prepareRelayFileStream(
+  filePath: string,
+  dispatcher: RelayDispatcher,
+  registry: RelayStreamRegistry,
+  context: RequestContext,
+  pumpOptions?: StreamPumpOptions
 ): Promise<StreamMetadata> {
   const stats = await stat(filePath)
   const mimeType = IMAGE_MIME_TYPES[extname(filePath).toLowerCase()]
@@ -99,7 +115,10 @@ export async function readRelayFileStreamMetadata(
   // Why: unlike the legacy single-shot path, streaming does not read the full
   // buffer before classifying content. Probe every unknown file so small binary
   // files do not get decoded as UTF-8 text over SSH.
-  if (!mimeType && (await isBinaryFilePrefix(filePath))) {
+  if (
+    !mimeType &&
+    (await isBinaryFilePrefix(filePath, (handle) => registry.releaseUnregisteredHandle(handle)))
+  ) {
     return { totalSize: 0, isBinary: true, empty: true }
   }
 
@@ -112,8 +131,13 @@ export async function readRelayFileStreamMetadata(
     handle = await open(filePath, 'r')
     streamId = registry.register(handle)
   } catch (err) {
-    await handle?.close()
-    releaseTerminalFrameSlot()
+    try {
+      if (handle) {
+        await registry.releaseUnregisteredHandle(handle)
+      }
+    } finally {
+      releaseTerminalFrameSlot()
+    }
     throw err
   }
 
@@ -123,6 +147,7 @@ export async function readRelayFileStreamMetadata(
   // setImmediate kicks the pump off the metadata-response task so the client
   // sees the response before the first chunk frame.
   const resolvedPumpOptions = pumpOptions ?? { paceWithAcks: false }
+  const finishPump = registry.beginOperation()
   setImmediate(() => {
     void pumpChunks(
       streamId,
@@ -133,6 +158,10 @@ export async function readRelayFileStreamMetadata(
       resolvedPumpOptions,
       releaseTerminalFrameSlot
     )
+      .catch((error: unknown) => {
+        process.stderr.write(`[relay] stream cleanup failed id=${streamId}: ${String(error)}\n`)
+      })
+      .finally(finishPump)
   })
 
   return {
@@ -146,48 +175,6 @@ export async function readRelayFileStreamMetadata(
   }
 }
 
-/**
- * Terminal frames (fs.streamEnd/fs.streamError) ride the control lane, which does not
- * drop on overflow — it destroys the link at 256 queued frames / 1 MB. A stream's
- * registry slot is gone the moment its last chunk is read, so on a socket that is not
- * draining, back-to-back reads could stack one undelivered terminal frame each until
- * that budget blew. Holding the slot until the frame settles keeps the number of queued
- * terminal frames at MAX_CONCURRENT_STREAMS, far below the killing threshold; the
- * overflow now costs one refused read (TooManyStreams, which clients already handle)
- * instead of the whole connection.
- *
- * Counted per client, because the control queue this protects is per client: one peer whose
- * socket stopped draining must not refuse reads for every other peer on the same relay.
- */
-const pendingTerminalFramesByClient = new WeakMap<RelayStreamRegistry, Map<number, number>>()
-
-function reserveTerminalFrameSlot(registry: RelayStreamRegistry, clientId: number): () => void {
-  let byClient = pendingTerminalFramesByClient.get(registry)
-  if (!byClient) {
-    byClient = new Map()
-    pendingTerminalFramesByClient.set(registry, byClient)
-  }
-  const pending = byClient.get(clientId) ?? 0
-  if (pending >= MAX_CONCURRENT_STREAMS) {
-    throw new TooManyStreamsError()
-  }
-  byClient.set(clientId, pending + 1)
-  let released = false
-  return () => {
-    if (released) {
-      return
-    }
-    released = true
-    const remaining = (byClient.get(clientId) ?? 1) - 1
-    // Drop the entry at zero so a long-lived registry cannot accumulate one per detached client.
-    if (remaining <= 0) {
-      byClient.delete(clientId)
-      return
-    }
-    byClient.set(clientId, remaining)
-  }
-}
-
 async function pumpChunks(
   streamId: number,
   totalSize: number,
@@ -302,6 +289,11 @@ async function pumpChunks(
           releaseTerminalFrameSlot
         )
       }
+      if (registry.isAborted(streamId)) {
+        endReason = 'aborted'
+      } else if (context.isStale()) {
+        endReason = 'stale'
+      }
       if (endReason === 'end') {
         publishTerminal('fs.streamEnd', { streamId })
         process.stderr.write(`[relay] stream end id=${streamId}\n`)
@@ -325,9 +317,12 @@ async function pumpChunks(
   } finally {
     // Why: the fd goes back first — a terminal frame that can never be delivered must not
     // strand it. Cancelled/stale streams publish nothing, so nothing else frees their slot.
-    await registry.release(streamId)
-    if (!slotReleaseDeferred) {
-      releaseTerminalFrameSlot()
+    try {
+      await registry.release(streamId)
+    } finally {
+      if (!slotReleaseDeferred) {
+        releaseTerminalFrameSlot()
+      }
     }
   }
 }

src/relay/fs-handler-stream.test.ts

@@ -492,6 +492,29 @@ describe('FsHandler readFileStream', () => {
     await waitFor(() => terminalFrames().length === MAX_CONCURRENT_STREAMS + 2)
   })
 
+  it('holds file stream shutdown until an undelivered terminal frame settles', async () => {
+    const filePath = path.join(tmpDir, 'shutdown-terminal.png')
+    writeFileSync(filePath, Buffer.alloc(STREAM_CHUNK_SIZE, 0x42))
+    dispatcher.holdControlSettlements()
+    const context = { clientId: 4, isStale: () => false }
+    await dispatcher.callRequest('fs.readFileStream', { filePath }, context)
+    await waitFor(() => collectStream(dispatcher).end !== null)
+
+    let drained = false
+    const drain = handler.disposeFileStreams().then(() => {
+      drained = true
+    })
+    await flush(10)
+    expect(drained).toBe(false)
+    await expect(
+      dispatcher.callRequest('fs.readFileStream', { filePath }, context)
+    ).rejects.toThrow('relay_file_stream_shutdown_fenced')
+
+    dispatcher.settleHeldControlFrames()
+    await drain
+    expect(drained).toBe(true)
+  })
+
   it('rejects the 17th concurrent stream with TooManyStreams', async () => {
     const paths: string[] = []
     for (let i = 0; i < 17; i++) {

src/relay/fs-handler-utils.ts

@@ -64,14 +64,18 @@ export function isBinaryBuffer(buffer: Buffer): boolean {
   return false
 }
 
-export async function isBinaryFilePrefix(filePath: string): Promise<boolean> {
+export async function isBinaryFilePrefix(
+  filePath: string,
+  releaseHandle: (handle: Awaited<ReturnType<typeof open>>) => Promise<void> = (handle) =>
+    handle.close()
+): Promise<boolean> {
   const handle = await open(filePath, 'r')
   try {
     const probe = Buffer.alloc(BINARY_PROBE_BYTES)
     const { bytesRead } = await handle.read(probe, 0, probe.length, 0)
     return isBinaryBuffer(probe.subarray(0, bytesRead))
   } finally {
-    await handle.close()
+    await releaseHandle(handle)
   }
 }
 

src/relay/fs-handler.ts

@@ -267,6 +267,12 @@ export class FsHandler {
 
   dispose(): void {
     this.watchRegistry.dispose()
-    void this.streamRegistry.disposeAll()
+    void this.disposeFileStreams().catch((error: unknown) => {
+      process.stderr.write(`[relay] file stream shutdown failed: ${String(error)}\n`)
+    })
+  }
+
+  disposeFileStreams(): Promise<void> {
+    return this.streamRegistry.disposeAll()
   }
 }

src/relay/fs-stream-registry-shutdown.test.ts

@@ -0,0 +1,171 @@
+import type { FileHandle } from 'node:fs/promises'
+import { describe, expect, it, vi } from 'vitest'
+import { RelayStreamRegistry } from './fs-stream-registry'
+
+function deferred(): { promise: Promise<void>; resolve: () => void } {
+  let resolve!: () => void
+  const promise = new Promise<void>((done) => {
+    resolve = done
+  })
+  return { promise, resolve }
+}
+
+async function flushTasks(): Promise<void> {
+  await new Promise<void>((resolve) => setImmediate(resolve))
+}
+
+function handle(close: ReturnType<typeof vi.fn>): FileHandle {
+  // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the registry only ever calls close() on a handle.
+  return { close } as unknown as FileHandle
+}
+
+describe('RelayStreamRegistry shutdown', () => {
+  it('retains failed late pre-registration cleanup after the initial disposal snapshot', async () => {
+    const registry = new RelayStreamRegistry()
+    const finishMetadata = registry.beginOperation()
+    const failure = new Error('late handle close failed')
+    const close = vi.fn().mockRejectedValueOnce(failure).mockResolvedValue(undefined)
+    const result = registry.disposeAll().catch((error: unknown) => error)
+    await flushTasks()
+    const openedHandle = handle(close)
+    expect(() => registry.register(openedHandle)).toThrow('relay_file_stream_shutdown_fenced')
+    await expect(registry.releaseUnregisteredHandle(openedHandle)).rejects.toBe(failure)
+    finishMetadata()
+    expect(await result).toMatchObject({
+      message: 'relay_file_stream_shutdown_incomplete',
+      errors: [failure]
+    })
+    expect(registry.size()).toBe(1)
+    await registry.disposeAll()
+    expect(registry.size()).toBe(0)
+    expect(close).toHaveBeenCalledTimes(2)
+  })
+
+  it('waits for a late pre-registration close owned by an admitted metadata operation', async () => {
+    const registry = new RelayStreamRegistry()
+    const finishMetadata = registry.beginOperation()
+    const pending = deferred()
+    const close = vi.fn(() => pending.promise)
+    const finished = vi.fn()
+    const drain = registry.disposeAll().then(finished)
+    await flushTasks()
+    const cleanup = registry.releaseUnregisteredHandle(handle(close)).finally(finishMetadata)
+    await flushTasks()
+    expect(finished).not.toHaveBeenCalled()
+    pending.resolve()
+    await Promise.all([cleanup, drain])
+    expect(registry.size()).toBe(0)
+  })
+
+  it('joins a release already waiting for the file handle to close', async () => {
+    const registry = new RelayStreamRegistry()
+    const pending = deferred()
+    const close = vi.fn(() => pending.promise)
+    const id = registry.register(handle(close))
+    const release = registry.release(id)
+    let disposed = false
+    const dispose = registry.disposeAll().then(() => {
+      disposed = true
+    })
+
+    try {
+      await flushTasks()
+      expect(disposed).toBe(false)
+      expect(close).toHaveBeenCalledTimes(1)
+      expect(registry.isAborted(id)).toBe(true)
+    } finally {
+      pending.resolve()
+      await Promise.all([release, dispose])
+    }
+    expect(disposed).toBe(true)
+    expect(registry.size()).toBe(0)
+  })
+
+  it('coalesces concurrent release calls until the single close settles', async () => {
+    const registry = new RelayStreamRegistry()
+    const pending = deferred()
+    const close = vi.fn(() => pending.promise)
+    const id = registry.register(handle(close))
+    const first = registry.release(id)
+    let secondSettled = false
+    const second = registry.release(id).then(() => {
+      secondSettled = true
+    })
+
+    try {
+      await flushTasks()
+      expect(close).toHaveBeenCalledTimes(1)
+      expect(secondSettled).toBe(false)
+    } finally {
+      pending.resolve()
+      await Promise.all([first, second])
+    }
+    expect(registry.size()).toBe(0)
+  })
+
+  it('retains a failed close for a later disposal retry', async () => {
+    const registry = new RelayStreamRegistry()
+    const failure = Object.assign(new Error('close failed'), { code: 'EIO' })
+    const close = vi.fn().mockRejectedValueOnce(failure).mockResolvedValueOnce(undefined)
+    const id = registry.register(handle(close))
+
+    await expect(registry.release(id)).rejects.toBe(failure)
+    expect(registry.size()).toBe(1)
+    expect(registry.isAborted(id)).toBe(true)
+    await registry.disposeAll()
+
+    expect(close).toHaveBeenCalledTimes(2)
+    expect(registry.size()).toBe(0)
+  })
+
+  it('waits for every close before rejecting disposal and retries only failed handles', async () => {
+    const registry = new RelayStreamRegistry()
+    const failure = Object.assign(new Error('close failed'), { code: 'EIO' })
+    const failedClose = vi.fn().mockRejectedValueOnce(failure).mockResolvedValueOnce(undefined)
+    const pending = deferred()
+    const pendingClose = vi.fn(() => pending.promise)
+    registry.register(handle(failedClose))
+    registry.register(handle(pendingClose))
+    let settled = false
+    const result = registry.disposeAll().then(
+      () => {
+        settled = true
+        return undefined
+      },
+      (error: unknown) => {
+        settled = true
+        return error
+      }
+    )
+
+    try {
+      await flushTasks()
+      expect(settled).toBe(false)
+      expect(failedClose).toHaveBeenCalledTimes(1)
+      expect(pendingClose).toHaveBeenCalledTimes(1)
+    } finally {
+      pending.resolve()
+    }
+    const error = await result
+    expect(error).toBeInstanceOf(AggregateError)
+    expect(error).toMatchObject({ errors: [failure] })
+    expect(registry.size()).toBe(1)
+
+    await registry.disposeAll()
+    expect(failedClose).toHaveBeenCalledTimes(2)
+    expect(pendingClose).toHaveBeenCalledTimes(1)
+    expect(registry.size()).toBe(0)
+  })
+
+  it('tolerates EBADF as an already-closed handle', async () => {
+    const registry = new RelayStreamRegistry()
+    const close = vi.fn().mockRejectedValue(Object.assign(new Error('closed'), { code: 'EBADF' }))
+    const id = registry.register(handle(close))
+
+    await registry.release(id)
+    await registry.disposeAll()
+
+    expect(close).toHaveBeenCalledTimes(1)
+    expect(registry.size()).toBe(0)
+  })
+})