fix(peer): release client request state once a streamed response finishes

packages/peer/src/client.test.ts

@@ -1,4 +1,5 @@
 import type { StandardLazyResponse, StandardRequest } from '@standard-server/core'
+import { ErrorEvent } from '@standard-server/core'
 import {
   AbortError,
   AsyncIteratorClass,
@@ -819,6 +820,65 @@ describe('clientPeer', () => {
           new TypeError('Failed to read body: body stream already read'),
         )
       })
+
+      it('releases the request once the stream closes, even if the body is never resolved', async () => {
+        const controller = new AbortController()
+        const removeSpy = vi.spyOn(controller.signal, 'removeEventListener')
+
+        const { id, promise } = await requestAndGetId(makeRequest({ signal: controller.signal }))
+        await peer.message(makeStreamingResponse(id, 'event-stream'))
+        const response = await promise
+
+        await peer.message(makeEventStreamMessage(id, 'evt1'))
+        expect(getPeerSize(peer)).toBe(1)
+        expect(removeSpy).not.toHaveBeenCalled()
+
+        await peer.message(makeEventStreamMessage(id, 'bye', 'close'))
+        expect(getPeerSize(peer)).toBe(0)
+        expect(removeSpy).toHaveBeenCalledOnce()
+
+        const iter = (await response.resolveBody()) as AsyncIterator<unknown>
+        await expect(iter.next()).resolves.toEqual({ done: false, value: 'evt1' })
+        await expect(iter.next()).resolves.toEqual({ done: true, value: 'bye' })
+        await expect(iter.next()).resolves.toEqual({ done: true, value: undefined })
+        expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['request'])
+      })
+
+      it('releases the request once the stream errors, even if the body is never resolved', async () => {
+        const { id, promise } = await requestAndGetId()
+        await peer.message(makeStreamingResponse(id, 'event-stream'))
+        const response = await promise
+
+        await peer.message(makeEventStreamMessage(id, 'evt1'))
+        await peer.message(makeEventStreamMessage(id, { code: 'BOOM' }, 'error'))
+        expect(getPeerSize(peer)).toBe(0)
+
+        const iter = (await response.resolveBody()) as AsyncIterator<unknown>
+        await expect(iter.next()).resolves.toEqual({ done: false, value: 'evt1' })
+        await expect(iter.next()).rejects.toSatisfy((error: ErrorEvent) => {
+          expect(error).toBeInstanceOf(ErrorEvent)
+          expect(error.data).toEqual({ code: 'BOOM' })
+          return true
+        })
+        await expect(iter.next()).resolves.toEqual({ done: true, value: undefined })
+        expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['request'])
+      })
+
+      it('does not send cancel message when canceling the iterator of a finished stream', async () => {
+        const { id, promise } = await requestAndGetId()
+        await peer.message(makeStreamingResponse(id, 'event-stream'))
+        const response = await promise
+
+        await peer.message(makeEventStreamMessage(id, 'evt1'))
+        await peer.message(makeEventStreamMessage(id, 'evt2'))
+        await peer.message(makeEventStreamMessage(id, 'bye', 'close'))
+
+        const iter = (await response.resolveBody()) as AsyncIterator<unknown>
+        await expect(iter.next()).resolves.toEqual({ done: false, value: 'evt1' })
+        await iter.return?.()
+
+        expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['request'])
+      })
     })
   })
 
@@ -1177,6 +1237,45 @@ describe('clientPeer', () => {
           new TypeError('Failed to read body: body stream already read'),
         )
       })
+
+      it('releases the request once the stream closes, even if the body is never resolved', async () => {
+        const controller = new AbortController()
+        const removeSpy = vi.spyOn(controller.signal, 'removeEventListener')
+
+        const { id, promise } = await requestAndGetId(makeRequest({ signal: controller.signal }))
+        await peer.message(makeStreamingResponse(id, 'octet-stream'))
+        const response = await promise
+
+        await peer.message(makeOctetStreamMessage(id, false, new Uint8Array([1, 2])))
+        expect(getPeerSize(peer)).toBe(1)
+        expect(removeSpy).not.toHaveBeenCalled()
+
+        await peer.message(makeOctetStreamMessage(id, true, new Uint8Array([3])))
+        expect(getPeerSize(peer)).toBe(0)
+        expect(removeSpy).toHaveBeenCalledOnce()
+
+        const reader = ((await response.resolveBody()) as ReadableStream<Uint8Array>).getReader()
+        expect(await reader.read()).toEqual({ done: false, value: new Uint8Array([1, 2]) })
+        expect(await reader.read()).toEqual({ done: false, value: new Uint8Array([3]) })
+        expect(await reader.read()).toEqual({ done: true, value: undefined })
+        expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['request'])
+      })
+
+      it('does not send cancel message when canceling the ReadableStream of a finished stream', async () => {
+        const { id, promise } = await requestAndGetId()
+        await peer.message(makeStreamingResponse(id, 'octet-stream'))
+        const response = await promise
+
+        await peer.message(makeOctetStreamMessage(id, false, new Uint8Array([1])))
+        await peer.message(makeOctetStreamMessage(id, false, new Uint8Array([2])))
+        await peer.message(makeOctetStreamMessage(id, true))
+
+        const reader = ((await response.resolveBody()) as ReadableStream<Uint8Array>).getReader()
+        expect(await reader.read()).toEqual({ done: false, value: new Uint8Array([1]) })
+        await reader.cancel()
+
+        expect(send.mock.calls.map(([m]) => m.kind)).toEqual(['request'])
+      })
     })
   })
 

