wrpc_transport/
serve.rs

1//! wRPC transport server handle
2
3use core::future::Future;
4use core::mem;
5use core::pin::Pin;
6
7use std::sync::Arc;
8
9use anyhow::{Context as _, bail};
10use futures::{SinkExt as _, Stream, TryStreamExt as _};
11use tokio::io::AsyncWriteExt as _;
12use tokio_util::codec::{FramedRead, FramedWrite};
13use tracing::{Instrument as _, Span, debug, instrument, trace};
14
15use crate::frame::{Incoming, Outgoing};
16use crate::{BufferedIncoming, Deferred as _, TupleDecode, TupleEncode};
17
18/// Server-side handle to a wRPC transport
19///
20/// Invocations are always multiplexed over the wRPC framing layer, so the outgoing and
21/// incoming byte streams are the framed [`Outgoing`] and [`Incoming`] streams regardless of
22/// the underlying transport.
23pub trait Serve: Sync {
24    /// Transport-specific invocation context
25    type Context: Send + Sync + 'static;
26
27    /// Serve function `func` from instance `instance`
28    fn serve(
29        &self,
30        instance: &str,
31        func: &str,
32        paths: Arc<[Box<[Option<usize>]>]>,
33    ) -> impl Future<
34        Output = anyhow::Result<
35            impl Stream<Item = anyhow::Result<(Self::Context, Outgoing, Incoming)>>
36            + Send
37            + 'static
38            + use<Self>,
39        >,
40    > + Send;
41}
42
43/// Extension trait for [Serve]
44pub trait ServeExt: Serve {
45    /// Serve function `func` from instance `instance` using typed `Params` and `Results`
46    #[instrument(level = "trace", skip(self, paths))]
47    fn serve_values<Params, Results>(
48        &self,
49        instance: &str,
50        func: &str,
51        paths: Arc<[Box<[Option<usize>]>]>,
52    ) -> impl Future<
53        Output = anyhow::Result<
54            impl Stream<
55                Item = anyhow::Result<(
56                    Self::Context,
57                    Params,
58                    Option<
59                        impl Future<Output = std::io::Result<()>>
60                        + Send
61                        + Unpin
62                        + 'static
63                        + use<Self, Params, Results>,
64                    >,
65                    impl FnOnce(
66                        Results,
67                    )
68                        -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + 'static>>
69                    + Send
70                    + 'static
71                    + use<Self, Params, Results>,
72                )>,
73            >
74            + Send
75            + 'static
76            + use<Self, Params, Results>,
77        >,
78    > + Send
79    where
80        Params: TupleDecode + Send + 'static,
81        Results: TupleEncode + Send + 'static,
82        <Params::Decoder as tokio_util::codec::Decoder>::Error:
83            std::error::Error + Send + Sync + 'static,
84        <Results::Encoder as tokio_util::codec::Encoder<Results>>::Error:
85            std::error::Error + Send + Sync + 'static,
86    {
87        let span = Span::current();
88        async {
89            let invocations = self.serve(instance, func, paths).await?;
90            Ok(invocations.and_then(move |(cx, outgoing, incoming)| {
91                async {
92                    let mut dec = FramedRead::new(incoming, Params::Decoder::default());
93                    debug!("receiving sync parameters");
94                    let Some(params) = dec
95                        .try_next()
96                        .await
97                        .context("failed to receive sync parameters")?
98                    else {
99                        bail!("incomplete sync parameters")
100                    };
101                    trace!("received sync parameters");
102                    let rx = dec.decoder_mut().take_deferred();
103                    let buffer = mem::take(dec.read_buffer_mut());
104                    let span = Span::current();
105                    Ok((
106                        cx,
107                        params,
108                        rx.map(|f| {
109                            f(
110                                BufferedIncoming {
111                                    buffer,
112                                    inner: dec.into_inner(),
113                                },
114                                Vec::default(),
115                            )
116                        }),
117                        move |results| {
118                            Box::pin(
119                                async {
120                                    let mut enc =
121                                        FramedWrite::new(outgoing, Results::Encoder::default());
122                                    debug!("transmitting sync results");
123                                    enc.send(results)
124                                        .await
125                                        .context("failed to transmit synchronous results")?;
126                                    let tx = enc.encoder_mut().take_deferred();
127                                    let mut outgoing = enc.into_inner();
128                                    outgoing
129                                        .shutdown()
130                                        .await
131                                        .context("failed to shutdown synchronous return channel")?;
132                                    if let Some(tx) = tx {
133                                        debug!("transmitting async results");
134                                        tx(outgoing, Vec::default())
135                                            .await
136                                            .context("failed to write async results")?;
137                                    }
138                                    Ok(())
139                                }
140                                .instrument(span),
141                            ) as Pin<_>
142                        },
143                    ))
144                }
145                .instrument(span.clone())
146            }))
147        }
148    }
149}
150
151impl<T: Serve> ServeExt for T {}
152
153#[allow(dead_code)]
154#[cfg(test)]
155mod tests {
156    use bytes::Bytes;
157    use futures::{StreamExt as _, TryStreamExt as _, stream};
158
159    use super::*;
160
161    async fn call_serve<T: Serve>(s: &T) -> anyhow::Result<Vec<(T::Context, Outgoing, Incoming)>> {
162        let st = stream::empty()
163            .chain({
164                s.serve(
165                    "foo",
166                    "bar",
167                    Arc::from([Box::from([Some(42), None]), Box::from([None])]),
168                )
169                .await
170                .unwrap()
171            })
172            .chain({
173                s.serve(
174                    "foo",
175                    "bar",
176                    Arc::from([Box::from([Some(42), None]), Box::from([None])]),
177                )
178                .await
179                .unwrap()
180            })
181            .chain({
182                s.serve(
183                    "foo",
184                    "bar",
185                    Arc::from([Box::from([Some(42), None]), Box::from([None])]),
186                )
187                .await
188                .unwrap()
189            });
190        tokio::spawn(async move { st.try_collect().await })
191            .await
192            .unwrap()
193    }
194
195    fn serve_lifetime<T: Serve>(
196        s: &T,
197    ) -> impl Future<
198        Output = anyhow::Result<Pin<Box<dyn Stream<Item = anyhow::Result<T::Context>> + 'static>>>,
199    > {
200        let fut = s.serve(
201            "foo",
202            "bar",
203            Arc::from([Box::from([Some(42), None]), Box::from([None])]),
204        );
205        async move {
206            let st = fut.await.unwrap();
207            Ok(Box::pin(st.and_then(|(cx, _, _)| async { Ok(cx) }))
208                as Pin<Box<dyn Stream<Item = _>>>)
209        }
210    }
211
212    fn serve_values_lifetime<T: Serve>(
213        s: &T,
214    ) -> impl Future<
215        Output = anyhow::Result<Pin<Box<dyn Stream<Item = anyhow::Result<T::Context>> + 'static>>>,
216    > {
217        let fut = s.serve_values::<(Bytes,), (Bytes,)>(
218            "foo",
219            "bar",
220            Arc::from([Box::from([Some(42), None]), Box::from([None])]),
221        );
222        async move {
223            let st = fut.await.unwrap();
224            Ok(Box::pin(st.and_then(|(cx, _, _, tx)| async {
225                tx((Bytes::from("test"),)).await.unwrap();
226                Ok(cx)
227            })) as Pin<Box<dyn Stream<Item = _>>>)
228        }
229    }
230}
wrpc_transport/serve.rs

wrpc_transport/
serve.rs
