starpc 0.52.0 → 0.52.1

package/srpc/stream.test.ts

@@ -1,5 +1,9 @@
-import { describe, expect, it } from 'vitest'
+import { describe, expect, it, vi } from 'vitest'
 import { pipe } from 'it-pipe'
+import { pushable } from 'it-pushable'
+import type { Source } from 'it-stream-types'
+import type { Stream } from './stream-muxer.js'
+import { streamToPacketStream } from './stream.js'
 
 import {
   ChannelStream,
@@ -71,6 +75,133 @@ describe('StreamConn packet stream', () => {
     expect(serverError).toBeUndefined()
   })
 
+  it('settles a blocked packet source when closed', async () => {
+    const { clientConn, cleanup } = connectStreamConns({
+      handlePacketStream() {},
+    })
+
+    try {
+      const stream = await clientConn.openStream()
+      const pending = stream.source.next()
+
+      await stream.close()
+
+      await expect(pending).resolves.toEqual({ done: true, value: undefined })
+    } finally {
+      cleanup()
+    }
+  })
+
+  it('settles a blocked packet source when aborted', async () => {
+    const { clientConn, cleanup } = connectStreamConns({
+      handlePacketStream() {},
+    })
+
+    try {
+      const stream = await clientConn.openStream()
+      const pending = stream.source.next()
+      const error = new Error('stopped')
+
+      stream.abort(error)
+
+      await expect(pending).rejects.toBe(error)
+    } finally {
+      cleanup()
+    }
+  })
+
+  it('settles a blocked packet sink when closed', async () => {
+    const { clientConn, cleanup } = connectStreamConns({
+      handlePacketStream() {},
+    })
+
+    try {
+      const stream = await clientConn.openStream()
+      const input = pushable<Uint8Array>({ objectMode: true })
+      const pending = stream.sink(input)
+
+      await stream.close()
+
+      await expect(pending).resolves.toBeUndefined()
+    } finally {
+      cleanup()
+    }
+  })
+
+  it('settles a packet sink blocked in the underlying write', async () => {
+    const sinkStarted = Promise.withResolvers<void>()
+    const transport = {
+      source: (async function* () {})(),
+      sink: async (source: Source<Uint8Array>) => {
+        for await (const _chunk of source) {
+          sinkStarted.resolve()
+          await new Promise<void>(() => {})
+        }
+      },
+      close: vi.fn(async () => {}),
+      closeRead: vi.fn(async () => {}),
+      closeWrite: vi.fn(async () => {}),
+      abort: vi.fn(),
+    } satisfies Stream
+    const stream = streamToPacketStream(transport)
+    const input = pushable<Uint8Array>({ objectMode: true })
+    input.push(new Uint8Array([1]))
+    const pending = stream.sink(input)
+    await sinkStarted.promise
+
+    await stream.close()
+
+    await expect(pending).resolves.toBeUndefined()
+  })
+
+  it('rejects a blocked packet sink with the abort error', async () => {
+    const { clientConn, cleanup } = connectStreamConns({
+      handlePacketStream() {},
+    })
+
+    try {
+      const stream = await clientConn.openStream()
+      const input = pushable<Uint8Array>({ objectMode: true })
+      const pending = stream.sink(input)
+      const error = new Error('stopped')
+
+      stream.abort(error)
+
+      await expect(pending).rejects.toBe(error)
+    } finally {
+      cleanup()
+    }
+  })
+
+  it('does not write ready input after close', async () => {
+    const serverStream = Promise.withResolvers<PacketStream>()
+    const { clientConn, cleanup } = connectStreamConns({
+      handlePacketStream(stream) {
+        serverStream.resolve(stream)
+      },
+    })
+
+    try {
+      const stream = await clientConn.openStream()
+      const peer = await serverStream.promise
+      const input = pushable<Uint8Array>({ objectMode: true })
+      input.push(new Uint8Array([1]))
+
+      const pending = stream.sink(input)
+      await stream.close()
+
+      await expect(pending).resolves.toBeUndefined()
+      await expect(nextWithTimeout(peer.source, 'server eof')).resolves.toEqual(
+        {
+          done: true,
+          value: undefined,
+        },
+      )
+    } finally {
+      cleanup()
+    }
+  })
+
   it('aborts the yamux stream when the packet source errors', async () => {
     const request = new TextEncoder().encode('request')
     const sourceError = new Error('source failed')

package/srpc/stream.ts

@@ -8,16 +8,26 @@ import {
   parseLengthPrefixTransform,
   prependLengthPrefixTransform,
 } from './packet.js'
+import {
+  closeIterator,
+  sourceIterator,
+  TerminationGate,
+} from './termination.js'
 
 // PacketHandler handles incoming packets.
 export type PacketHandler = (packet: Packet) => Promise<void>
 
 // PacketStream represents a stream of packets where each Uint8Array represents one packet.
-export type PacketStream = Duplex<
+export interface PacketStream extends Duplex<
   AsyncGenerator<Uint8Array>,
   Source<Uint8Array>,
   Promise<void>
->
+> {
+  // close cleanly ends both directions of the stream.
+  close(): Promise<void>
+  // abort ends both directions of the stream with err.
+  abort(err: Error): void
+}
 
 // OpenStreamFunc is a function to start a new RPC by opening a Stream.
 export type OpenStreamFunc = () => Promise<PacketStream>
@@ -29,20 +39,66 @@ export type HandleStreamFunc = (ch: PacketStream) => Promise<void>
 
 // streamToPacketStream converts a Stream into a PacketStream using length-prefix framing.
 export function streamToPacketStream(stream: Stream): PacketStream {
+  const termination = new TerminationGate()
   return {
-    source: pipe(
-      stream,
-      parseLengthPrefixTransform(),
-      combineUint8ArrayListTransform(),
-    ),
+    close: async () => {
+      if (termination.terminate()) await stream.close()
+    },
+    abort: (err: Error) => {
+      if (termination.terminate(err)) stream.abort(err)
+    },
+    source: (async function* () {
+      const packets = pipe(
+        stream,
+        parseLengthPrefixTransform(),
+        combineUint8ArrayListTransform(),
+      )[Symbol.asyncIterator]()
+      try {
+        while (true) {
+          const next = await termination.next(packets)
+          if ('terminated' in next) {
+            if (next.error) throw next.error
+            return
+          }
+          if ('error' in next) throw next.error
+          if (next.result.done) return
+          yield next.result.value
+        }
+      } finally {
+        closeIterator(packets)
+      }
+    })(),
     sink: async (source: Source<Uint8Array>): Promise<void> => {
+      const iterator = sourceIterator(
+        pipe(source, prependLengthPrefixTransform()),
+      )
+      const gatedSource = (async function* () {
+        while (true) {
+          const next = await termination.next(iterator)
+          if ('terminated' in next) {
+            if (next.error) throw next.error
+            return
+          }
+          if ('error' in next) throw next.error
+          if (next.result.done) return
+          if (termination.terminated) return
+          yield next.result.value
+        }
+      })()
       try {
-        await pipe(source, prependLengthPrefixTransform(), stream)
+        const result = await termination.wait(stream.sink(gatedSource))
+        if ('terminated' in result) {
+          if (result.error) throw result.error
+          return
+        }
+        if ('error' in result) throw result.error
         await stream.closeWrite()
       } catch (err) {
         const error = err instanceof Error ? err : new Error(String(err))
-        stream.abort(error)
+        if (termination.terminate(error)) stream.abort(error)
         throw error
+      } finally {
+        closeIterator(iterator)
       }
     },
   }

package/srpc/termination.test.ts

@@ -0,0 +1,30 @@
+import { describe, expect, it } from 'vitest'
+
+import { TerminationGate } from './termination.js'
+
+describe('TerminationGate', () => {
+  it('prioritizes termination over an already-ready result', async () => {
+    const gate = new TerminationGate()
+    const iterator = (async function* () {
+      yield 1
+    })()
+    const pending = gate.next(iterator)
+    const error = new Error('stopped')
+
+    gate.terminate(error)
+
+    await expect(pending).resolves.toEqual({ terminated: true, error })
+  })
+
+  it('keeps the first termination result', async () => {
+    const gate = new TerminationGate()
+    const error = new Error('stopped')
+
+    expect(gate.terminate(error)).toBe(true)
+    expect(gate.terminate()).toBe(false)
+    await expect(gate.wait(new Promise<void>(() => {}))).resolves.toEqual({
+      terminated: true,
+      error,
+    })
+  })
+})

package/srpc/termination.ts

@@ -0,0 +1,70 @@
+import type { Source } from 'it-stream-types'
+
+type Terminated = { terminated: true; error?: Error }
+type Received<T> = { result: IteratorResult<T> } | { error: unknown }
+type WaitResult<T> = { result: T } | { error: unknown }
+
+// TerminationGate makes the first clean close or abort visible to blocked I/O.
+export class TerminationGate {
+  private readonly _waiters = new Set<(error?: Error) => void>()
+  private _error: Error | undefined
+  private _terminated = false
+
+  public get terminated(): boolean {
+    return this._terminated
+  }
+
+  public terminate(error?: Error): boolean {
+    if (this._terminated) return false
+    this._terminated = true
+    this._error = error
+    for (const waiter of this._waiters) waiter(error)
+    this._waiters.clear()
+    return true
+  }
+
+  public async wait<T>(
+    promise: Promise<T>,
+  ): Promise<WaitResult<T> | Terminated> {
+    if (this._terminated) return { terminated: true, error: this._error }
+
+    let notify!: (error?: Error) => void
+    const terminated = new Promise<Terminated>((resolve) => {
+      notify = (error) => resolve({ terminated: true, error })
+      this._waiters.add(notify)
+    })
+    const received: Promise<WaitResult<T>> = promise.then(
+      (result) => ({ result }),
+      (error: unknown) => ({ error }),
+    )
+
+    try {
+      const result = await Promise.race([received, terminated])
+      if (this._terminated) {
+        return { terminated: true, error: this._error }
+      }
+      return result
+    } finally {
+      this._waiters.delete(notify)
+    }
+  }
+
+  public next<T>(
+    iterator: AsyncIterator<T>,
+  ): Promise<Received<T> | Terminated> {
+    return this.wait(iterator.next())
+  }
+}
+
+export function sourceIterator<T>(source: Source<T>): AsyncIterator<T> {
+  if (Symbol.asyncIterator in source) return source[Symbol.asyncIterator]()
+  const iterator = source[Symbol.iterator]()
+  return { next: async () => iterator.next() }
+}
+// closeIterator asks an abandoned iterator to release its upstream resources.
+// Async generators may wait for an active next call before running return;
+// rejection is observed here so cleanup cannot create an unhandled promise.
+export function closeIterator<T>(iterator: AsyncIterator<T>): void {
+  if (!iterator.return) return
+  void Promise.resolve(iterator.return()).catch(() => undefined)
+}