starpc 0.49.14 → 0.49.15

package/README.md

@@ -232,7 +232,7 @@ import { WebSocketConn } from 'starpc'
 import { EchoerClient } from './echo/index.js'
 
 const ws = new WebSocket('ws://localhost:8080/api')
-const conn = new WebSocketConn(ws)
+const conn = new WebSocketConn(ws, 'outbound')
 const client = conn.buildClient()
 const echoer = new EchoerClient(client)
 

package/dist/srpc/server.test.js

@@ -94,6 +94,71 @@ describe('srpc server', () => {
         }
         expect([...body.value.data]).toEqual([...response]);
     });
+    it('keeps StreamConn server-streaming responses open after Go-style request close', async () => {
+        const mux = createMux();
+        const response = new TextEncoder().encode('streamconn init');
+        mux.registerLookupMethod(async (service, method) => {
+            if (service !== 'test.ResourceService' || method !== 'ResourceClient') {
+                return null;
+            }
+            return async (_dataSource, dataSink) => {
+                await dataSink((async function* () {
+                    await new Promise((resolve) => setTimeout(resolve, 10));
+                    yield response;
+                })());
+            };
+        });
+        const server = new Server(mux.lookupMethod);
+        const clientConn = new StreamConn();
+        const serverConn = new StreamConn(server, { direction: 'inbound' });
+        const { port1: clientPort, port2: serverPort } = new MessageChannel();
+        const clientChannelStream = new ChannelStream('client', clientPort);
+        const serverChannelStream = new ChannelStream('server', serverPort);
+        pipe(clientChannelStream, clientConn, combineUint8ArrayListTransform(), clientChannelStream)
+            .catch((err) => clientConn.close(err))
+            .then(() => clientConn.close());
+        pipe(serverChannelStream, serverConn, combineUint8ArrayListTransform(), serverChannelStream)
+            .catch((err) => serverConn.close(err))
+            .then(() => serverConn.close());
+        try {
+            const stream = await clientConn.openStream();
+            const firstResponse = (async () => {
+                for await (const packetData of stream.source) {
+                    const packet = Packet.fromBinary(packetData);
+                    if (packet.body?.case === 'callData') {
+                        return packet;
+                    }
+                }
+                throw new Error('server response stream ended before call data');
+            })();
+            await stream.sink((async function* () {
+                yield Packet.toBinary({
+                    body: {
+                        case: 'callStart',
+                        value: {
+                            rpcService: 'test.ResourceService',
+                            rpcMethod: 'ResourceClient',
+                            data: new Uint8Array(0),
+                            dataIsZero: true,
+                        },
+                    },
+                });
+            })());
+            const packet = await promiseWithTimeout(firstResponse, 'StreamConn server-streaming response');
+            const body = packet.body;
+            expect(body?.case).toBe('callData');
+            if (body?.case !== 'callData' || !body.value?.data) {
+                throw new Error('expected callData packet');
+            }
+            expect([...body.value.data]).toEqual([...response]);
+        }
+        finally {
+            clientConn.close();
+            serverConn.close();
+            clientChannelStream.close();
+            serverChannelStream.close();
+        }
+    });
     it('removes abort listeners after a request completes', async () => {
         const controller = new AbortController();
         const removeEventListener = vi.spyOn(controller.signal, 'removeEventListener');

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "starpc",
-  "version": "0.49.14",
+  "version": "0.49.15",
   "description": "Streaming protobuf RPC service protocol over any two-way channel.",
   "license": "MIT",
   "author": {

package/srpc/client-rpc.go

@@ -109,10 +109,14 @@ func (r *ClientRPC) HandleCallStart(pkt *CallStart) error {
 // Close releases any resources held by the ClientRPC.
 func (r *ClientRPC) Close() {
 	locked := r.bcast.Lock()
+	var writer PacketWriter
 	// call did not start yet if writer is nil.
 	if r.writer != nil {
 		_ = r.WriteCallCancel()
-		r.closeLocked(&locked)
+		writer = r.closeLocked(&locked)
 	}
 	locked.Unlock()
+	if writer != nil {
+		_ = writer.Close()
+	}
 }

package/srpc/client.rs

@@ -262,6 +262,10 @@ pub mod transport {
 
     /// Re-export TransportPacketWriter for direct use.
     pub use crate::transport::TransportPacketWriter;
+
+    /// Re-export YamuxStreamOpener for multiplexed transports.
+    #[cfg(feature = "yamux")]
+    pub use crate::yamux::YamuxStreamOpener;
 }
 
 #[cfg(test)]

package/srpc/common-rpc.go

@@ -26,6 +26,8 @@ type commonRPC struct {
 	bcast broadcast.Broadcast
 	// writer is the writer to write messages to
 	writer PacketWriter
+	// writerClosed is set after writer has been closed locally.
+	writerClosed bool
 	// dataQueue contains incoming data packets.
 	// note: packets may be len() == 0
 	dataQueue [][]byte
@@ -86,8 +88,11 @@ func (c *commonRPC) ReadOne() ([]byte, error) {
 		locked := c.bcast.Lock()
 		if ctxDone && !c.dataClosed {
 			// context must have been canceled locally
-			c.closeLocked(&locked)
+			writer := c.closeLocked(&locked)
 			locked.Unlock()
+			if writer != nil {
+				_ = writer.Close()
+			}
 			return nil, context.Canceled
 		}
 
@@ -144,7 +149,7 @@ func (c *commonRPC) WriteCallData(data []byte, dataIsZero, complete bool, err er
 func (c *commonRPC) HandleStreamClose(closeErr error) {
 	var writer PacketWriter
 	locked := c.bcast.Lock()
-	if c.dataClosed {
+	if c.dataClosed && c.writerClosed {
 		locked.Unlock()
 		return
 	}
@@ -153,7 +158,7 @@ func (c *commonRPC) HandleStreamClose(closeErr error) {
 	}
 	c.dataClosed = true
 	c.ctxCancel()
-	writer = c.writer
+	writer = c.closeWriterLocked()
 	locked.Broadcast()
 	locked.Unlock()
 	if writer != nil {
@@ -212,16 +217,22 @@ func (c *commonRPC) WriteCallCancel() error {
 }
 
 // closeLocked releases resources held by the RPC.
-func (c *commonRPC) closeLocked(locked *broadcast.Locked) {
-	if c.dataClosed {
-		return
-	}
+func (c *commonRPC) closeLocked(locked *broadcast.Locked) PacketWriter {
 	c.dataClosed = true
 	c.localCompleted.Store(true)
 	if c.remoteErr == nil {
 		c.remoteErr = context.Canceled
 	}
-	_ = c.writer.Close()
+	writer := c.closeWriterLocked()
 	locked.Broadcast()
 	c.ctxCancel()
+	return writer
+}
+
+func (c *commonRPC) closeWriterLocked() PacketWriter {
+	if c.writerClosed || c.writer == nil {
+		return nil
+	}
+	c.writerClosed = true
+	return c.writer
 }

package/srpc/common-rpc_test.go

@@ -69,6 +69,44 @@ func TestCommonRPCHandleStreamCloseClosesWriterOutsideBroadcastLock(t *testing.T
 	}
 }
 
+func TestClientRPCCloseAfterRemoteCompleteClosesWriter(t *testing.T) {
+	writer := &closeCountingPacketWriter{}
+	rpc := NewClientRPC(context.Background(), "service", "method")
+	if err := rpc.Start(writer, false, nil); err != nil {
+		t.Fatalf("start: %v", err)
+	}
+	if err := rpc.HandleCallData(NewCallDataPacket([]byte("ok"), false, true, nil).GetCallData()); err != nil {
+		t.Fatalf("handle call data: %v", err)
+	}
+
+	rpc.Close()
+
+	if got := writer.closed.Load(); got != 1 {
+		t.Fatalf("expected writer closed once, got %d", got)
+	}
+}
+
+func TestClientRPCCloseClosesWriterOutsideBroadcastLock(t *testing.T) {
+	var rpc *ClientRPC
+	writerClosedOutsideLock := false
+	writer := &closeCallbackPacketWriter{
+		closeFn: func() {
+			ok := rpc.bcast.TryHoldLock(func(func(), func() <-chan struct{}) {})
+			writerClosedOutsideLock = ok
+		},
+	}
+	rpc = NewClientRPC(context.Background(), "service", "method")
+	if err := rpc.Start(writer, false, nil); err != nil {
+		t.Fatalf("start: %v", err)
+	}
+
+	rpc.Close()
+
+	if !writerClosedOutsideLock {
+		t.Fatal("expected writer close outside broadcast lock")
+	}
+}
+
 func TestCommonRPCReadOneQueuedDoesNotAllocate(t *testing.T) {
 	msg := []byte("message")
 	queue := make([][]byte, 1)

package/srpc/lib.rs

@@ -75,6 +75,10 @@ pub mod transport;
 
 #[cfg(feature = "build")]
 pub mod build;
+#[cfg(feature = "websocket")]
+pub mod websocket;
+#[cfg(feature = "yamux")]
+pub mod yamux;
 
 // Re-exports for convenience.
 pub use client::{BoxClient, Client, OpenStream, SrpcClient};
@@ -91,6 +95,10 @@ pub use stream::{ArcStream, BoxStream, Context, Stream, StreamExt};
 pub use transport::{
     create_packet_channel, decode_optional_data, encode_optional_data, TransportPacketWriter,
 };
+#[cfg(feature = "websocket")]
+pub use websocket::websocket_byte_stream;
+#[cfg(feature = "yamux")]
+pub use yamux::YamuxStreamOpener;
 
 // Re-export async_trait for use in generated code.
 pub use async_trait::async_trait;

package/srpc/server.rs

@@ -104,7 +104,7 @@ impl<I: Invoker + 'static> Server<I> {
     }
 
     /// Reports an error through the error handler, if configured.
-    fn report_error(&self, err: Error) {
+    pub(crate) fn report_error(&self, err: Error) {
         if let Some(ref handler) = self.error_handler {
             handler(err);
         }
@@ -214,6 +214,44 @@ impl<I: Invoker + 'static> Server<I> {
         Ok(())
     }
 
+    /// handle_yamux handles a yamux connection by routing each accepted substream as a
+    /// Starpc stream.
+    #[cfg(feature = "yamux")]
+    pub async fn handle_yamux<T>(&self, transport: T) -> Result<()>
+    where
+        T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        self.handle_yamux_with_config(transport, yamux::Config::default())
+            .await
+    }
+
+    /// handle_yamux_with_config handles a yamux connection with the provided yamux
+    /// configuration.
+    #[cfg(feature = "yamux")]
+    pub async fn handle_yamux_with_config<T>(
+        &self,
+        transport: T,
+        config: yamux::Config,
+    ) -> Result<()>
+    where
+        T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        crate::yamux::handle_server_connection(self, transport, config).await
+    }
+
+    /// handle_websocket_yamux handles a WebSocket connection carrying yamux substreams.
+    #[cfg(all(feature = "websocket", feature = "yamux"))]
+    pub async fn handle_websocket_yamux<T>(
+        &self,
+        socket: tokio_tungstenite::WebSocketStream<T>,
+    ) -> Result<()>
+    where
+        T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        self.handle_yamux(crate::websocket::websocket_byte_stream(socket))
+            .await
+    }
+
     /// Accepts and handles connections in a loop.
     ///
     /// This is a convenience method that accepts connections from a listener
@@ -246,7 +284,7 @@ impl<I: Invoker + 'static> Server<I> {
     }
 
     /// Creates a clone of the server for spawning tasks.
-    fn clone_for_spawn(&self) -> Server<I> {
+    pub(crate) fn clone_for_spawn(&self) -> Server<I> {
         Server {
             invoker: self.invoker.clone(),
             config: self.config.clone(),

package/srpc/server.test.ts

@@ -140,6 +140,94 @@ describe('srpc server', () => {
     expect([...body.value.data]).toEqual([...response])
   })
 
+  it('keeps StreamConn server-streaming responses open after Go-style request close', async () => {
+    const mux = createMux()
+    const response = new TextEncoder().encode('streamconn init')
+    mux.registerLookupMethod(async (service, method) => {
+      if (service !== 'test.ResourceService' || method !== 'ResourceClient') {
+        return null
+      }
+      return async (_dataSource, dataSink) => {
+        await dataSink(
+          (async function* () {
+            await new Promise((resolve) => setTimeout(resolve, 10))
+            yield response
+          })(),
+        )
+      }
+    })
+
+    const server = new Server(mux.lookupMethod)
+    const clientConn = new StreamConn()
+    const serverConn = new StreamConn(server, { direction: 'inbound' })
+    const { port1: clientPort, port2: serverPort } = new MessageChannel()
+    const clientChannelStream = new ChannelStream('client', clientPort)
+    const serverChannelStream = new ChannelStream('server', serverPort)
+
+    pipe(
+      clientChannelStream,
+      clientConn,
+      combineUint8ArrayListTransform(),
+      clientChannelStream,
+    )
+      .catch((err: Error) => clientConn.close(err))
+      .then(() => clientConn.close())
+
+    pipe(
+      serverChannelStream,
+      serverConn,
+      combineUint8ArrayListTransform(),
+      serverChannelStream,
+    )
+      .catch((err: Error) => serverConn.close(err))
+      .then(() => serverConn.close())
+
+    try {
+      const stream = await clientConn.openStream()
+      const firstResponse = (async () => {
+        for await (const packetData of stream.source) {
+          const packet = Packet.fromBinary(packetData)
+          if (packet.body?.case === 'callData') {
+            return packet
+          }
+        }
+        throw new Error('server response stream ended before call data')
+      })()
+
+      await stream.sink(
+        (async function* () {
+          yield Packet.toBinary({
+            body: {
+              case: 'callStart',
+              value: {
+                rpcService: 'test.ResourceService',
+                rpcMethod: 'ResourceClient',
+                data: new Uint8Array(0),
+                dataIsZero: true,
+              },
+            },
+          })
+        })(),
+      )
+
+      const packet = await promiseWithTimeout(
+        firstResponse,
+        'StreamConn server-streaming response',
+      )
+      const body = packet.body
+      expect(body?.case).toBe('callData')
+      if (body?.case !== 'callData' || !body.value?.data) {
+        throw new Error('expected callData packet')
+      }
+      expect([...body.value.data]).toEqual([...response])
+    } finally {
+      clientConn.close()
+      serverConn.close()
+      clientChannelStream.close()
+      serverChannelStream.close()
+    }
+  })
+
   it('removes abort listeners after a request completes', async () => {
     const controller = new AbortController()
     const removeEventListener = vi.spyOn(

package/srpc/websocket.rs

@@ -0,0 +1,70 @@
+//! WebSocket transport adapters for starpc.
+
+use bytes::Bytes;
+use futures::{SinkExt, StreamExt};
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, DuplexStream};
+use tokio_tungstenite::{tungstenite::Message, WebSocketStream};
+
+/// WEBSOCKET_BYTE_STREAM_BUFFER is the duplex buffer used by WebSocket byte adapters.
+pub const WEBSOCKET_BYTE_STREAM_BUFFER: usize = 64 * 1024;
+
+/// WEBSOCKET_MESSAGE_BUFFER is the maximum byte chunk sent in one WebSocket message.
+pub const WEBSOCKET_MESSAGE_BUFFER: usize = 16 * 1024;
+
+/// websocket_byte_stream adapts a binary WebSocket stream to Tokio byte I/O.
+///
+/// The adapter treats incoming binary WebSocket messages as consecutive bytes
+/// and writes outbound byte chunks as binary WebSocket messages. Control and
+/// text messages are ignored except that close terminates the byte stream.
+pub fn websocket_byte_stream<S>(socket: WebSocketStream<S>) -> DuplexStream
+where
+    S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
+    let (mut socket_writer, mut socket_reader) = socket.split();
+    let (stream, peer) = tokio::io::duplex(WEBSOCKET_BYTE_STREAM_BUFFER);
+    let (mut peer_reader, mut peer_writer) = tokio::io::split(peer);
+
+    tokio::spawn(async move {
+        while let Some(result) = socket_reader.next().await {
+            match result {
+                Ok(message) if message.is_binary() => {
+                    let data = message.into_data();
+                    if peer_writer.write_all(data.as_ref()).await.is_err() {
+                        break;
+                    }
+                    if peer_writer.flush().await.is_err() {
+                        break;
+                    }
+                }
+                Ok(message) if message.is_close() => break,
+                Ok(_) => {}
+                Err(_) => break,
+            }
+        }
+        let _ = peer_writer.shutdown().await;
+    });
+
+    tokio::spawn(async move {
+        let mut buf = vec![0; WEBSOCKET_MESSAGE_BUFFER];
+        loop {
+            match peer_reader.read(&mut buf).await {
+                Ok(0) => {
+                    let _ = socket_writer.close().await;
+                    break;
+                }
+                Ok(n) => {
+                    let message = Message::Binary(Bytes::copy_from_slice(&buf[..n]));
+                    if socket_writer.send(message).await.is_err() {
+                        break;
+                    }
+                }
+                Err(_) => {
+                    let _ = socket_writer.close().await;
+                    break;
+                }
+            }
+        }
+    });
+
+    stream
+}

package/srpc/yamux.rs

@@ -0,0 +1,198 @@
+//! Yamux transport adapters for starpc.
+
+use async_trait::async_trait;
+use futures::future::poll_fn;
+use std::io;
+use std::pin::Pin;
+use std::sync::Arc;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::sync::{mpsc, oneshot};
+use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
+
+use crate::client::{OpenStream, PacketReceiver};
+use crate::error::{Error, Result};
+use crate::invoker::Invoker;
+use crate::rpc::PacketWriter;
+use crate::server::Server;
+use crate::transport::{create_packet_channel, DEFAULT_CHANNEL_BUFFER};
+
+type OpenResult = std::result::Result<::yamux::Stream, ::yamux::ConnectionError>;
+type OpenRequest = oneshot::Sender<OpenResult>;
+
+/// YamuxStreamOpener opens Starpc packet streams over one yamux connection.
+#[derive(Clone)]
+pub struct YamuxStreamOpener {
+    requests: mpsc::Sender<OpenRequest>,
+}
+
+impl YamuxStreamOpener {
+    /// client creates a client-mode yamux opener over a Tokio transport.
+    pub fn client<T>(transport: T) -> Self
+    where
+        T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        Self::client_with_config(transport, ::yamux::Config::default())
+    }
+
+    /// client_with_config creates a client-mode yamux opener with a custom config.
+    pub fn client_with_config<T>(transport: T, config: ::yamux::Config) -> Self
+    where
+        T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        let (requests, request_rx) = mpsc::channel(DEFAULT_CHANNEL_BUFFER);
+        spawn_client_driver(transport.compat(), config, request_rx);
+        Self { requests }
+    }
+
+    /// client_websocket creates a client-mode yamux opener over a WebSocket.
+    #[cfg(feature = "websocket")]
+    pub fn client_websocket<S>(socket: tokio_tungstenite::WebSocketStream<S>) -> Self
+    where
+        S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        Self::client(crate::websocket::websocket_byte_stream(socket))
+    }
+
+    /// client_websocket_with_config creates a client-mode yamux WebSocket opener with a custom config.
+    #[cfg(feature = "websocket")]
+    pub fn client_websocket_with_config<S>(
+        socket: tokio_tungstenite::WebSocketStream<S>,
+        config: ::yamux::Config,
+    ) -> Self
+    where
+        S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+    {
+        Self::client_with_config(crate::websocket::websocket_byte_stream(socket), config)
+    }
+}
+
+#[async_trait]
+impl OpenStream for YamuxStreamOpener {
+    async fn open_stream(&self) -> Result<(Arc<dyn PacketWriter>, PacketReceiver)> {
+        let (tx, rx) = oneshot::channel();
+        self.requests
+            .send(tx)
+            .await
+            .map_err(|_| Error::StreamClosed)?;
+
+        let stream = rx
+            .await
+            .map_err(|_| Error::StreamClosed)?
+            .map_err(connection_error)?;
+        let stream = stream.compat();
+        let (read_half, write_half) = tokio::io::split(stream);
+        Ok(create_packet_channel(read_half, write_half))
+    }
+}
+
+/// handle_server_connection accepts yamux streams and serves each as Starpc.
+pub async fn handle_server_connection<I, T>(
+    server: &Server<I>,
+    transport: T,
+    config: ::yamux::Config,
+) -> Result<()>
+where
+    I: Invoker + 'static,
+    T: AsyncRead + AsyncWrite + Send + Unpin + 'static,
+{
+    let mut connection =
+        ::yamux::Connection::new(transport.compat(), config, ::yamux::Mode::Server);
+
+    loop {
+        match poll_fn(|cx| connection.poll_next_inbound(cx)).await {
+            Some(Ok(stream)) => {
+                let server = server.clone_for_spawn();
+                tokio::spawn(async move {
+                    if let Err(err) = server.handle_stream(stream.compat()).await {
+                        server.report_error(err);
+                    }
+                });
+            }
+            Some(Err(err)) => return Err(connection_error(err)),
+            None => return Ok(()),
+        }
+    }
+}
+
+enum DriverEvent {
+    Opened(OpenRequest, OpenResult),
+    Closed {
+        pending: Option<OpenRequest>,
+        err: Option<::yamux::ConnectionError>,
+    },
+}
+
+fn spawn_client_driver<T>(
+    transport: T,
+    config: ::yamux::Config,
+    mut requests: mpsc::Receiver<OpenRequest>,
+) -> tokio::task::JoinHandle<()>
+where
+    T: futures::io::AsyncRead + futures::io::AsyncWrite + Send + Unpin + 'static,
+{
+    tokio::spawn(async move {
+        let mut connection = ::yamux::Connection::new(transport, config, ::yamux::Mode::Client);
+        let mut pending = None;
+        let mut requests_closed = false;
+
+        loop {
+            let event = poll_fn(|cx| {
+                if pending.is_none() && !requests_closed {
+                    match Pin::new(&mut requests).poll_recv(cx) {
+                        std::task::Poll::Ready(Some(request)) => pending = Some(request),
+                        std::task::Poll::Ready(None) => requests_closed = true,
+                        std::task::Poll::Pending => {}
+                    }
+                }
+
+                if let Some(request) = pending.take() {
+                    match connection.poll_new_outbound(cx) {
+                        std::task::Poll::Ready(result) => {
+                            return std::task::Poll::Ready(DriverEvent::Opened(request, result));
+                        }
+                        std::task::Poll::Pending => pending = Some(request),
+                    }
+                }
+
+                loop {
+                    match connection.poll_next_inbound(cx) {
+                        std::task::Poll::Ready(Some(Ok(_stream))) => continue,
+                        std::task::Poll::Ready(Some(Err(err))) => {
+                            return std::task::Poll::Ready(DriverEvent::Closed {
+                                pending: pending.take(),
+                                err: Some(err),
+                            });
+                        }
+                        std::task::Poll::Ready(None) => {
+                            return std::task::Poll::Ready(DriverEvent::Closed {
+                                pending: pending.take(),
+                                err: None,
+                            });
+                        }
+                        std::task::Poll::Pending => return std::task::Poll::Pending,
+                    }
+                }
+            })
+            .await;
+
+            match event {
+                DriverEvent::Opened(request, result) => {
+                    let _ = request.send(result);
+                }
+                DriverEvent::Closed { pending, err } => {
+                    if let Some(request) = pending {
+                        let _ = request.send(Err(err.unwrap_or(::yamux::ConnectionError::Closed)));
+                    }
+                    break;
+                }
+            }
+        }
+    })
+}
+
+fn connection_error(err: ::yamux::ConnectionError) -> Error {
+    match err {
+        ::yamux::ConnectionError::Io(err) => Error::Io(err),
+        other => Error::Io(io::Error::new(io::ErrorKind::Other, other.to_string())),
+    }
+}