starpc 0.52.0 → 0.52.2

package/srpc/stream.hpp

@@ -1,61 +1,122 @@
 #pragma once
 
-#include <functional>
-#include <memory>
-
 #include "errors.hpp"
 #include "message.hpp"
 
+#include <atomic>
+#include <functional>
+#include <memory>
+#include <stop_token>
+#include <string>
+
 namespace starpc {
 
-// Stream is a handle to an on-going bi-directional or one-directional stream
-// RPC handle. Matches Go Stream interface in stream.go
+/*
+ * Stream is one live RPC stream between the caller and the remote. Matches
+ * the Go Stream interface in stream.go.
+ */
 class Stream {
 public:
-  virtual ~Stream() = default;
+	virtual ~Stream() = default;
+
+	/* StopToken interrupts work when the RPC ends or is canceled. */
+	virtual std::stop_token StopToken() const
+	{
+		return {};
+	}
 
-  // MsgSend sends the message to the remote.
-  virtual Error MsgSend(const Message &msg) = 0;
+	/*
+	 * RemoteErrorMessage returns the diagnostic accompanying
+	 * Error::RemoteError.
+	 */
+	virtual std::string RemoteErrorMessage() const
+	{
+		return {};
+	}
 
-  // MsgRecv receives an incoming message from the remote.
-  // Parses the message into the object at msg.
-  virtual Error MsgRecv(Message *msg) = 0;
+	/* MsgSend serializes msg and writes it to the remote. */
+	virtual Error MsgSend(const Message &msg) = 0;
 
-  // CloseSend signals to the remote that we will no longer send any messages.
-  virtual Error CloseSend() = 0;
+	/*
+	 * MsgRecv reads the next message from the remote and parses it into
+	 * msg. Returns EOF_ after the call completes.
+	 */
+	virtual Error MsgRecv(Message *msg) = 0;
 
-  // Close closes the stream for reading and writing.
-  virtual Error Close() = 0;
+	/*
+	 * CloseSend half-closes: no further messages will be sent. Once the
+	 * request is written, the half-close is only a notice: it fails when
+	 * the call already ended, and MsgRecv then returns the buffered
+	 * messages followed by the call's outcome.
+	 */
+	virtual Error CloseSend() = 0;
+
+	/* Close ends the stream for both reading and writing. */
+	virtual Error Close() = 0;
 };
 
-// StreamWithClose wraps a Stream with a close function to call when Close is
-// called. Matches streamWithClose in stream.go
+/*
+ * StreamWithClose forwards a Stream and runs an extra callback on the first
+ * Close. Matches streamWithClose in stream.go.
+ */
 class StreamWithClose : public Stream {
 public:
-  StreamWithClose(Stream *inner, std::function<Error()> close_fn)
-      : inner_(inner), close_fn_(std::move(close_fn)) {}
+	/*
+	 * StreamWithClose borrows inner; it must outlive the wrapper.
+	 * close_fn runs once, after inner->Close.
+	 */
+	StreamWithClose(Stream *inner, std::function<Error()> close_fn)
+		: inner(inner),
+		  close_fn(std::move(close_fn))
+	{
+	}
+
+	Error MsgSend(const Message &msg) override
+	{
+		return inner->MsgSend(msg);
+	}
+	Error MsgRecv(Message *msg) override
+	{
+		return inner->MsgRecv(msg);
+	}
+	Error CloseSend() override
+	{
+		return inner->CloseSend();
+	}
+	std::stop_token StopToken() const override
+	{
+		return inner->StopToken();
+	}
+	std::string RemoteErrorMessage() const override
+	{
+		return inner->RemoteErrorMessage();
+	}
 
-  Error MsgSend(const Message &msg) override { return inner_->MsgSend(msg); }
-  Error MsgRecv(Message *msg) override { return inner_->MsgRecv(msg); }
-  Error CloseSend() override { return inner_->CloseSend(); }
+	Error Close() override
+	{
+		if (closed.exchange(true))
+			return Error::OK;
 
-  Error Close() override {
-    Error err = inner_->Close();
-    Error err2 = close_fn_();
-    if (err != Error::OK)
-      return err;
-    return err2;
-  }
+		/* Report the inner close even when the callback also fails. */
+		Error err = inner->Close();
+		Error err2 = close_fn();
+		if (err != Error::OK)
+			return err;
+		return err2;
+	}
 
 private:
-  Stream *inner_;
-  std::function<Error()> close_fn_;
+	/* inner is borrowed; the caller owns it. */
+	Stream *inner;
+	std::function<Error()> close_fn;
+	std::atomic<bool> closed{false};
 };
 
-// NewStreamWithClose wraps a Stream with a close function.
+/* NewStreamWithClose wraps strm so close_fn runs on the first Close. */
 inline std::unique_ptr<Stream>
-NewStreamWithClose(Stream *strm, std::function<Error()> close_fn) {
-  return std::make_unique<StreamWithClose>(strm, std::move(close_fn));
+NewStreamWithClose(Stream *strm, std::function<Error()> close_fn)
+{
+	return std::make_unique<StreamWithClose>(strm, std::move(close_fn));
 }
 
 } // namespace starpc

package/srpc/stream.rs

@@ -110,7 +110,7 @@ impl Context {
 ///
 /// - All methods are async and may block waiting for I/O or messages
 /// - `send_bytes` may return `Error::Completed` if `close_send` was already called
-/// - `recv_bytes` returns `Error::StreamClosed` when the remote closes
+/// - `recv_bytes` returns `Error::StreamClosed` when the remote completes
 /// - `close` cancels the context and releases all resources
 #[async_trait]
 pub trait Stream: Send + Sync {
@@ -133,8 +133,9 @@ pub trait Stream: Send + Sync {
     ///
     /// # Errors
     ///
-    /// - `Error::StreamClosed` if the remote closed the stream
+    /// - `Error::StreamClosed` if the remote completed the call
     /// - `Error::Remote` if the remote sent an error
+    /// - `Error::ClosedBeforeCompletion` if the transport closed first
     /// - `Error::Cancelled` if the context was cancelled
     async fn recv_bytes(&self) -> Result<Bytes>;
 
@@ -142,6 +143,10 @@ pub trait Stream: Send + Sync {
     ///
     /// After calling this, `send_bytes` will return `Error::Completed`.
     /// The receive side remains open until the remote closes.
+    ///
+    /// Once the request is written, the half-close is only a notice. It fails
+    /// when the call already ended, and `recv_bytes` then returns the buffered
+    /// messages followed by the call's outcome.
     async fn close_send(&self) -> Result<()>;
 
     /// Closes both sides of the stream.

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/testing.rs

@@ -13,22 +13,8 @@ use crate::error::{Error, Result};
 use crate::rpc::PacketWriter;
 use crate::transport::create_packet_channel;
 
-/// Creates a pair of connected in-memory streams.
-///
-/// Returns (client_stream, server_stream) that can be used for testing
-/// client-server communication.
-///
-/// # Arguments
-/// * `buffer_size` - Size of the internal buffer for each direction
-///
-/// # Example
-///
-/// ```ignore
-/// let (client_stream, server_stream) = create_pipe(64 * 1024);
-///
-/// // Use client_stream with a client
-/// // Use server_stream with a server
-/// ```
+/// Creates a pair of connected in-memory streams for testing client-server
+/// communication, with `buffer_size` bytes of buffer in each direction.
 pub fn create_pipe(buffer_size: usize) -> (DuplexStream, DuplexStream) {
     duplex(buffer_size)
 }
@@ -38,21 +24,8 @@ pub fn create_pipe_default() -> (DuplexStream, DuplexStream) {
     create_pipe(64 * 1024)
 }
 
-/// In-memory stream opener for testing.
-///
-/// This opener can hold multiple streams and returns them one at a time
-/// on each call to `open_stream`.
-///
-/// # Example
-///
-/// ```ignore
-/// let (client_stream, server_stream) = create_pipe_default();
-///
-/// let opener = InMemoryOpener::new(vec![client_stream]);
-/// let client = SrpcClient::new(opener);
-///
-/// // The client will use client_stream for its RPC
-/// ```
+/// In-memory stream opener for testing that holds multiple streams and
+/// returns them one at a time on each call to `open_stream`.
 pub struct InMemoryOpener {
     streams: Arc<Mutex<Vec<DuplexStream>>>,
 }
@@ -108,40 +81,15 @@ impl SingleInMemoryOpener {
 #[async_trait]
 impl OpenStream for SingleInMemoryOpener {
     async fn open_stream(&self) -> Result<(Arc<dyn PacketWriter>, PacketReceiver)> {
-        let stream = self
-            .stream
-            .lock()
-            .await
-            .take()
-            .ok_or(Error::StreamClosed)?;
+        let stream = self.stream.lock().await.take().ok_or(Error::StreamClosed)?;
 
         let (read_half, write_half) = tokio::io::split(stream);
         Ok(create_packet_channel(read_half, write_half))
     }
 }
 
-/// Creates a connected client-server test setup.
-///
-/// Returns (client_opener, server_stream) where:
-/// - `client_opener` can be used to create a `SrpcClient`
-/// - `server_stream` can be passed to `Server::handle_stream`
-///
-/// # Example
-///
-/// ```ignore
-/// let (opener, server_stream) = create_test_pair();
-///
-/// let client = SrpcClient::new(opener);
-/// let server = Server::new(mux);
-///
-/// // Spawn server
-/// tokio::spawn(async move {
-///     server.handle_stream(server_stream).await
-/// });
-///
-/// // Use client
-/// let response = client.exec_call("svc", "method", &request).await?;
-/// ```
+/// Creates a connected client-server test setup: an opener for a
+/// `SrpcClient` and the server stream for `Server::handle_stream`.
 pub fn create_test_pair() -> (SingleInMemoryOpener, DuplexStream) {
     let (client_stream, server_stream) = create_pipe_default();
     (SingleInMemoryOpener::new(client_stream), server_stream)
@@ -225,13 +173,10 @@ mod tests {
         stream.close_send().await.unwrap();
 
         // Read response with a timeout to handle potential races.
-        let response = tokio::time::timeout(
-            std::time::Duration::from_secs(5),
-            stream.recv_bytes(),
-        )
-        .await
-        .expect("timeout")
-        .expect("recv_bytes failed");
+        let response = tokio::time::timeout(std::time::Duration::from_secs(5), stream.recv_bytes())
+            .await
+            .expect("timeout")
+            .expect("recv_bytes failed");
         assert_eq!(&response[..], b"hello");
 
         // Wait for server to complete.