wrpc_transport/frame/
oneshot.rs

1//! wRPC transport stream framing
2
3use bytes::Bytes;
4use tokio::io::{
5    AsyncRead, AsyncWrite, BufReader, DuplexStream, ReadHalf, WriteHalf, duplex, split,
6};
7use tracing::instrument;
8
9use crate::Invoke;
10use crate::frame::{Incoming, Outgoing, invoke};
11
12/// [Invoke] implementation in terms of a single stream pair.
13///
14/// [`Invoke::invoke`] can only be called at most once on [Oneshot],
15/// repeated calls will return an error.
16#[derive(Debug)]
17pub struct Oneshot<I, O>(std::sync::Mutex<Option<(I, O)>>);
18
19impl<I, O> From<(I, O)> for Oneshot<I, O> {
20    fn from((rx, tx): (I, O)) -> Self {
21        Self(std::sync::Mutex::new(Some((rx, tx))))
22    }
23}
24
25impl From<DuplexStream> for Oneshot<ReadHalf<DuplexStream>, WriteHalf<DuplexStream>> {
26    fn from(stream: DuplexStream) -> Self {
27        split(stream).into()
28    }
29}
30
31impl Oneshot<ReadHalf<DuplexStream>, WriteHalf<DuplexStream>> {
32    /// Creates a pair of [Oneshot] and server-side [`DuplexStream`] using [`tokio::io::duplex`].
33    #[must_use]
34    pub fn duplex(max_buf_size: usize) -> (Self, DuplexStream) {
35        let (clt, srv) = duplex(max_buf_size);
36        (clt.into(), srv)
37    }
38}
39
40impl<I, O> Oneshot<I, O> {
41    /// Returns the inner stream pair if [Oneshot] has not been used yet or an error.
42    pub fn try_take_inner(&self) -> std::io::Result<(I, O)> {
43        match self.0.try_lock().map(|mut stream| stream.take()) {
44            Ok(Some((rx, tx))) => Ok((rx, tx)),
45            Ok(None) | Err(std::sync::TryLockError::WouldBlock) => Err(std::io::Error::new(
46                std::io::ErrorKind::UnexpectedEof,
47                "stream was already used",
48            )),
49            Err(std::sync::TryLockError::Poisoned(..)) => {
50                Err(std::io::Error::other("stream lock poisoned"))
51            }
52        }
53    }
54}
55
56impl<I, O> Invoke for Oneshot<I, O>
57where
58    I: AsyncRead + Send + Unpin + 'static,
59    O: AsyncWrite + Send + Unpin + 'static,
60{
61    type Context = ();
62
63    async fn invoke<P>(
64        &self,
65        cx: Self::Context,
66        instance: &str,
67        func: &str,
68        params: Bytes,
69        paths: impl AsRef<[P]> + Send,
70    ) -> anyhow::Result<(Outgoing, Incoming)>
71    where
72        P: AsRef<[Option<usize>]> + Send + Sync,
73    {
74        (&self).invoke(cx, instance, func, params, paths).await
75    }
76}
77
78impl<I, O> Invoke for &Oneshot<I, O>
79where
80    I: AsyncRead + Send + Unpin + 'static,
81    O: AsyncWrite + Send + Unpin + 'static,
82{
83    type Context = ();
84
85    #[instrument(level = "trace", skip(self, paths, params), fields(params = format!("{params:02x?}")))]
86    fn invoke<P>(
87        &self,
88        (): Self::Context,
89        instance: &str,
90        func: &str,
91        params: Bytes,
92        paths: impl AsRef<[P]> + Send,
93    ) -> impl Future<Output = anyhow::Result<(Outgoing, Incoming)>>
94    where
95        P: AsRef<[Option<usize>]> + Send + Sync,
96    {
97        let stream = self.try_take_inner();
98        async move {
99            let (rx, tx) = stream?;
100            invoke(tx, BufReader::new(rx), instance, func, params, paths).await
101        }
102    }
103}
wrpc_transport/frame/oneshot.rs

wrpc_transport/frame/
oneshot.rs
