wrpc_transport/
serve.rs

1//! wRPC transport server handle
2
3use core::mem;
4use core::pin::Pin;
5
6use std::sync::Arc;
7
8use anyhow::Context as _;
9use bytes::BytesMut;
10use futures::{Stream, TryStreamExt as _};
11use tokio::io::AsyncWriteExt as _;
12use tokio_util::codec::{Decoder as _, Encoder as _, FramedRead};
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            core::error::Error + Send + Sync + 'static,
84        <Results::Encoder as tokio_util::codec::Encoder<Results>>::Error:
85            core::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, mut outgoing, incoming)| {
91                async {
92                    let mut dec = FramedRead::new(incoming, Params::Decoder::default());
93                    debug!("receiving sync parameters");
94                    let params = match dec
95                        .decoder_mut()
96                        .decode(&mut BytesMut::default())
97                        .context("failed to decode sync parameters")?
98                    {
99                        Some(params) => params,
100                        None => {
101                            let params = dec
102                                .try_next()
103                                .await
104                                .context("failed to receive sync parameters")?;
105                            params.context("incomplete sync parameters")?
106                        }
107                    };
108                    trace!("received sync parameters");
109                    let rx = dec.decoder_mut().take_deferred();
110                    let incoming = BufferedIncoming {
111                        buffer: mem::take(dec.read_buffer_mut()),
112                        inner: dec.into_inner(),
113                    };
114                    let span = Span::current();
115                    Ok((
116                        cx,
117                        params,
118                        rx.map(|f| f(incoming, Vec::default())),
119                        move |results| {
120                            Box::pin(
121                                async move {
122                                    let mut enc = Results::Encoder::default();
123                                    let mut buf = BytesMut::default();
124                                    enc.encode(results, &mut buf)
125                                        .context("failed to encode synchronous results")?;
126                                    debug!("transmitting sync results");
127                                    outgoing
128                                        .write_bytes(buf.freeze())
129                                        .await
130                                        .context("failed to transmit synchronous results")?;
131                                    let tx = enc.take_deferred();
132                                    outgoing
133                                        .shutdown()
134                                        .await
135                                        .context("failed to shutdown synchronous return channel")?;
136                                    if let Some(tx) = tx {
137                                        debug!("transmitting async results");
138                                        tx(outgoing, Vec::default())
139                                            .await
140                                            .context("failed to write async results")?;
141                                    }
142                                    Ok(())
143                                }
144                                .instrument(span),
145                            ) as Pin<_>
146                        },
147                    ))
148                }
149                .instrument(span.clone())
150            }))
151        }
152    }
153}
154
155impl<T: Serve> ServeExt for T {}
156
157#[allow(dead_code)]
158#[cfg(test)]
159mod tests {
160    use bytes::Bytes;
161    use futures::{StreamExt as _, stream};
162
163    use super::*;
164
165    async fn call_serve<T: Serve>(s: &T) -> anyhow::Result<Vec<(T::Context, Outgoing, Incoming)>> {
166        let st = stream::empty()
167            .chain({
168                s.serve(
169                    "foo",
170                    "bar",
171                    Arc::from([Box::from([Some(42), None]), Box::from([None])]),
172                )
173                .await
174                .unwrap()
175            })
176            .chain({
177                s.serve(
178                    "foo",
179                    "bar",
180                    Arc::from([Box::from([Some(42), None]), Box::from([None])]),
181                )
182                .await
183                .unwrap()
184            })
185            .chain({
186                s.serve(
187                    "foo",
188                    "bar",
189                    Arc::from([Box::from([Some(42), None]), Box::from([None])]),
190                )
191                .await
192                .unwrap()
193            });
194        tokio::spawn(async move { st.try_collect().await })
195            .await
196            .unwrap()
197    }
198
199    fn serve_lifetime<T: Serve>(
200        s: &T,
201    ) -> impl Future<
202        Output = anyhow::Result<Pin<Box<dyn Stream<Item = anyhow::Result<T::Context>> + 'static>>>,
203    > {
204        let fut = s.serve(
205            "foo",
206            "bar",
207            Arc::from([Box::from([Some(42), None]), Box::from([None])]),
208        );
209        async move {
210            let st = fut.await.unwrap();
211            Ok(Box::pin(st.and_then(|(cx, _, _)| async { Ok(cx) }))
212                as Pin<Box<dyn Stream<Item = _>>>)
213        }
214    }
215
216    fn serve_values_lifetime<T: Serve>(
217        s: &T,
218    ) -> impl Future<
219        Output = anyhow::Result<Pin<Box<dyn Stream<Item = anyhow::Result<T::Context>> + 'static>>>,
220    > {
221        let fut = s.serve_values::<(Bytes,), (Bytes,)>(
222            "foo",
223            "bar",
224            Arc::from([Box::from([Some(42), None]), Box::from([None])]),
225        );
226        async move {
227            let st = fut.await.unwrap();
228            Ok(Box::pin(st.and_then(|(cx, _, _, tx)| async {
229                tx((Bytes::from("test"),)).await.unwrap();
230                Ok(cx)
231            })) as Pin<Box<dyn Stream<Item = _>>>)
232        }
233    }
234}
wrpc_transport/serve.rs

wrpc_transport/
serve.rs
