Skip to main content

wrpc_quic/
lib.rs

1//! wRPC QUIC transport
2
3use anyhow::Context as _;
4use bytes::Bytes;
5use quinn::{Connection, RecvStream, SendStream, VarInt};
6use tokio::io::BufReader;
7use tracing::{debug, error, trace, warn};
8use wrpc_transport::Invoke;
9use wrpc_transport::frame::{Incoming, InvokeBuilder, Outgoing};
10
11/// QUIC server with graceful stream shutdown handling
12pub type Server = wrpc_transport::Server<(), BufReader<RecvStream>, SendStream, ConnHandler>;
13
14/// QUIC wRPC client
15#[derive(Clone, Debug)]
16pub struct Client(Connection);
17
18/// Graceful stream shutdown handler
19pub struct ConnHandler;
20
21const DONE: VarInt = VarInt::from_u32(1);
22
23impl wrpc_transport::frame::ConnHandler<BufReader<RecvStream>, SendStream> for ConnHandler {
24    const DRAIN: bool = false;
25
26    async fn on_ingress(mut rx: BufReader<RecvStream>, res: std::io::Result<()>) {
27        if let Err(err) = res {
28            error!(?err, "ingress failed");
29        } else {
30            debug!("ingress successfully complete");
31        }
32        if let Err(err) = rx.get_mut().stop(DONE) {
33            debug!(?err, "failed to close incoming stream");
34        }
35    }
36
37    async fn on_egress(tx: SendStream, res: std::io::Result<()>) {
38        if let Err(err) = res {
39            error!(?err, "egress failed");
40        } else {
41            debug!("egress successfully complete");
42        }
43        match tx.stopped().await {
44            Ok(None) => {
45                trace!("stream successfully closed");
46            }
47            Ok(Some(code)) => {
48                if code == DONE {
49                    trace!("stream successfully closed");
50                } else {
51                    warn!(?code, "stream closed with code");
52                }
53            }
54            Err(err) => {
55                error!(?err, "failed to await stream close");
56            }
57        }
58    }
59}
60
61impl From<Connection> for Client {
62    fn from(conn: Connection) -> Self {
63        Self(conn)
64    }
65}
66
67impl Invoke for &Client {
68    type Context = ();
69
70    async fn invoke<P>(
71        &self,
72        (): Self::Context,
73        instance: &str,
74        func: &str,
75        params: Bytes,
76        paths: impl AsRef<[P]> + Send,
77    ) -> anyhow::Result<(Outgoing, Incoming)>
78    where
79        P: AsRef<[Option<usize>]> + Send + Sync,
80    {
81        let (tx, rx) = self
82            .0
83            .open_bi()
84            .await
85            .context("failed to open parameter stream")?;
86        InvokeBuilder::<ConnHandler>::default()
87            .invoke(tx, BufReader::new(rx), instance, func, params, paths)
88            .await
89    }
90}
91
92impl Invoke for Client {
93    type Context = ();
94
95    async fn invoke<P>(
96        &self,
97        (): Self::Context,
98        instance: &str,
99        func: &str,
100        params: Bytes,
101        paths: impl AsRef<[P]> + Send,
102    ) -> anyhow::Result<(Outgoing, Incoming)>
103    where
104        P: AsRef<[Option<usize>]> + Send + Sync,
105    {
106        (&self).invoke((), instance, func, params, paths).await
107    }
108}