feat(transport): make conn I/O handler behavior configurable

Signed-off-by: Roman Volosatovs <[email protected]>

Author: rvolosatovs

crates/transport-quic/Cargo.toml

@@ -13,11 +13,11 @@ repository.workspace = true
 anyhow = { workspace = true, features = ["std"] }
 bytes = { workspace = true }
 quinn = { workspace = true, features = ["runtime-tokio"] }
+tracing = { workspace = true }
 wrpc-transport = { workspace = true }
 
 [dev-dependencies]
 futures = { workspace = true }
 test-log = { workspace = true, features = ["color", "log", "trace"] }
 tokio = { workspace = true, features = ["rt-multi-thread"] }
-tracing = { workspace = true }
 wrpc-test = { workspace = true, features = ["quic"] }

crates/transport-quic/src/lib.rs

@@ -2,14 +2,62 @@
 
 use anyhow::Context as _;
 use bytes::Bytes;
-use quinn::{Connection, RecvStream, SendStream};
-use wrpc_transport::frame::{invoke, Accept, Incoming, Outgoing};
+use quinn::{Connection, RecvStream, SendStream, VarInt};
+use tracing::{debug, error, trace, warn};
+use wrpc_transport::frame::{Accept, Incoming, InvokeBuilder, Outgoing};
 use wrpc_transport::Invoke;
 
-/// QUIC transport client
+/// QUIC server with graceful stream shutdown handling
+pub type Server = wrpc_transport::Server<(), RecvStream, SendStream, ConnHandler>;
+
+/// QUIC wRPC client
 #[derive(Clone, Debug)]
 pub struct Client(Connection);
 
