starpc 0.49.14 → 0.49.16

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/rwc-conn.go

@@ -7,6 +7,8 @@ import (
 	"os"
 	"sync"
 	"time"
+
+	"github.com/aperturerobotics/starpc/internal/contextutil"
 )
 
 // connPktSize is the size of the buffers to use for packets for the RwcConn.
@@ -68,7 +70,7 @@ func NewRwcConn(
 	laddr, raddr net.Addr,
 	bufferPacketN int,
 ) *RwcConn {
-	ctx, ctxCancel := context.WithCancel(ctx)
+	ctx, ctxCancel := contextutil.WithCancel(ctx)
 	if bufferPacketN <= 0 {
 		bufferPacketN = 10
 	}

package/srpc/server-rpc.go

@@ -91,13 +91,13 @@ func (r *ServerRPC) HandleCallStart(pkt *CallStart) error {
 // invokeRPC invokes the RPC after CallStart is received.
 func (r *ServerRPC) invokeRPC(serviceID, methodID string) {
 	// on the server side, the writer is closed by invokeRPC.
-	strm := NewMsgStream(r.ctx, r, r.ctxCancel)
+	strm := NewMsgStream(r.ctx, r, r.cancelContext)
 	ok, err := r.invoker.InvokeMethod(serviceID, methodID, strm)
 	if err == nil && !ok {
 		err = ErrUnimplemented
 	}
+	r.beginLocalCompletion()
 	outPkt := NewCallDataPacket(nil, false, true, err)
 	_ = r.writer.WritePacket(outPkt)
-	_ = r.writer.Close()
-	r.ctxCancel()
+	r.finishLocalCompletion()
 }

package/srpc/server.go

@@ -3,6 +3,8 @@ package srpc
 import (
 	"context"
 	"io"
+
+	"github.com/aperturerobotics/starpc/internal/contextutil"
 )
 
 // Server handles incoming RPC streams with a mux.
@@ -26,7 +28,7 @@ func (s *Server) GetInvoker() Invoker {
 // HandleStream handles an incoming stream and runs the read loop.
 // Uses length-prefixed packets.
 func (s *Server) HandleStream(ctx context.Context, rwc io.ReadWriteCloser) {
-	subCtx, subCtxCancel := context.WithCancel(ctx)
+	subCtx, subCtxCancel := contextutil.WithCancel(ctx)
 	defer subCtxCancel()
 	prw := NewPacketReadWriter(rwc)
 	serverRPC := NewServerRPC(subCtx, s.invoker, prw)

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/stream-pipe.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"io"
 	"sync"
+
+	"github.com/aperturerobotics/starpc/internal/contextutil"
 )
 
 // pipeStream implements an in-memory stream.
@@ -22,9 +24,9 @@ type pipeStream struct {
 // NewPipeStream constructs a new in-memory stream.
 func NewPipeStream(ctx context.Context) (Stream, Stream) {
 	s1 := &pipeStream{dataCh: make(chan []byte, 5)}
-	s1.ctx, s1.ctxCancel = context.WithCancel(ctx)
+	s1.ctx, s1.ctxCancel = contextutil.WithCancel(ctx)
 	s2 := &pipeStream{other: s1, dataCh: make(chan []byte, 5)}
-	s2.ctx, s2.ctxCancel = context.WithCancel(ctx)
+	s2.ctx, s2.ctxCancel = contextutil.WithCancel(ctx)
 	s1.other = s2
 	return s1, s2
 }

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