wrpc_transport/frame/conn/
mod.rs

1use core::fmt::{Debug, Display};
2use core::future::Future;
3use core::mem;
4use core::pin::Pin;
5use core::task::{Context, Poll, ready};
6
7use std::sync::Arc;
8
9use bytes::{Buf as _, BufMut as _, Bytes, BytesMut};
10use futures::Sink as _;
11use pin_project_lite::pin_project;
12use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt as _};
13use tokio::sync::mpsc;
14use tokio::task::JoinSet;
15use tokio_stream::wrappers::ReceiverStream;
16use tokio_util::codec::Encoder;
17use tokio_util::io::StreamReader;
18use tokio_util::sync::PollSender;
19use tracing::{Instrument as _, Span, debug, error, instrument, trace};
20use wasm_tokio::{AsyncReadCore as _, AsyncReadLeb128 as _, Leb128Encoder};
21
22mod client;
23mod server;
24
25pub use client::*;
26pub use server::*;
27
28/// Error returned by [`Header::read`]
29pub enum HeaderReadError {
30    /// I/O error
31    IO(std::io::Error),
32    /// Protocol version is not supported
33    UnsupportedVersion(u8),
34}
35
36impl Debug for HeaderReadError {
37    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        match self {
39            Self::IO(err) => Debug::fmt(err, f),
40            Self::UnsupportedVersion(v) => write!(f, "unsupported version byte: {v}"),
41        }
42    }
43}
44
45impl Display for HeaderReadError {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        match self {
48            Self::IO(err) => Display::fmt(err, f),
49            Self::UnsupportedVersion(v) => write!(f, "unsupported version byte: {v}"),
50        }
51    }
52}
53
54impl core::error::Error for HeaderReadError {}
55
56/// wRPC invocation header
57#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
58pub struct Header {
59    /// Instance name of the function being called
60    pub instance: String,
61
62    /// Name of the function being called
63    pub name: String,
64}
65
66impl Header {
67    /// Reads the wRPC header from a byte stream
68    ///
69    /// # Errors
70    ///
71    /// Returns an error if reading the header fails
72    #[instrument(level = "trace", skip_all, ret(level = "trace"))]
73    pub async fn read(mut rx: impl AsyncRead + Unpin) -> Result<Self, HeaderReadError> {
74        let mut instance = String::default();
75        let mut name = String::default();
76        match rx.read_u8().await.map_err(HeaderReadError::IO)? {
77            0x00 => {
78                rx.read_core_name(&mut instance)
79                    .await
80                    .map_err(HeaderReadError::IO)?;
81                rx.read_core_name(&mut name)
82                    .await
83                    .map_err(HeaderReadError::IO)?;
84            }
85            v => return Err(HeaderReadError::UnsupportedVersion(v)),
86        }
87        Ok(Self { instance, name })
88    }
89}
90
91/// Index trie containing async stream subscriptions
92#[derive(Debug, Default)]
93enum IndexTrie {
94    #[default]
95    Empty,
96    Leaf {
97        tx: Option<mpsc::Sender<std::io::Result<Bytes>>>,
98        rx: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
99    },
100    IndexNode {
101        tx: Option<mpsc::Sender<std::io::Result<Bytes>>>,
102        rx: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
103        nested: Vec<Option<IndexTrie>>,
104    },
105    // TODO: Add partially-indexed `WildcardIndexNode`
106    WildcardNode {
107        tx: Option<mpsc::Sender<std::io::Result<Bytes>>>,
108        rx: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
109        nested: Option<Box<IndexTrie>>,
110    },
111}
112
113impl<'a>
114    From<(
115        &'a [Option<usize>],
116        mpsc::Sender<std::io::Result<Bytes>>,
117        Option<mpsc::Receiver<std::io::Result<Bytes>>>,
118    )> for IndexTrie
119{
120    fn from(
121        (path, tx, rx): (
122            &'a [Option<usize>],
123            mpsc::Sender<std::io::Result<Bytes>>,
124            Option<mpsc::Receiver<std::io::Result<Bytes>>>,
125        ),
126    ) -> Self {
127        match path {
128            [] => Self::Leaf { tx: Some(tx), rx },
129            [None, path @ ..] => Self::WildcardNode {
130                tx: None,
131                rx: None,
132                nested: Some(Box::new(Self::from((path, tx, rx)))),
133            },
134            [Some(i), path @ ..] => Self::IndexNode {
135                tx: None,
136                rx: None,
137                nested: {
138                    let n = i.saturating_add(1);
139                    let mut nested = Vec::with_capacity(n);
140                    nested.resize_with(n, Option::default);
141                    nested[*i] = Some(Self::from((path, tx, rx)));
142                    nested
143                },
144            },
145        }
146    }
147}
148
149impl<'a>
150    From<(
151        &'a [Option<usize>],
152        mpsc::Sender<std::io::Result<Bytes>>,
153        mpsc::Receiver<std::io::Result<Bytes>>,
154    )> for IndexTrie
155{
156    fn from(
157        (path, tx, rx): (
158            &'a [Option<usize>],
159            mpsc::Sender<std::io::Result<Bytes>>,
160            mpsc::Receiver<std::io::Result<Bytes>>,
161        ),
162    ) -> Self {
163        Self::from((path, tx, Some(rx)))
164    }
165}
166
167impl<'a> From<(&'a [Option<usize>], mpsc::Sender<std::io::Result<Bytes>>)> for IndexTrie {
168    fn from((path, tx): (&'a [Option<usize>], mpsc::Sender<std::io::Result<Bytes>>)) -> Self {
169        Self::from((path, tx, None))
170    }
171}
172
173impl<P: AsRef<[Option<usize>]>> FromIterator<P> for IndexTrie {
174    fn from_iter<T: IntoIterator<Item = P>>(iter: T) -> Self {
175        let mut root = Self::Empty;
176        for path in iter {
177            let (tx, rx) = mpsc::channel(16);
178            if !root.insert(path.as_ref(), tx, Some(rx)) {
179                return Self::Empty;
180            }
181        }
182        root
183    }
184}
185
186impl IndexTrie {
187    /// Takes the receiver
188    #[instrument(level = "trace", skip(self), ret(level = "trace"))]
189    fn take_rx(&mut self, path: &[usize]) -> Option<mpsc::Receiver<std::io::Result<Bytes>>> {
190        let Some((i, path)) = path.split_first() else {
191            return match self {
192                Self::Empty => None,
193                Self::Leaf { rx, .. } => rx.take(),
194                Self::IndexNode { tx, rx, nested } => {
195                    let rx = rx.take();
196                    if nested.is_empty() && tx.is_none() {
197                        *self = Self::Empty;
198                    }
199                    rx
200                }
201                Self::WildcardNode { tx, rx, nested } => {
202                    let rx = rx.take();
203                    if nested.is_none() && tx.is_none() {
204                        *self = Self::Empty;
205                    }
206                    rx
207                }
208            };
209        };
210        match self {
211            Self::Empty | Self::Leaf { .. } | Self::WildcardNode { .. } => None,
212            Self::IndexNode { nested, .. } => nested
213                .get_mut(*i)
214                .and_then(|nested| nested.as_mut().and_then(|nested| nested.take_rx(path))),
215            // TODO: Demux the subscription
216            //Self::WildcardNode { ref mut nested, .. } => {
217            //    nested.as_mut().and_then(|nested| nested.take(path))
218            //}
219        }
220    }
221
222    /// Gets a sender
223    #[instrument(level = "trace", skip(self), ret(level = "trace"))]
224    fn get_tx(&mut self, path: &[usize]) -> Option<mpsc::Sender<std::io::Result<Bytes>>> {
225        let Some((i, path)) = path.split_first() else {
226            return match self {
227                Self::Empty => None,
228                Self::Leaf { tx, .. } => tx.clone(),
229                Self::IndexNode { tx, .. } | Self::WildcardNode { tx, .. } => tx.clone(),
230            };
231        };
232        match self {
233            Self::Empty | Self::Leaf { .. } | Self::WildcardNode { .. } => None,
234            Self::IndexNode { nested, .. } => {
235                let nested = nested.get_mut(*i)?;
236                let nested = nested.as_mut()?;
237                nested.get_tx(path)
238            } // TODO: Demux the subscription
239              //Self::WildcardNode { ref mut nested, .. } => {
240              //    nested.as_mut().and_then(|nested| nested.take(path))
241              //}
242        }
243    }
244
245    /// Closes all senders in the trie
246    #[instrument(level = "trace", skip(self), ret(level = "trace"))]
247    fn close_tx(&mut self) {
248        match self {
249            Self::Empty => {}
250            Self::Leaf { tx, .. } => {
251                mem::take(tx);
252            }
253            Self::IndexNode { tx, nested, .. } => {
254                mem::take(tx);
255                for nested in nested.iter_mut().flatten() {
256                    nested.close_tx();
257                }
258            }
259            Self::WildcardNode { tx, nested, .. } => {
260                mem::take(tx);
261                if let Some(nested) = nested {
262                    nested.close_tx();
263                }
264            }
265        }
266    }
267
268    /// Inserts `sender` and `receiver` under a `path` - returns `false` if it failed and `true` if it succeeded.
269    /// Tree state after `false` is returned is undefined
270    #[instrument(level = "trace", skip(self, sender, receiver), ret(level = "trace"))]
271    fn insert(
272        &mut self,
273        path: &[Option<usize>],
274        sender: mpsc::Sender<std::io::Result<Bytes>>,
275        receiver: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
276    ) -> bool {
277        match self {
278            Self::Empty => {
279                *self = Self::from((path, sender, receiver));
280                true
281            }
282            Self::Leaf { .. } => {
283                let Some((i, path)) = path.split_first() else {
284                    return false;
285                };
286                let Self::Leaf { tx, rx } = mem::take(self) else {
287                    return false;
288                };
289                if let Some(i) = i {
290                    let n = i.saturating_add(1);
291                    let mut nested = Vec::with_capacity(n);
292                    nested.resize_with(n, Option::default);
293                    nested[*i] = Some(Self::from((path, sender, receiver)));
294                    *self = Self::IndexNode { tx, rx, nested };
295                } else {
296                    *self = Self::WildcardNode {
297                        tx,
298                        rx,
299                        nested: Some(Box::new(Self::from((path, sender, receiver)))),
300                    };
301                }
302                true
303            }
304            Self::IndexNode { tx, rx, nested } => match (&tx, &rx, path) {
305                (None, None, []) => {
306                    *tx = Some(sender);
307                    *rx = receiver;
308                    true
309                }
310                (_, _, [Some(i), path @ ..]) => {
311                    let cap = i.saturating_add(1);
312                    if nested.len() < cap {
313                        nested.resize_with(cap, Option::default);
314                    }
315                    let nested = &mut nested[*i];
316                    if let Some(nested) = nested {
317                        nested.insert(path, sender, receiver)
318                    } else {
319                        *nested = Some(Self::from((path, sender, receiver)));
320                        true
321                    }
322                }
323                _ => false,
324            },
325            Self::WildcardNode { tx, rx, nested } => match (&tx, &rx, path) {
326                (None, None, []) => {
327                    *tx = Some(sender);
328                    *rx = receiver;
329                    true
330                }
331                (_, _, [None, path @ ..]) => {
332                    if let Some(nested) = nested {
333                        nested.insert(path, sender, receiver)
334                    } else {
335                        *nested = Some(Box::new(Self::from((path, sender, receiver))));
336                        true
337                    }
338                }
339                _ => false,
340            },
341        }
342    }
343}
344
345pin_project! {
346    /// Incoming framed stream
347    #[project = IncomingProj]
348    pub struct Incoming {
349        #[pin]
350        rx: Option<StreamReader<ReceiverStream<std::io::Result<Bytes>>, Bytes>>,
351        path: Arc<[usize]>,
352        index: Arc<std::sync::Mutex<IndexTrie>>,
353        io: Arc<JoinSet<()>>,
354    }
355}
356
357impl Incoming {
358    /// Creates a new [Incoming] given an [`AsyncRead`], [`ConnHandler`] and a set of async paths.
359    /// `on_ingress` will be called once data ingress is complete.
360    pub fn new<T, P, Fut>(
361        mut rx: T,
362        paths: impl IntoIterator<Item = P>,
363        on_ingress: impl FnOnce(T, std::io::Result<()>) -> Fut + Send + 'static,
364    ) -> Self
365    where
366        T: AsyncRead + Unpin + Send + 'static,
367        P: AsRef<[Option<usize>]>,
368        Fut: Future<Output = ()> + Send,
369    {
370        let index = Arc::new(std::sync::Mutex::new(paths.into_iter().collect()));
371        let (rx_tx, rx_rx) = mpsc::channel(128);
372        let mut rx_io = JoinSet::new();
373        let span = Span::current();
374        rx_io.spawn({
375            let index = Arc::clone(&index);
376            async move {
377                let res = ingress(&mut rx, &index, rx_tx).await;
378                on_ingress(rx, res).await;
379                let Ok(mut index) = index.lock() else {
380                    error!("failed to lock index trie");
381                    return;
382                };
383                trace!("shutting down index trie");
384                index.close_tx();
385            }
386            .instrument(span.clone())
387        });
388        Self {
389            rx: Some(StreamReader::new(ReceiverStream::new(rx_rx))),
390            path: Arc::from([]),
391            index: Arc::clone(&index),
392            io: Arc::new(rx_io),
393        }
394    }
395
396    /// Index the incoming stream using a structural `path`, returning a handle to the
397    /// multiplexed sub-stream addressed by it.
398    #[instrument(level = "trace", skip(self), fields(path = ?self.path))]
399    pub fn index(&self, path: &[usize]) -> std::io::Result<Self> {
400        if path.is_empty() {
401            return Err(std::io::Error::new(
402                std::io::ErrorKind::InvalidInput,
403                "path cannot be empty",
404            ));
405        }
406        let path = if self.path.is_empty() {
407            Arc::from(path)
408        } else {
409            Arc::from([self.path.as_ref(), path].concat())
410        };
411        trace!("locking index trie");
412        let mut index = self
413            .index
414            .lock()
415            .map_err(|err| std::io::Error::other(err.to_string()))?;
416        trace!(?path, "taking index subscription");
417        let rx = index
418            .take_rx(&path)
419            .map(|rx| StreamReader::new(ReceiverStream::new(rx)));
420        Ok(Self {
421            rx,
422            path,
423            index: Arc::clone(&self.index),
424            io: Arc::clone(&self.io),
425        })
426    }
427}
428
429impl AsyncRead for Incoming {
430    #[instrument(level = "trace", skip_all, fields(path = ?self.path), ret(level = "trace"))]
431    fn poll_read(
432        mut self: Pin<&mut Self>,
433        cx: &mut Context<'_>,
434        buf: &mut tokio::io::ReadBuf<'_>,
435    ) -> Poll<std::io::Result<()>> {
436        if buf.remaining() == 0 {
437            return Poll::Ready(Ok(()));
438        }
439        trace!("reading");
440        let this = self.as_mut().project();
441        let Some(rx) = this.rx.as_pin_mut() else {
442            trace!("reader is closed");
443            return Poll::Ready(Ok(()));
444        };
445        ready!(rx.poll_read(cx, buf))?;
446        trace!(buf = ?buf.filled(), "read buffer");
447        if buf.filled().is_empty() {
448            self.rx.take();
449        }
450        Poll::Ready(Ok(()))
451    }
452}
453
454pin_project! {
455    /// Outgoing framed stream
456    #[project = OutgoingProj]
457    pub struct Outgoing {
458        #[pin]
459        tx: PollSender<(Bytes, Bytes)>,
460        path: Arc<[usize]>,
461        path_buf: Bytes,
462    }
463}
464
465impl Outgoing {
466    /// Creates a new [Outgoing] given an [`AsyncWrite`].
467    pub fn new<T, Fut>(
468        mut tx: T,
469        on_egress: impl FnOnce(T, std::io::Result<()>) -> Fut + Send + 'static,
470    ) -> Self
471    where
472        T: AsyncWrite + Unpin + Send + 'static,
473        Fut: Future<Output = ()> + Send,
474    {
475        let span = Span::current();
476        let (tx_tx, tx_rx) = mpsc::channel(128);
477        tokio::spawn(
478            async {
479                let res = egress(&mut tx, tx_rx).await;
480                on_egress(tx, res).await;
481            }
482            .instrument(span.clone()),
483        );
484        Self {
485            tx: PollSender::new(tx_tx),
486            path: Arc::from([]),
487            path_buf: Bytes::from_static(&[0]),
488        }
489    }
490
491    /// Index the outgoing stream using a structural `path`, returning a handle that writes
492    /// to the multiplexed sub-stream addressed by it.
493    #[instrument(level = "trace", skip(self), fields(path = ?self.path))]
494    pub fn index(&self, path: &[usize]) -> std::io::Result<Self> {
495        if path.is_empty() {
496            return Err(std::io::Error::new(
497                std::io::ErrorKind::InvalidInput,
498                "path cannot be empty",
499            ));
500        }
501        let path: Arc<[usize]> = if self.path.is_empty() {
502            Arc::from(path)
503        } else {
504            Arc::from([self.path.as_ref(), path].concat())
505        };
506        let mut buf = BytesMut::with_capacity(path.len().saturating_add(5));
507        let n = u32::try_from(path.len())
508            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
509        trace!(n, "encoding path length");
510        Leb128Encoder.encode(n, &mut buf)?;
511        for p in path.as_ref() {
512            let p = u32::try_from(*p)
513                .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
514            trace!(p, "encoding path element");
515            Leb128Encoder.encode(p, &mut buf)?;
516        }
517        Ok(Self {
518            tx: self.tx.clone(),
519            path,
520            path_buf: buf.freeze(),
521        })
522    }
523}
524
525impl AsyncWrite for Outgoing {
526    #[instrument(level = "trace", skip_all, fields(path = ?self.path, buf = format!("{buf:02x?}")), ret(level = "trace"))]
527    fn poll_write(
528        self: Pin<&mut Self>,
529        cx: &mut Context<'_>,
530        buf: &[u8],
531    ) -> Poll<std::io::Result<usize>> {
532        trace!("writing outgoing chunk");
533        let mut this = self.project();
534        ready!(this.tx.as_mut().poll_ready(cx))
535            .map_err(|err| std::io::Error::new(std::io::ErrorKind::BrokenPipe, err))?;
536        this.tx
537            .start_send((this.path_buf.clone(), Bytes::copy_from_slice(buf)))
538            .map_err(|err| std::io::Error::new(std::io::ErrorKind::BrokenPipe, err))?;
539        Poll::Ready(Ok(buf.len()))
540    }
541
542    #[instrument(level = "trace", skip_all, fields(path = ?self.path), ret(level = "trace"))]
543    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
544        Poll::Ready(Ok(()))
545    }
546
547    #[instrument(level = "trace", skip_all, fields(path = ?self.path), ret(level = "trace"))]
548    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
549        Poll::Ready(Ok(()))
550    }
551}
552
553#[instrument(level = "trace", skip_all, ret(level = "trace"))]
554async fn ingress(
555    mut rx: impl AsyncRead + Unpin,
556    index: &std::sync::Mutex<IndexTrie>,
557    param_tx: mpsc::Sender<std::io::Result<Bytes>>,
558) -> std::io::Result<()> {
559    loop {
560        trace!("reading path length");
561        let b = match rx.read_u8().await {
562            Ok(b) => b,
563            Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(()),
564            Err(err) => return Err(err),
565        };
566        let n = AsyncReadExt::chain([b].as_slice(), &mut rx)
567            .read_u32_leb128()
568            .await?;
569        let n = n
570            .try_into()
571            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
572        trace!(n, "read path length");
573        let tx = if n == 0 {
574            &param_tx
575        } else {
576            let mut path = Vec::with_capacity(n);
577            for i in 0..n {
578                trace!(i, "reading path element");
579                let p = rx.read_u32_leb128().await?;
580                let p = usize::try_from(p)
581                    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
582                path.push(p);
583            }
584            trace!(?path, "read path");
585
586            trace!("locking index trie");
587            let mut index = index
588                .lock()
589                .map_err(|err| std::io::Error::other(err.to_string()))?;
590            &index.get_tx(&path).ok_or_else(|| {
591                std::io::Error::new(
592                    std::io::ErrorKind::NotFound,
593                    format!("`{path:?}` subscription not found"),
594                )
595            })?
596        };
597        trace!("reading data length");
598        let n = rx.read_u32_leb128().await?;
599        let n = n
600            .try_into()
601            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
602        trace!(n, "read data length");
603        let mut buf = BytesMut::with_capacity(n);
604        buf.put_bytes(0, n);
605        trace!("reading data");
606        rx.read_exact(&mut buf).await?;
607        trace!(?buf, "read data");
608        tx.send(Ok(buf.freeze())).await.map_err(|_| {
609            std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stream receiver closed")
610        })?;
611    }
612}
613
614#[instrument(level = "trace", skip_all)]
615async fn egress(
616    mut tx: impl AsyncWrite + Unpin,
617    mut rx: mpsc::Receiver<(Bytes, Bytes)>,
618) -> std::io::Result<()> {
619    let mut buf = BytesMut::with_capacity(5);
620    trace!("waiting for next frame");
621    while let Some((path, data)) = rx.recv().await {
622        let data_len = u32::try_from(data.len())
623            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
624        buf.clear();
625        Leb128Encoder.encode(data_len, &mut buf)?;
626        let mut frame = path.chain(&mut buf).chain(data);
627        trace!(?frame, "writing egress frame");
628        tx.write_all_buf(&mut frame).await?;
629        tx.flush().await?;
630    }
631    trace!("shutting down outgoing stream");
632    tx.shutdown().await
633}
634
635/// Connection handler defines the connection I/O behavior.
636/// It is mostly useful for transports that may require additional clean up not already covered
637/// by [`AsyncWrite::shutdown`], for example.
638/// This API is experimental and may change in backwards-incompatible ways in the future.
639pub trait ConnHandler<Rx, Tx> {
640    /// Handle ingress completion
641    fn on_ingress(rx: Rx, res: std::io::Result<()>) -> impl Future<Output = ()> + Send {
642        _ = rx;
643        if let Err(err) = res {
644            error!(?err, "ingress failed");
645        } else {
646            debug!("ingress successfully complete");
647        }
648        async {}
649    }
650
651    /// Handle egress completion
652    fn on_egress(tx: Tx, res: std::io::Result<()>) -> impl Future<Output = ()> + Send {
653        _ = tx;
654        if let Err(err) = res {
655            error!(?err, "egress failed");
656        } else {
657            debug!("egress successfully complete");
658        }
659        async {}
660    }
661}
662
663impl<Rx, Tx> ConnHandler<Rx, Tx> for () {}
wrpc_transport/frame/conn/mod.rs

wrpc_transport/frame/conn/
mod.rs