packages/peer/src/client.ts

@@ -178,13 +178,26 @@ export class ClientPeer {
       return
     }
 
+    // by the last stream message the server has already released the request, so release it too
     if (message.kind === 'event-stream') {
-      state.eventStreamMessageQueue?.push(message)
+      const queue = state.eventStreamMessageQueue
+      queue?.push(message)
+
+      if (queue && (message.json.event === 'close' || message.json.event === 'error')) {
+        await this.closeById(id)
+      }
+
       return
     }
 
     if (message.kind === 'octet-stream') {
-      state.octetStreamMessageQueue?.push(message)
+      const queue = state.octetStreamMessageQueue
+      queue?.push(message)
+
+      if (queue && message.json.close) {
+        await this.closeById(id)
+      }
+
       return
     }
 

packages/peer/tests/peer.test.ts

@@ -1,6 +1,8 @@
+import { getEventListeners } from 'node:events'
+
 import type { StandardLazyRequest, StandardResponse } from '@standard-server/core'
-import { ErrorEvent, getEventMeta, withEventMeta } from '@standard-server/core'
-import { promiseWithResolvers } from '@standard-server/shared'
+import { cancelStandardBody, ErrorEvent, getEventMeta, withEventMeta } from '@standard-server/core'
+import { promiseWithResolvers, sleep } from '@standard-server/shared'
 
 import type { ClientPeer, ServerPeer } from '../src'
 import {
@@ -357,6 +359,78 @@ describe('peer integration (client <-> server over encoded wire)', () => {
     await vi.waitFor(() => expect(cancel).toHaveBeenCalled())
   })
 
+  it('releases both peers when finished streamed response bodies are never resolved', async () => {
+    const { client, server } = connect(async (request) => ({
+      status: 200,
+      headers: {},
+      body:
+        request.url === '/events'
+          ? (async function* () {
+              yield 1
+              yield 2
+            })()
+          : new ReadableStream<Uint8Array>({
+              start(controller) {
+                controller.enqueue(new Uint8Array([1, 2]))
+                controller.close()
+              },
+            }),
+    }))
+
+    const controller = new AbortController()
+
+    for (let i = 0; i < 100; i++) {
+      const response = await client.request({
+        url: i % 2 === 0 ? '/events' : '/download',
+        method: 'GET',
+        headers: {},
+        signal: controller.signal,
+      })
+      expect(response.status).toBe(200)
+    }
+
+    await vi.waitFor(() => {
+      expect((client as any).requests.size).toBe(0)
+      expect((server as any).requests.size).toBe(0)
+    })
+    expect(getEventListeners(controller.signal, 'abort')).toHaveLength(0)
+  })
+
+  it('releases both peers when an endless streamed response body is cancelled without reading', async () => {
+    const finalized = vi.fn()
+
+    const { client, server } = connect(async () => ({
+      status: 200,
+      headers: {},
+      body: (async function* () {
+        try {
+          while (true) {
+            yield 'tick'
+            await sleep(1)
+          }
+        } finally {
+          finalized()
+        }
+      })(),
+    }))
+
+    const controller = new AbortController()
+    const response = await client.request({
+      url: '/events',
+      method: 'GET',
+      headers: {},
+      signal: controller.signal,
+    })
+    await cancelStandardBody(await response.resolveBody())
+
+    await vi.waitFor(() => {
+      expect((client as any).requests.size).toBe(0)
+      expect((server as any).requests.size).toBe(0)
+      expect(finalized).toHaveBeenCalledOnce()
+    })
+    expect(getEventListeners(controller.signal, 'abort')).toHaveLength(0)
+  })
+
   it('propagates client aborts to the server handler signal', async () => {
     let serverSignal: AbortSignal | undefined