+/// Graceful stream shutdown handler
+pub struct ConnHandler;
+
+const DONE: VarInt = VarInt::from_u32(1);
+
+impl wrpc_transport::frame::ConnHandler<RecvStream, SendStream> for ConnHandler {
+    async fn on_ingress(mut rx: RecvStream, res: std::io::Result<()>) {
+        if let Err(err) = res {
+            error!(?err, "ingress failed");
+        } else {
+            debug!("ingress successfully complete");
+        }
+        if let Err(err) = rx.stop(DONE) {
+            error!(?err, "failed to close stream");
+        }
+    }
+
+    async fn on_egress(mut tx: SendStream, res: std::io::Result<()>) {
+        if let Err(err) = res {
+            error!(?err, "egress failed");
+        } else {
+            debug!("egress successfully complete");
+        }
+        if let Err(err) = tx.finish() {
+            error!(?err, "failed to close stream");
+        }
+        match tx.stopped().await {
+            Ok(None) => {
+                trace!("stream successfully closed")
+            }
+            Ok(Some(code)) => {
+                if code == DONE {
+                    trace!("stream successfully closed")
+                } else {
+                    warn!(?code, "stream closed with code")
+                }
+            }
+            Err(err) => {
+                error!(?err, "failed to await stream close");
+            }
+        }
+    }
+}
+
 impl From<Connection> for Client {
     fn from(conn: Connection) -> Self {
         Self(conn)
@@ -37,7 +85,9 @@ impl Invoke for &Client {
             .open_bi()
             .await
             .context("failed to open parameter stream")?;
-        invoke(tx, rx, instance, func, params, paths).await
+        InvokeBuilder::<ConnHandler>::default()
+            .invoke(tx, rx, instance, func, params, paths)
+            .await
     }
 }
 

crates/transport-quic/tests/loopback.rs

@@ -15,7 +15,7 @@ async fn loopback() -> anyhow::Result<()> {
     wrpc_test::with_quic(|clt, srv| async {
         let clt = Client::from(clt);
         let srv_conn = Client::from(srv);
-        let srv = Arc::new(wrpc_transport::frame::Server::default());
+        let srv = Arc::new(wrpc_transport_quic::Server::new());
         let invocations = srv
             .serve("foo", "bar", [Box::from([Some(42), Some(0)])])
             .await

crates/transport/Cargo.toml

@@ -20,11 +20,11 @@ anyhow = { workspace = true, features = ["std"] }
 bytes = { workspace = true }
 futures = { workspace = true, features = ["std"] }
 pin-project-lite = { workspace = true }
+send-future = { workspace = true }
 tokio = { workspace = true, features = ["macros", "rt", "time"] }
 tokio-stream = { workspace = true }
 tokio-util = { workspace = true, features = ["codec", "io"] }
 tracing = { workspace = true, features = ["attributes"] }
-send-future = { workspace = true }
 wasm-tokio = { workspace = true, features = ["tracing"] }
 
 [target.'cfg(target_family = "wasm")'.dependencies]

crates/transport/src/frame/conn/client.rs

@@ -1,24 +1,70 @@
-use std::sync::Arc;
+use core::marker::PhantomData;
 
 use anyhow::Context as _;
 use bytes::{BufMut as _, Bytes, BytesMut};
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt as _};
-use tokio::sync::mpsc;
-use tokio::task::JoinSet;
-use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::codec::Encoder;
-use tokio_util::io::StreamReader;
-use tokio_util::sync::PollSender;
-use tracing::{debug, error, instrument, trace, Instrument as _};
+use tracing::{instrument, trace};
 use wasm_tokio::{CoreNameEncoder, CoreVecEncoderBytes};
 
-use crate::frame::conn::{egress, ingress, Incoming, Outgoing};
-use crate::frame::PROTOCOL;
+use crate::frame::conn::{Incoming, Outgoing};
+use crate::frame::{Conn, ConnHandler, PROTOCOL};
+
+/// Defines invocation behavior
+#[derive(Clone)]
+pub struct InvokeBuilder<H = ()>(PhantomData<H>)
+where
+    H: ?Sized;
+
+impl<H> InvokeBuilder<H> {
+    /// Invoke function `func` on instance `instance`
+    #[instrument(level = "trace", skip_all)]
+    pub async fn invoke<P, I, O>(
+        self,
+        mut tx: O,
+        rx: I,
+        instance: &str,
+        func: &str,
+        params: Bytes,
+        paths: impl AsRef<[P]> + Send,
+    ) -> anyhow::Result<(Outgoing, Incoming)>
+    where
+        P: AsRef<[Option<usize>]> + Send + Sync,
+        I: AsyncRead + Unpin + Send + 'static,
+        O: AsyncWrite + Unpin + Send + 'static,
+        H: ConnHandler<I, O>,
+    {
+        let mut buf = BytesMut::with_capacity(
+            17_usize // len(PROTOCOL) + len(instance) + len(func) + len([]) + len(params)
+                .saturating_add(instance.len())
+                .saturating_add(func.len())
+                .saturating_add(params.len()),
+        );
+        buf.put_u8(PROTOCOL);
+        CoreNameEncoder.encode(instance, &mut buf)?;
+        CoreNameEncoder.encode(func, &mut buf)?;
+        buf.put_u8(0);
+        CoreVecEncoderBytes.encode(params, &mut buf)?;
+        trace!(?buf, "writing invocation");
+        tx.write_all(&buf)
+            .await
+            .context("failed to initialize connection")?;
+
+        let Conn { tx, rx } = Conn::new::<H, _, _, _>(rx, tx, paths.as_ref());
+        Ok((tx, rx))
+    }
+}
+
+impl<H> Default for InvokeBuilder<H> {
+    fn default() -> Self {
+        Self(PhantomData)
+    }
+}
 
 /// Invoke function `func` on instance `instance`
 #[instrument(level = "trace", skip_all)]
 pub async fn invoke<P, I, O>(
-    mut tx: O,
+    tx: O,
     rx: I,
     instance: &str,
     func: &str,
@@ -30,65 +76,7 @@ where
     I: AsyncRead + Unpin + Send + 'static,
     O: AsyncWrite + Unpin + Send + 'static,
 {
-    let mut buf = BytesMut::with_capacity(
-        17_usize // len(PROTOCOL) + len(instance) + len(func) + len([]) + len(params)
-            .saturating_add(instance.len())
-            .saturating_add(func.len())
-            .saturating_add(params.len()),
-    );
-    buf.put_u8(PROTOCOL);
-    CoreNameEncoder.encode(instance, &mut buf)?;
-    CoreNameEncoder.encode(func, &mut buf)?;
-    buf.put_u8(0);
-    CoreVecEncoderBytes.encode(params, &mut buf)?;
-    trace!(?buf, "writing invocation");
-    tx.write_all(&buf)
+    InvokeBuilder::<()>::default()
+        .invoke(tx, rx, instance, func, params, paths)
         .await
-        .context("failed to initialize connection")?;
-
-    let index = Arc::new(std::sync::Mutex::new(paths.as_ref().iter().collect()));
-    let (results_tx, results_rx) = mpsc::channel(128);
-    let mut results_io = JoinSet::new();
-    results_io.spawn({
-        let index = Arc::clone(&index);
-        async move {
-            if let Err(err) = ingress(rx, &index, results_tx).await {
-                error!(?err, "result ingress failed");
-            } else {
-                debug!("result ingress successfully complete");
-            }
-            let Ok(mut index) = index.lock() else {
-                error!("failed to lock index trie");
-                return;
-            };
-            trace!("shutting down index trie");
-            index.close_tx();
-        }
-        .in_current_span()
-    });
-
-    let (params_tx, params_rx) = mpsc::channel(128);
-    tokio::spawn(
-        async {
-            if let Err(err) = egress(params_rx, tx).await {
-                error!(?err, "parameter egress failed");
-            } else {
-                debug!("parameter egress successfully complete");
-            }
-        }
-        .in_current_span(),
-    );
-    Ok((
-        Outgoing {
-            tx: PollSender::new(params_tx),
-            path: Arc::from([]),
-            path_buf: Bytes::from_static(&[0]),
-        },
-        Incoming {
-            rx: Some(StreamReader::new(ReceiverStream::new(results_rx))),
-            path: Arc::from([]),
-            index,
-            io: Arc::new(results_io),
-        },
-    ))
 }

crates/transport/src/frame/conn/mod.rs

@@ -1,3 +1,4 @@
+use core::future::Future;
 use core::mem;
 use core::pin::Pin;
 use core::task::{ready, Context, Poll};
@@ -15,7 +16,7 @@ use tokio_stream::wrappers::ReceiverStream;
 use tokio_util::codec::Encoder;
 use tokio_util::io::StreamReader;
 use tokio_util::sync::PollSender;
-use tracing::{instrument, trace};
+use tracing::{debug, error, instrument, trace, Instrument as _, Span};
 use wasm_tokio::{AsyncReadLeb128 as _, Leb128Encoder};
 
 use crate::Index;
@@ -484,10 +485,10 @@ async fn ingress(
     }
 }
 
-#[instrument(level = "trace", skip_all, ret(level = "trace"))]
+#[instrument(level = "trace", skip_all)]
 async fn egress(
-    mut rx: mpsc::Receiver<(Bytes, Bytes)>,
     mut tx: impl AsyncWrite + Unpin,
+    mut rx: mpsc::Receiver<(Bytes, Bytes)>,
 ) -> std::io::Result<()> {
     let mut buf = BytesMut::with_capacity(5);
     trace!("waiting for next frame");
@@ -503,3 +504,90 @@ async fn egress(
     trace!("shutting down outgoing stream");
     tx.shutdown().await
 }
+
+/// Connection handler defines the connection I/O behavior.
+/// It is mostly useful for transports that may require additional clean up not already covered
+/// by [AsyncWrite::shutdown], for example.
+/// This API is experimental and may change in backwards-incompatible ways in the future.
+pub trait ConnHandler<Rx, Tx> {
+    /// Handle ingress completion
+    fn on_ingress(rx: Rx, res: std::io::Result<()>) -> impl Future<Output = ()> + Send {
+        _ = rx;
+        if let Err(err) = res {
+            error!(?err, "ingress failed");
+        } else {
+            debug!("ingress successfully complete");
+        }
+        async {}
+    }
+
+    /// Handle egress completion
+    fn on_egress(tx: Tx, res: std::io::Result<()>) -> impl Future<Output = ()> + Send {
+        _ = tx;
+        if let Err(err) = res {
+            error!(?err, "egress failed");
+        } else {
+            debug!("egress successfully complete");
+        }
+        async {}
+    }
+}
+
+impl<Rx, Tx> ConnHandler<Rx, Tx> for () {}
+
+/// Peer connection
+pub(crate) struct Conn {
+    rx: Incoming,
+    tx: Outgoing,
+}
+
+impl Conn {
+    /// Creates a new [Conn] given an [AsyncRead], [ConnHandler] and a set of async paths
+    fn new<H, Rx, Tx, P>(mut rx: Rx, mut tx: Tx, paths: impl IntoIterator<Item = P>) -> Self
+    where
+        Rx: AsyncRead + Unpin + Send + 'static,
+        Tx: AsyncWrite + Unpin + Send + 'static,
+        H: ConnHandler<Rx, Tx>,
+        P: AsRef<[Option<usize>]>,
+    {
+        let index = Arc::new(std::sync::Mutex::new(paths.into_iter().collect()));
+        let (rx_tx, rx_rx) = mpsc::channel(128);
+        let mut rx_io = JoinSet::new();
+        let span = Span::current();
+        rx_io.spawn({
+            let index = Arc::clone(&index);
+            async move {
+                let res = ingress(&mut rx, &index, rx_tx).await;
+                H::on_ingress(rx, res).await;
+                let Ok(mut index) = index.lock() else {
+                    error!("failed to lock index trie");
+                    return;
+                };
+                trace!("shutting down index trie");
+                index.close_tx();
+            }
+            .instrument(span.clone())
+        });
+        let (tx_tx, tx_rx) = mpsc::channel(128);
+        tokio::spawn(
+            async {
+                let res = egress(&mut tx, tx_rx).await;
+                H::on_egress(tx, res).await;
+            }
+            .instrument(span.clone()),
+        );
+        Conn {
+            tx: Outgoing {
+                tx: PollSender::new(tx_tx),
+                path: Arc::from([]),
+                path_buf: Bytes::from_static(&[0]),
+            },
+            rx: Incoming {
+                rx: Some(StreamReader::new(ReceiverStream::new(rx_rx))),
+                path: Arc::from([]),
+                index: Arc::clone(&index),
+                io: Arc::new(rx_io),
+            },
+        }
+    }
+}