wrpc_transport/
invoke.rs

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

wrpc_transport/
invoke.rs
