wrpc_transport/
invoke.rs

1//! wRPC transport client handle
2
3use core::future::Future;
4use core::mem;
5use core::pin::pin;
6use core::time::Duration;
7
8use anyhow::Context as _;
9use bytes::{Bytes, BytesMut};
10use futures::TryStreamExt as _;
11use tokio::io::AsyncWriteExt as _;
12use tokio::{select, try_join};
13use tokio_util::codec::{Encoder as _, FramedRead};
14use tracing::{Instrument as _, debug, instrument, trace};
15
16use crate::frame::{Incoming, Outgoing};
17use crate::{BufferedIncoming, Deferred as _, TupleDecode, TupleEncode};
18
19/// Client-side handle to a wRPC transport
20///
21/// Invocations are always multiplexed over the wRPC framing layer, so the outgoing and
22/// incoming byte streams are the framed [`Outgoing`] and [`Incoming`] streams regardless of
23/// the underlying transport.
24pub trait Invoke: Send + Sync {
25    /// Transport-specific invocation context
26    type Context: Send + Sync;
27
28    /// Invoke function `func` on instance `instance`
29    fn invoke<P>(
30        &self,
31        cx: Self::Context,
32        instance: &str,
33        func: &str,
34        params: Bytes,
35        paths: impl AsRef<[P]> + Send,
36    ) -> impl Future<Output = anyhow::Result<(Outgoing, Incoming)>> + Send
37    where
38        P: AsRef<[Option<usize>]> + Send + Sync;
39}
40
41/// Wrapper struct returned by [`InvokeExt::timeout`]
42#[derive(Clone, Copy, Debug, Eq, PartialEq)]
43pub struct Timeout<'a, T: ?Sized> {
44    /// Inner [Invoke]
45    pub inner: &'a T,
46    /// Invocation timeout
47    pub timeout: Duration,
48}
49
50impl<T: Invoke> Invoke for Timeout<'_, T> {
51    type Context = T::Context;
52
53    #[instrument(level = "trace", skip(self, cx, params, paths))]
54    async fn invoke<P>(
55        &self,
56        cx: Self::Context,
57        instance: &str,
58        func: &str,
59        params: Bytes,
60        paths: impl AsRef<[P]> + Send,
61    ) -> anyhow::Result<(Outgoing, Incoming)>
62    where
63        P: AsRef<[Option<usize>]> + Send + Sync,
64    {
65        tokio::time::timeout(
66            self.timeout,
67            self.inner.invoke(cx, instance, func, params, paths),
68        )
69        .await
70        .context("invocation timed out")?
71    }
72}
73
74/// Wrapper struct returned by [`InvokeExt::timeout_owned`]
75#[derive(Clone, Copy, Debug, Eq, PartialEq)]
76pub struct TimeoutOwned<T> {
77    /// Inner [Invoke]
78    pub inner: T,
79    /// Invocation timeout
80    pub timeout: Duration,
81}
82
83impl<T: Invoke> Invoke for TimeoutOwned<T> {
84    type Context = T::Context;
85
86    #[instrument(level = "trace", skip(self, cx, params, paths))]
87    async fn invoke<P>(
88        &self,
89        cx: Self::Context,
90        instance: &str,
91        func: &str,
92        params: Bytes,
93        paths: impl AsRef<[P]> + Send,
94    ) -> anyhow::Result<(Outgoing, Incoming)>
95    where
96        P: AsRef<[Option<usize>]> + Send + Sync,
97    {
98        self.inner
99            .timeout(self.timeout)
100            .invoke(cx, instance, func, params, paths)
101            .await
102    }
103}
104
105/// Extension trait for [Invoke]
106pub trait InvokeExt: Invoke {
107    /// Invoke function `func` on instance `instance` using typed `Params` and `Results`
108    #[instrument(level = "trace", skip(self, cx, params, paths))]
109    fn invoke_values<P, Params, Results, Paths>(
110        &self,
111        cx: Self::Context,
112        instance: &str,
113        func: &str,
114        params: Params,
115        paths: Paths,
116    ) -> impl Future<
117        Output = anyhow::Result<(
118            Results,
119            Option<
120                impl Future<Output = anyhow::Result<()>>
121                + Send
122                + 'static
123                + use<Self, P, Params, Results, Paths>,
124            >,
125        )>,
126    > + Send
127    where
128        Paths: AsRef<[P]> + Send,
129        P: AsRef<[Option<usize>]> + Send + Sync,
130        Params: TupleEncode + Send,
131        Results: TupleDecode + Send,
132        <Params::Encoder as tokio_util::codec::Encoder<Params>>::Error:
133            std::error::Error + Send + Sync + 'static,
134        <Results::Decoder as tokio_util::codec::Decoder>::Error:
135            std::error::Error + Send + Sync + 'static,
136    {
137        async {
138            let mut buf = BytesMut::default();
139            let mut enc = Params::Encoder::default();
140            trace!("encoding parameters");
141            enc.encode(params, &mut buf)
142                .context("failed to encode parameters")?;
143            debug!("invoking function");
144            let (mut outgoing, incoming) = self
145                .invoke(cx, instance, func, buf.freeze(), paths)
146                .await
147                .context("failed to invoke function")?;
148            trace!("shutdown synchronous parameter channel");
149            outgoing
150                .shutdown()
151                .await
152                .context("failed to shutdown synchronous parameter channel")?;
153            let mut tx = enc.take_deferred().map(|tx| {
154                tokio::spawn(
155                    async {
156                        debug!("transmitting async parameters");
157                        tx(outgoing, Vec::default())
158                            .await
159                            .context("failed to write async parameters")
160                    }
161                    .in_current_span(),
162                )
163            });
164
165            let mut dec = FramedRead::new(incoming, Results::Decoder::default());
166            let results = async {
167                debug!("receiving sync results");
168                dec.try_next()
169                    .await
170                    .context("failed to receive sync results")?
171                    .context("incomplete results")
172            };
173            let results = match tx.take() {
174                Some(mut fut) => {
175                    let mut results = pin!(results);
176                    select! {
177                        res = &mut results => {
178                            tx = Some(fut);
179                            res?
180                        }
181                        res = &mut fut => {
182                            res??;
183                            results.await?
184                        }
185                    }
186                }
187                _ => results.await?,
188            };
189            trace!("received sync results");
190            let buffer = mem::take(dec.read_buffer_mut());
191            let rx = dec.decoder_mut().take_deferred();
192            let incoming = BufferedIncoming {
193                buffer,
194                inner: dec.into_inner(),
195            };
196            Ok((
197                results,
198                (tx.is_some() || rx.is_some()).then_some(
199                    async {
200                        match (tx, rx) {
201                            (Some(tx), Some(rx)) => {
202                                try_join!(
203                                    async {
204                                        debug!("receiving async results");
205                                        rx(incoming, Vec::default())
206                                            .await
207                                            .context("receiving async results failed")
208                                    },
209                                    async {
210                                        tx.await.context("transmitting async parameters failed")?
211                                    }
212                                )?;
213                            }
214                            (Some(tx), None) => {
215                                tx.await.context("transmitting async parameters failed")??;
216                            }
217                            (None, Some(rx)) => {
218                                debug!("receiving async results");
219                                rx(incoming, Vec::default())
220                                    .await
221                                    .context("receiving async results failed")?;
222                            }
223                            _ => {}
224                        }
225                        Ok(())
226                    }
227                    .in_current_span(),
228                ),
229            ))
230        }
231    }
232
233    /// Invoke function `func` on instance `instance` using typed `Params` and `Results`
234    /// This is like [`Self::invoke_values`], but it only results once all I/O is done
235    #[instrument(level = "trace", skip_all)]
236    fn invoke_values_blocking<P, Params, Results>(
237        &self,
238        cx: Self::Context,
239        instance: &str,
240        func: &str,
241        params: Params,
242        paths: impl AsRef<[P]> + Send,
243    ) -> impl Future<Output = anyhow::Result<Results>> + Send
244    where
245        P: AsRef<[Option<usize>]> + Send + Sync,
246        Params: TupleEncode + Send,
247        Results: TupleDecode + Send,
248        <Params::Encoder as tokio_util::codec::Encoder<Params>>::Error:
249            std::error::Error + Send + Sync + 'static,
250        <Results::Decoder as tokio_util::codec::Decoder>::Error:
251            std::error::Error + Send + Sync + 'static,
252    {
253        async {
254            let (ret, io) = self
255                .invoke_values(cx, instance, func, params, paths)
256                .await?;
257            if let Some(io) = io {
258                trace!("awaiting I/O completion");
259                io.await?;
260            }
261            Ok(ret)
262        }
263    }
264
265    /// Returns a [`Timeout`], wrapping [Self] with an implementation of [Invoke], which will
266    /// error, if call to [`Invoke::invoke`] does not return within a supplied `timeout`
267    fn timeout(&self, timeout: Duration) -> Timeout<'_, Self> {
268        Timeout {
269            inner: self,
270            timeout,
271        }
272    }
273
274    /// This is like [`InvokeExt::timeout`], but moves [Self] and returns corresponding [`TimeoutOwned`]
275    fn timeout_owned(self, timeout: Duration) -> TimeoutOwned<Self>
276    where
277        Self: Sized,
278    {
279        TimeoutOwned {
280            inner: self,
281            timeout,
282        }
283    }
284}
285
286impl<T: Invoke> InvokeExt for T {}
287
288#[allow(dead_code)]
289#[cfg(test)]
290mod tests {
291    use core::future::Future;
292    use core::pin::Pin;
293
294    use std::sync::Arc;
295
296    use bytes::Bytes;
297    use futures::{Stream, StreamExt as _};
298
299    use super::*;
300
301    #[allow(clippy::manual_async_fn)]
302    fn invoke_values_send<T>() -> impl Future<
303        Output = anyhow::Result<(
304            Pin<Box<dyn Stream<Item = Vec<Pin<Box<dyn Future<Output = String> + Send>>>> + Send>>,
305        )>,
306    > + Send
307    where
308        T: Invoke<Context = ()> + Default,
309    {
310        async {
311            let wrpc = T::default();
312            let ((r0,), _) = wrpc
313                .invoke_values(
314                    (),
315                    "wrpc-test:integration/async",
316                    "with-streams",
317                    (),
318                    [[None].as_slice()],
319                )
320                .await?;
321            Ok(r0)
322        }
323    }
324
325    async fn call_invoke<T: Invoke>(
326        i: &T,
327        cx: T::Context,
328        paths: Arc<[Arc<[Option<usize>]>]>,
329    ) -> anyhow::Result<(Outgoing, Incoming)> {
330        i.invoke(cx, "foo", "bar", Bytes::default(), &paths).await
331    }
332
333    async fn call_invoke_async<T>() -> anyhow::Result<(Pin<Box<dyn Stream<Item = Bytes> + Send>>,)>
334    where
335        T: Invoke<Context = ()> + Default,
336    {
337        let wrpc = T::default();
338        let ((r0,), _) = wrpc
339            .invoke_values(
340                (),
341                "wrpc-test:integration/async",
342                "with-streams",
343                (),
344                [
345                    [Some(1), Some(2)].as_slice(),
346                    [None].as_slice(),
347                    [Some(42)].as_slice(),
348                ],
349            )
350            .await?;
351        Ok(r0)
352    }
353
354    trait Handler {
355        fn foo() -> impl Future<Output = anyhow::Result<()>>;
356    }
357
358    impl<T> Handler for T
359    where
360        T: Invoke<Context = ()> + Default,
361    {
362        async fn foo() -> anyhow::Result<()> {
363            let (st,) = call_invoke_async::<Self>().await?;
364            st.collect::<Vec<_>>().await;
365            Ok(())
366        }
367    }
368}
wrpc_transport/invoke.rs

wrpc_transport/
invoke.rs
