starpc 0.51.0 → 0.52.1

package/srpc/server.test.ts

@@ -1,5 +1,7 @@
 import { describe, it, beforeEach, expect, vi } from 'vitest'
 import { pipe } from 'it-pipe'
+import { pushable } from 'it-pushable'
+import type { Source } from 'it-stream-types'
 import {
   createHandler,
   createMux,
@@ -10,6 +12,10 @@ import {
   combineUint8ArrayListTransform,
   ChannelStreamOpts,
   Packet,
+  createContextKey,
+  serverContextValue,
+  withServerContextValue,
+  type ServerContext,
 } from '../srpc/index.js'
 import {
   EchoerDefinition,
@@ -78,12 +84,21 @@ describe('srpc server', () => {
   it('should pass rpc stream tests', async () => {
     await runRpcStreamTest(client)
   })
-  it('passes the exact invocation signal after async request decode', async () => {
+  it('passes the exact invocation context after async request decode', async () => {
     const controller = new AbortController()
-    let observedSignal: AbortSignal | undefined
+    const callerKey = createContextKey<string>()
+    let observedAbortSignal: AbortSignal | undefined
+    let observedContextSignal: AbortSignal | undefined
+    let observedCaller: string | undefined
     const handler = createHandler(EchoerDefinition, {
-      Echo: async (request: EchoMsgType, signal?: AbortSignal) => {
-        observedSignal = signal
+      Echo: async (
+        request: EchoMsgType,
+        abortSignal: AbortSignal,
+        context: ServerContext,
+      ) => {
+        observedAbortSignal = abortSignal
+        observedContextSignal = context.signal
+        observedCaller = serverContextValue(context, callerKey)
         return request
       },
     })
@@ -104,9 +119,41 @@ describe('srpc server', () => {
         }
         drained.resolve()
       },
-      controller.signal,
+      withServerContextValue(
+        { signal: controller.signal },
+        callerKey,
+        'caller-1',
+      ),
     )
     await drained.promise
+    expect(observedAbortSignal).toBe(controller.signal)
+    expect(observedContextSignal).toBe(controller.signal)
+    expect(observedCaller).toBe('caller-1')
+  })
+
+  it('keeps two-argument server handlers compatible', async () => {
+    const controller = new AbortController()
+    let observedSignal: AbortSignal | undefined
+    const handler = createHandler(EchoerDefinition, {
+      Echo: async (request: EchoMsgType, abortSignal?: AbortSignal) => {
+        observedSignal = abortSignal
+        return request
+      },
+    })
+    const invokeFn = await handler.lookupMethod(EchoerServiceName, 'Echo')
+    if (!invokeFn) throw new Error('Echo method was not found')
+    const request = EchoMsg.create({ body: 'legacy handler' })
+    await invokeFn(
+      (async function* () {
+        yield EchoMsg.toBinary(request)
+      })(),
+      async (source) => {
+        for await (const _data of source) {
+          // Drain the response.
+        }
+      },
+      { signal: controller.signal },
+    )
     expect(observedSignal).toBe(controller.signal)
   })
 
@@ -130,6 +177,8 @@ describe('srpc server', () => {
     const server = new Server(mux.lookupMethod)
     const firstResponse = new Promise<Packet>((resolve, reject) => {
       server.handlePacketStream({
+        close: async () => {},
+        abort: () => {},
         source: (async function* () {
           yield Packet.toBinary({
             body: {
@@ -297,6 +346,52 @@ describe('srpc server', () => {
     await Promise.resolve()
   })
 
+  it('closes the packet stream when the server pipeline completes', async () => {
+    const server = new Server(createMux().lookupMethod)
+    const close = vi.fn(async () => {})
+    const source = pushable<Uint8Array>({ objectMode: true })
+    const stream = {
+      close,
+      abort: vi.fn(),
+      source,
+      sink: async (output: Source<Uint8Array>) => {
+        for await (const _packet of output) {
+          // Drain the response pipeline.
+        }
+      },
+    }
+    const rpc = server.handlePacketStream(stream)
+
+    await rpc.close()
+    source.end()
+
+    await vi.waitFor(() => expect(close).toHaveBeenCalledOnce())
+    expect(stream.abort).not.toHaveBeenCalled()
+  })
+
+  it('aborts the packet stream when the server pipeline fails', async () => {
+    const server = new Server(createMux().lookupMethod)
+    const error = new Error('input failed')
+    const abort = vi.fn()
+    const source = pushable<Uint8Array>({ objectMode: true })
+    const stream = {
+      close: vi.fn(async () => {}),
+      abort,
+      source,
+      sink: async (output: Source<Uint8Array>) => {
+        for await (const _packet of output) {
+          // Drain the response pipeline.
+        }
+      },
+    }
+
+    server.handlePacketStream(stream)
+    source.end(error)
+
+    await vi.waitFor(() => expect(abort).toHaveBeenCalledWith(error))
+    expect(stream.close).not.toHaveBeenCalled()
+  })
+
   it('tears down passive channel close state', async () => {
     const { port1, port2 } = new MessageChannel()
     const opts: ChannelStreamOpts = { idleTimeoutMs: 1000, keepAliveMs: 1000 }

package/srpc/server.ts

@@ -20,9 +20,7 @@ export class Server implements StreamHandler {
   public get rpcStreamHandler(): HandleStreamFunc {
     return async (stream: PacketStream) => {
       const rpc = this.startRpc()
-      return pipe(stream, decodePacketSource, rpc, encodePacketSource, stream)
-        .catch((err: Error) => rpc.close(err))
-        .then(() => rpc.close())
+      return runPacketStream(stream, rpc)
     }
   }
 
@@ -36,9 +34,27 @@ export class Server implements StreamHandler {
   // the stream has one Uint8Array per packet w/o length prefix.
   public handlePacketStream(stream: PacketStream): ServerRPC {
     const rpc = this.startRpc()
-    pipe(stream, decodePacketSource, rpc, encodePacketSource, stream)
-      .catch((err: Error) => rpc.close(err))
-      .then(() => rpc.close())
+    void runPacketStream(stream, rpc).catch(() => undefined)
     return rpc
   }
 }
+
+async function runPacketStream(
+  stream: PacketStream,
+  rpc: ServerRPC,
+): Promise<void> {
+  try {
+    await pipe(stream, decodePacketSource, rpc, encodePacketSource, stream)
+    if (rpc.isClosed instanceof Error) {
+      stream.abort(rpc.isClosed)
+      throw rpc.isClosed
+    }
+    await stream.close()
+    await rpc.close()
+  } catch (err) {
+    const error = err instanceof Error ? err : new Error(String(err))
+    stream.abort(error)
+    await rpc.close(error)
+    throw error
+  }
+}

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)
+}

package/srpc/watchdog.test.ts

@@ -7,6 +7,7 @@ describe('Watchdog', () => {
   })
 
   afterEach(() => {
+    vi.useRealTimers()
     vi.restoreAllMocks()
   })