fix(peer, core): cancel stream bodies when a request ends before they are sent

packages/core/README.md

@@ -275,7 +275,7 @@ Use it when building a custom adapter, or when you need to know how a body will
 
 ## Utilities
 
-The package also exports a small set of helpers for common header and URL operations.
+The package also exports a small set of helpers for common header, URL, and body operations.
 
 ### Content-Disposition helpers
 
@@ -333,6 +333,20 @@ const [pathname, search, hash] = parseStandardUrl('/users/123?tab=settings#profi
 // hash => '#profile'
 ```
 
+### Cancelling a body
+
+`cancelStandardBody()` cancels a body that will not be consumed, whether it is never sent or never read. It cancels a `ReadableStream` with the given reason and calls `return()` on an async iterator, so their sources can clean up. Other bodies are left as is. It rejects if that cleanup fails, so you decide whether that is worth reporting.
+
+```ts
+import { cancelStandardBody } from '@standard-server/core'
+
+const body = await request.resolveBody()
+
+if (!authorized) {
+  await cancelStandardBody(body, new Error('Unauthorized'))
+}
+```
+
 ## Validators
 
 Runtime type guards are useful when requests or responses cross process, transport, or message boundaries.

packages/core/src/utils.test.ts

@@ -1,4 +1,5 @@
-import { flattenStandardHeader, generateContentDisposition, getFilenameFromContentDisposition, mergeStandardHeaders, parseStandardUrl, resolveStandardBodyHint } from './utils'
+import { AsyncIteratorClass } from '@standard-server/shared'
+import { cancelStandardBody, flattenStandardHeader, generateContentDisposition, getFilenameFromContentDisposition, mergeStandardHeaders, parseStandardUrl, resolveStandardBodyHint } from './utils'
 
 beforeEach(() => {
   vi.clearAllMocks()
@@ -174,6 +175,55 @@ describe('resolveStandardBodyHint', () => {
   })
 })
 
+describe('cancelStandardBody', () => {
+  it('cancels a ReadableStream with the reason', async () => {
+    const cancel = vi.fn()
+    const reason = new Error('reason')
+
+    await cancelStandardBody(new ReadableStream({ cancel }), reason)
+
+    expect(cancel).toHaveBeenCalledOnce()
+    expect(cancel).toHaveBeenCalledWith(reason)
+  })
+
+  it('returns an AsyncIterator', async () => {
+    const cleanup = vi.fn()
+
+    await cancelStandardBody(new AsyncIteratorClass(async () => ({ done: true, value: undefined }), cleanup))
+
+    expect(cleanup).toHaveBeenCalledOnce()
+    expect(cleanup).toHaveBeenCalledWith({ kind: 'cancelled' })
+  })
+
+  it('ignores an AsyncIterator without return()', async () => {
+    const iterator = {
+      next: async () => ({ done: true, value: undefined }),
+      [Symbol.asyncIterator]() {
+        return this
+      },
+    }
+
+    await expect(cancelStandardBody(iterator)).resolves.toBeUndefined()
+  })
+
+  it('ignores non-stream bodies', async () => {
+    await expect(cancelStandardBody({ key: 'val' })).resolves.toBeUndefined()
+    await expect(cancelStandardBody(new Blob(['x']))).resolves.toBeUndefined()
+  })
+
+  it('rejects when releasing fails', async () => {
+    const locked = new ReadableStream()
+    locked.getReader()
+
+    await expect(cancelStandardBody(locked)).rejects.toThrow(TypeError)
+
+    const error = new Error('cleanup failed')
+    await expect(cancelStandardBody(new AsyncIteratorClass(async () => ({ done: true, value: undefined }), async () => {
+      throw error
+    }))).rejects.toBe(error)
+  })
+})
+
 describe('mergeStandardHeaders', () => {
   afterEach(() => {
     expect(({} as any).polluted).toEqual(undefined)

packages/core/src/utils.ts

@@ -1,5 +1,5 @@
-import type { StandardBodyHint, StandardHeaders, StandardUrl } from './types'
-import { safeDecodeURIComponent, safeEncodeURIComponent, toArray } from '@standard-server/shared'
+import type { StandardBody, StandardBodyHint, StandardHeaders, StandardUrl } from './types'
+import { isAsyncIteratorObject, safeDecodeURIComponent, safeEncodeURIComponent, toArray } from '@standard-server/shared'
 
 export function generateContentDisposition(filename: string, type: 'inline' | 'attachment' = 'inline'): string {
   const encodedFilename = filename.replace(/[^\x20-\x7E]/g, '_').replace(/[\\"]/g, '\\$&')
@@ -111,6 +111,19 @@ export function resolveStandardBodyHint(headers: {
   return 'octet-stream'
 }
 
+/**
+ * Cancel a body that will not be consumed, so its stream or iterator source can clean up.
+ * Other bodies are left as is. Rejects if that cleanup fails.
+ */
+export async function cancelStandardBody(body: StandardBody, reason?: unknown): Promise<void> {
+  if (body instanceof ReadableStream) {
+    await body.cancel(reason)
+  }
+  else if (isAsyncIteratorObject(body)) {
+    await body.return?.()
+  }
+}
+
 export function mergeStandardHeaders(a: StandardHeaders, b: StandardHeaders): StandardHeaders {
   const merged = { ...a, ...b }
 

packages/peer/src/client.test.ts

@@ -264,6 +264,73 @@ describe('clientPeer', () => {
       await expect(peer.request(makeRequest())).rejects.toThrow(error)
     })
 
+    it('cancels an octet-stream request body when signal aborted during encode', async () => {
+      const controller = new AbortController()
+      const cancel = vi.fn()
+
+      const promise = peer.request(makeRequest({ body: new ReadableStream({ cancel }), signal: controller.signal }))
+      const error = new Error('aborted during encode')
+      controller.abort(error)
+
+      await expect(promise).rejects.toThrow(error)
+      await vi.waitFor(() => expect(cancel).toHaveBeenCalledOnce())
+      expect(cancel).toHaveBeenCalledWith(error)
+      expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['cancel'])
+    })
+
+    it('returns an event-stream request body when the peer is closed while encoding', async () => {
+      const cleanup = vi.fn()
+
+      const promise = peer.request(makeRequest({ body: new AsyncIteratorClass<unknown>(() => new Promise(() => {}), cleanup) }))
+      await peer.close()
+
+      await expect(promise).rejects.toThrow(AbortError)
+      await vi.waitFor(() => expect(cleanup).toHaveBeenCalledOnce())
+      expect(cleanup).toHaveBeenCalledWith({ kind: 'cancelled' })
+      expect(send).not.toHaveBeenCalled()
+    })
+
+    it('ignores a request body that fails to clean up after the request settled', async () => {
+      const controller = new AbortController()
+      const cancel = vi.fn(() => {
+        throw new Error('cleanup failed')
+      })
+
+      const promise = peer.request(makeRequest({ body: new ReadableStream({ cancel }), signal: controller.signal }))
+      const error = new Error('aborted during encode')
+      controller.abort(error)
+
+      // the request keeps its own outcome, and the failed release is not an unhandled rejection
+      await expect(promise).rejects.toBe(error)
+      await vi.waitFor(() => expect(cancel).toHaveBeenCalledOnce())
+    })
+
+    it('cancels the request body when send throws', async () => {
+      const error = new Error('send failed')
+      send.mockRejectedValueOnce(error)
+      const cancel = vi.fn()
+
+      await expect(peer.request(makeRequest({ body: new ReadableStream({ cancel }) }))).rejects.toThrow(error)
+      await vi.waitFor(() => expect(cancel).toHaveBeenCalledOnce())
+      expect(cancel).toHaveBeenCalledWith(error)
+    })
+
+    it('cancels the request body with the abort reason when signal aborted during send', async () => {
+      const controller = new AbortController()
+      const error = new Error('aborted during send')
+      send.mockImplementation(async (message) => {
+        if (message.kind === 'request') {
+          controller.abort(error)
+        }
+      })
+      const cancel = vi.fn()
+
+      await expect(peer.request(makeRequest({ body: new ReadableStream({ cancel }), signal: controller.signal }))).rejects.toBe(error)
+      await vi.waitFor(() => expect(cancel).toHaveBeenCalledOnce())
+      expect(cancel).toHaveBeenCalledWith(error)
+      expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['request', 'cancel'])
+    })
+
     it('rejects pending request on server abort', async () => {
       const { id, promise } = await requestAndGetId()
       await peer.message(makeCancelMessage(id))
@@ -474,6 +541,30 @@ describe('clientPeer', () => {
         expect(send).toHaveBeenNthCalledWith(1, expect.objectContaining({ kind: 'request' }))
         expect(send).not.toHaveBeenCalledWith(expect.objectContaining({ kind: 'cancel' }))
       })
+
+      it('does not send cancel message when transport fails after server already canceled the upload', async () => {
+        const transportError = new Error('transport failed')
+        send.mockImplementation(async (message) => {
+          if (message.kind === 'event-stream') {
+            // server stops consuming the upload right as the transport breaks
+            await peer.message(makeStreamCancelMessage(message.id))
+            throw transportError
+          }
+        })
+
+        const { id, promise } = await requestAndGetId(
+          makeRequest({ method: 'POST', headers: {}, body: makeAsyncIter(['event1']) }),
+        )
+
+        await vi.waitFor(() => expect(send.mock.calls.some(([m]) => m.kind === 'event-stream')).toBe(true))
+        await sleep(1)
+
+        await peer.message(makeResponseMessage(id, 'ok'))
+        const response = await promise
+        expect(await response.resolveBody()).toBe('ok')
+
+        expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['request', 'event-stream'])
+      })
     })
 
     describe('response body (incoming)', () => {

packages/peer/src/client.ts

@@ -1,6 +1,7 @@
-import type { StandardLazyResponse, StandardRequest } from '@standard-server/core'
+import type { StandardBody, StandardLazyResponse, StandardRequest } from '@standard-server/core'
 import type { Queue } from '@standard-server/shared'
 import type { ClientPeerSendMessage, PeerEventStreamMessage, PeerOctetStreamMessage, ServerPeerSendMessage } from './types'
+import { cancelStandardBody } from '@standard-server/core'
 import { AbortError, hasAnyDefinedValue, isAsyncIteratorObject, SequentialIdGenerator } from '@standard-server/shared'
 import { encodeAtomicStandardBody, toStandardBody } from './body'
 import { EventStreamTransmitter } from './event-stream'
@@ -61,6 +62,9 @@ export class ClientPeer {
     state: ClientPeerRequestStateInternal,
     request: StandardRequest,
   ): Promise<void> {
+    let untransmittedBody: StandardBody | undefined = request.body
+    let failure: unknown
+
     try {
       const encodedAtomicBody = await encodeAtomicStandardBody(request.body, request.headers)
 
@@ -85,42 +89,41 @@ export class ClientPeer {
         binary: encodedAtomicBody.binary,
       })
 
+      // The request can already be settled/cancelled while was in flight
+      if (this.requests.get(id) !== state) {
+        return
+      }
+
+      untransmittedBody = undefined
+
       if (isAsyncIteratorObject(request.body)) {
         const transmitter = new EventStreamTransmitter(request.body, id, this.send)
-
-        // The request can already be settled/cancelled while was in flight
-        if (this.requests.get(id) !== state) {
-          await transmitter.cancel()
-        }
-        else {
-          state.eventStreamTransmitter = transmitter
-          await transmitter.transmit().catch((error) => {
-            if (state.eventStreamTransmitter) {
-              return this.abortById(id, error)
-            }
-          })
-        }
+        state.eventStreamTransmitter = transmitter
+        await transmitter.transmit().catch((error) => {
+          if (state.eventStreamTransmitter) {
+            return this.abortById(id, error)
+          }
+        })
       }
       else if (request.body instanceof ReadableStream) {
         const transmitter = new OctetStreamTransmitter(request.body, id, this.send)
-
-        // The request can already be settled/cancelled while was in flight
-        if (this.requests.get(id) !== state) {
-          await transmitter.cancel()
-        }
-        else {
-          state.octetStreamTransmitter = transmitter
-          await transmitter.transmit().catch((error) => {
-            if (state.octetStreamTransmitter) {
-              return this.abortById(id, error)
-            }
-          })
-        }
+        state.octetStreamTransmitter = transmitter
+        await transmitter.transmit().catch((error) => {
+          if (state.octetStreamTransmitter) {
+            return this.abortById(id, error)
+          }
+        })
       }
     }
     catch (reason) {
+      failure = reason
       await this.closeById(id, reason)
     }
+    finally {
+      if (untransmittedBody !== undefined) {
+        await cancelStandardBody(untransmittedBody, failure ?? request.signal?.reason).catch(() => {})
+      }
+    }
   }
 
   /**

packages/peer/src/server.test.ts

@@ -195,13 +195,68 @@ describe('serverPeer', () => {
         await peer.message(makeCancelMessage('1'), vi.fn())
       })
 
+      const cancel = vi.fn()
       const handler = vi.fn<HandlerFn>().mockImplementation(async () => {
-        return octetStreamResponse(new ReadableStream({ }))
+        return octetStreamResponse(new ReadableStream({ cancel }))
       })
 
       await peer.message(makeRequestMessage(), handler)
       expect(send).toHaveBeenCalledTimes(1)
       expect(send).toHaveBeenCalledWith(expect.objectContaining({ kind: 'response' }))
+      expect(cancel).toHaveBeenCalledOnce()
+    })
+
+    it('cancels an octet-stream response body returned after abort', async () => {
+      const cancel = vi.fn()
+      const handler = vi.fn<HandlerFn>().mockImplementation(async () => {
+        await peer.message(makeCancelMessage('1'), vi.fn())
+        return octetStreamResponse(new ReadableStream({ cancel }))
+      })
+
+      await peer.message(makeRequestMessage(), handler)
+      expect(send).toHaveBeenCalledTimes(0)
+      expect(cancel).toHaveBeenCalledOnce()
+      expect(cancel).toHaveBeenCalledWith(expect.any(AbortError))
+    })
+
+    it('returns an event-stream response body returned after abort', async () => {
+      const cleanup = vi.fn()
+      const handler = vi.fn<HandlerFn>().mockImplementation(async () => {
+        await peer.close()
+        return eventStreamResponse(new AsyncIteratorClass(async () => ({ done: true, value: undefined }), cleanup))
+      })
+
+      await peer.message(makeRequestMessage(), handler)
+      expect(send).toHaveBeenCalledTimes(0)
+      expect(cleanup).toHaveBeenCalledOnce()
+      expect(cleanup).toHaveBeenCalledWith({ kind: 'cancelled' })
+    })
+
+    it('rejects when a response body returned after abort fails to clean up', async () => {
+      const error = new Error('cleanup failed')
+      const handler = vi.fn<HandlerFn>().mockImplementation(async () => {
+        await peer.message(makeCancelMessage('1'), vi.fn())
+        return eventStreamResponse(new AsyncIteratorClass(async () => ({ done: true, value: undefined }), async () => {
+          throw error
+        }))
+      })
+
+      await expect(peer.message(makeRequestMessage(), handler)).rejects.toBe(error)
+      expect(send).toHaveBeenCalledTimes(0)
+    })
+
+    it('cancels the response body when sending the response message fails', async () => {
+      const error = new Error('send failed')
+      send.mockRejectedValueOnce(error)
+
+      const cancel = vi.fn()
+      await expect(
+        peer.message(makeRequestMessage(), async () => octetStreamResponse(new ReadableStream({ cancel }))),
+      ).rejects.toThrow(error)
+
+      expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['response', 'cancel'])
+      expect(cancel).toHaveBeenCalledOnce()
+      expect(cancel).toHaveBeenCalledWith(error)
     })
   })