wrpc_transport/frame/conn/
mod.rs

1use core::fmt::{Debug, Display};
2use core::future::poll_fn;
3use core::mem;
4use core::pin::{Pin, pin};
5use core::sync::atomic::{AtomicBool, Ordering};
6use core::task::{Context, Poll, ready};
7
8use std::sync::{Arc, PoisonError};
9
10use bytes::{Buf as _, Bytes, BytesMut};
11use pin_project_lite::pin_project;
12use tokio::io::{AsyncBufRead, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt as _};
13use tokio::select;
14use tokio::sync::{mpsc, oneshot};
15use tokio_stream::wrappers::ReceiverStream;
16use tokio_util::codec::Encoder;
17use tokio_util::io::StreamReader;
18use tokio_util::sync::PollSender;
19use tokio_util::task::AbortOnDropHandle;
20use tracing::{Instrument as _, Span, debug, error, instrument, trace};
21use wasm_tokio::{
22    AsyncReadCore as _, AsyncReadLeb128 as _, DEFAULT_MAX_INITIAL_CAPACITY, Leb128Encoder,
23};
24
25use crate::frame::MAX_INITIAL_PATH_CAPACITY;
26use crate::{copy_io_error, error_receiver_closed};
27
28mod client;
29mod server;
30
31pub use client::*;
32pub use server::*;
33
34/// Error returned by [`Header::read`]
35pub enum HeaderReadError {
36    /// I/O error
37    IO(std::io::Error),
38    /// Protocol version is not supported
39    UnsupportedVersion(u8),
40}
41
42impl Debug for HeaderReadError {
43    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
44        match self {
45            Self::IO(err) => Debug::fmt(err, f),
46            Self::UnsupportedVersion(v) => write!(f, "unsupported version byte: {v}"),
47        }
48    }
49}
50
51impl Display for HeaderReadError {
52    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
53        match self {
54            Self::IO(err) => Display::fmt(err, f),
55            Self::UnsupportedVersion(v) => write!(f, "unsupported version byte: {v}"),
56        }
57    }
58}
59
60impl core::error::Error for HeaderReadError {}
61
62/// wRPC invocation header
63#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
64pub struct Header {
65    /// Instance name of the function being called
66    pub instance: String,
67
68    /// Name of the function being called
69    pub name: String,
70}
71
72impl Header {
73    /// Reads the wRPC header from a byte stream
74    ///
75    /// # Errors
76    ///
77    /// Returns an error if reading the header fails
78    #[instrument(level = "trace", skip_all, ret(level = "trace"))]
79    pub async fn read(mut rx: impl AsyncRead + Unpin) -> Result<Self, HeaderReadError> {
80        let mut instance = String::default();
81        let mut name = String::default();
82        match rx.read_u8().await.map_err(HeaderReadError::IO)? {
83            0x00 => {
84                rx.read_core_name(&mut instance)
85                    .await
86                    .map_err(HeaderReadError::IO)?;
87                rx.read_core_name(&mut name)
88                    .await
89                    .map_err(HeaderReadError::IO)?;
90            }
91            v => return Err(HeaderReadError::UnsupportedVersion(v)),
92        }
93        Ok(Self { instance, name })
94    }
95}
96
97/// Index trie containing async stream subscriptions
98#[derive(Debug, Default)]
99enum IndexTrie {
100    #[default]
101    Empty,
102    Leaf {
103        tx: Option<mpsc::Sender<std::io::Result<Bytes>>>,
104        rx: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
105    },
106    IndexNode {
107        tx: Option<mpsc::Sender<std::io::Result<Bytes>>>,
108        rx: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
109        nested: Vec<Option<IndexTrie>>,
110    },
111    // TODO: Add partially-indexed `WildcardIndexNode`
112    WildcardNode {
113        tx: Option<mpsc::Sender<std::io::Result<Bytes>>>,
114        rx: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
115        nested: Option<Box<IndexTrie>>,
116    },
117}
118
119impl<'a>
120    From<(
121        &'a [Option<usize>],
122        mpsc::Sender<std::io::Result<Bytes>>,
123        Option<mpsc::Receiver<std::io::Result<Bytes>>>,
124    )> for IndexTrie
125{
126    fn from(
127        (path, tx, rx): (
128            &'a [Option<usize>],
129            mpsc::Sender<std::io::Result<Bytes>>,
130            Option<mpsc::Receiver<std::io::Result<Bytes>>>,
131        ),
132    ) -> Self {
133        match path {
134            [] => Self::Leaf { tx: Some(tx), rx },
135            [None, path @ ..] => Self::WildcardNode {
136                tx: None,
137                rx: None,
138                nested: Some(Box::new(Self::from((path, tx, rx)))),
139            },
140            [Some(i), path @ ..] => Self::IndexNode {
141                tx: None,
142                rx: None,
143                nested: {
144                    let n = i.saturating_add(1);
145                    let mut nested = Vec::with_capacity(n);
146                    nested.resize_with(n, Option::default);
147                    nested[*i] = Some(Self::from((path, tx, rx)));
148                    nested
149                },
150            },
151        }
152    }
153}
154
155impl<'a>
156    From<(
157        &'a [Option<usize>],
158        mpsc::Sender<std::io::Result<Bytes>>,
159        mpsc::Receiver<std::io::Result<Bytes>>,
160    )> for IndexTrie
161{
162    fn from(
163        (path, tx, rx): (
164            &'a [Option<usize>],
165            mpsc::Sender<std::io::Result<Bytes>>,
166            mpsc::Receiver<std::io::Result<Bytes>>,
167        ),
168    ) -> Self {
169        Self::from((path, tx, Some(rx)))
170    }
171}
172
173impl<'a> From<(&'a [Option<usize>], mpsc::Sender<std::io::Result<Bytes>>)> for IndexTrie {
174    fn from((path, tx): (&'a [Option<usize>], mpsc::Sender<std::io::Result<Bytes>>)) -> Self {
175        Self::from((path, tx, None))
176    }
177}
178
179impl<P: AsRef<[Option<usize>]>> FromIterator<P> for IndexTrie {
180    fn from_iter<T: IntoIterator<Item = P>>(iter: T) -> Self {
181        let mut root = Self::Empty;
182        for path in iter {
183            let (tx, rx) = mpsc::channel(16);
184            if !root.insert(path.as_ref(), tx, Some(rx)) {
185                return Self::Empty;
186            }
187        }
188        root
189    }
190}
191
192impl IndexTrie {
193    /// Calls `f` with the sender and receiver of each node in the trie
194    fn for_each(
195        &mut self,
196        f: &mut impl FnMut(
197            &mut Option<mpsc::Sender<std::io::Result<Bytes>>>,
198            &mut Option<mpsc::Receiver<std::io::Result<Bytes>>>,
199        ),
200    ) {
201        match self {
202            Self::Empty => {}
203            Self::Leaf { tx, rx } => f(tx, rx),
204            Self::IndexNode { tx, rx, nested } => {
205                f(tx, rx);
206                for nested in nested.iter_mut().flatten() {
207                    nested.for_each(f);
208                }
209            }
210            Self::WildcardNode { tx, rx, nested } => {
211                f(tx, rx);
212                if let Some(nested) = nested {
213                    nested.for_each(f);
214                }
215            }
216        }
217    }
218
219    /// Closes all receivers in the trie, returning `true` if any of them held data
220    #[instrument(level = "trace", skip(self), ret(level = "trace"))]
221    fn close_rx(&mut self) -> bool {
222        let mut pending = false;
223        self.for_each(&mut |_, rx| {
224            pending |= rx.take().is_some_and(|rx| !rx.is_empty());
225        });
226        pending
227    }
228
229    /// Takes the receiver
230    #[instrument(level = "trace", skip(self), ret(level = "trace"))]
231    fn take_rx(&mut self, path: &[usize]) -> Option<mpsc::Receiver<std::io::Result<Bytes>>> {
232        let Some((i, path)) = path.split_first() else {
233            return match self {
234                Self::Empty => None,
235                Self::Leaf { rx, .. } => rx.take(),
236                Self::IndexNode { tx, rx, nested } => {
237                    let rx = rx.take();
238                    if nested.is_empty() && tx.is_none() {
239                        *self = Self::Empty;
240                    }
241                    rx
242                }
243                Self::WildcardNode { tx, rx, nested } => {
244                    let rx = rx.take();
245                    if nested.is_none() && tx.is_none() {
246                        *self = Self::Empty;
247                    }
248                    rx
249                }
250            };
251        };
252        match self {
253            Self::Empty | Self::Leaf { .. } | Self::WildcardNode { .. } => None,
254            Self::IndexNode { nested, .. } => nested
255                .get_mut(*i)
256                .and_then(|nested| nested.as_mut().and_then(|nested| nested.take_rx(path))),
257            // TODO: Demux the subscription
258            //Self::WildcardNode { ref mut nested, .. } => {
259            //    nested.as_mut().and_then(|nested| nested.take(path))
260            //}
261        }
262    }
263
264    /// Gets a sender
265    #[instrument(level = "trace", skip(self), ret(level = "trace"))]
266    fn get_tx(&mut self, path: &[usize]) -> Option<mpsc::Sender<std::io::Result<Bytes>>> {
267        let Some((i, path)) = path.split_first() else {
268            return match self {
269                Self::Empty => None,
270                Self::Leaf { tx, .. } => tx.clone(),
271                Self::IndexNode { tx, .. } | Self::WildcardNode { tx, .. } => tx.clone(),
272            };
273        };
274        match self {
275            Self::Empty | Self::Leaf { .. } | Self::WildcardNode { .. } => None,
276            Self::IndexNode { nested, .. } => {
277                let nested = nested.get_mut(*i)?;
278                let nested = nested.as_mut()?;
279                nested.get_tx(path)
280            } // TODO: Demux the subscription
281              //Self::WildcardNode { ref mut nested, .. } => {
282              //    nested.as_mut().and_then(|nested| nested.take(path))
283              //}
284        }
285    }
286
287    /// Closes all senders in the trie
288    #[instrument(level = "trace", skip(self), ret(level = "trace"))]
289    fn close_tx(&mut self) {
290        self.for_each(&mut |tx, _| {
291            tx.take();
292        });
293    }
294
295    /// Inserts `sender` and `receiver` under a `path` - returns `false` if it failed and `true` if it succeeded.
296    /// Tree state after `false` is returned is undefined
297    #[instrument(level = "trace", skip(self, sender, receiver), ret(level = "trace"))]
298    fn insert(
299        &mut self,
300        path: &[Option<usize>],
301        sender: mpsc::Sender<std::io::Result<Bytes>>,
302        receiver: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
303    ) -> bool {
304        match self {
305            Self::Empty => {
306                *self = Self::from((path, sender, receiver));
307                true
308            }
309            Self::Leaf { .. } => {
310                let Some((i, path)) = path.split_first() else {
311                    return false;
312                };
313                let Self::Leaf { tx, rx } = mem::take(self) else {
314                    return false;
315                };
316                if let Some(i) = i {
317                    let n = i.saturating_add(1);
318                    let mut nested = Vec::with_capacity(n);
319                    nested.resize_with(n, Option::default);
320                    nested[*i] = Some(Self::from((path, sender, receiver)));
321                    *self = Self::IndexNode { tx, rx, nested };
322                } else {
323                    *self = Self::WildcardNode {
324                        tx,
325                        rx,
326                        nested: Some(Box::new(Self::from((path, sender, receiver)))),
327                    };
328                }
329                true
330            }
331            Self::IndexNode { tx, rx, nested } => match (&tx, &rx, path) {
332                (None, None, []) => {
333                    *tx = Some(sender);
334                    *rx = receiver;
335                    true
336                }
337                (_, _, [Some(i), path @ ..]) => {
338                    let cap = i.saturating_add(1);
339                    if nested.len() < cap {
340                        nested.resize_with(cap, Option::default);
341                    }
342                    let nested = &mut nested[*i];
343                    if let Some(nested) = nested {
344                        nested.insert(path, sender, receiver)
345                    } else {
346                        *nested = Some(Self::from((path, sender, receiver)));
347                        true
348                    }
349                }
350                _ => false,
351            },
352            Self::WildcardNode { tx, rx, nested } => match (&tx, &rx, path) {
353                (None, None, []) => {
354                    *tx = Some(sender);
355                    *rx = receiver;
356                    true
357                }
358                (_, _, [None, path @ ..]) => {
359                    if let Some(nested) = nested {
360                        nested.insert(path, sender, receiver)
361                    } else {
362                        *nested = Some(Box::new(Self::from((path, sender, receiver))));
363                        true
364                    }
365                }
366                _ => false,
367            },
368        }
369    }
370}
371
372pin_project! {
373    /// Incoming framed stream
374    #[project = IncomingProj]
375    pub struct Incoming {
376        #[pin]
377        rx: Option<StreamReader<ReceiverStream<std::io::Result<Bytes>>, Bytes>>,
378        path: Arc<[usize]>,
379        io: Arc<IngressGuard>,
380        state: Arc<IngressState>,
381    }
382
383    impl PinnedDrop for Incoming {
384        fn drop(this: Pin<&mut Self>) {
385            let this = Pin::into_inner(this);
386            if let Some(rx) = this.rx.take() {
387                let (mut rx, chunk) = rx.into_inner_with_chunk();
388                rx.close();
389                if chunk.is_some() || !rx.as_ref().is_empty() {
390                    this.mark_unprocessed();
391                }
392            }
393        }
394    }
395}
396
397impl Incoming {
398    /// Creates a new [Incoming] given an [`AsyncBufRead`], [`ConnHandler`] and a set of async paths.
399    /// `on_ingress` will be called once data ingress is complete.
400    /// If `drain` is `true`, data ingress continues after all handles are dropped until the peer
401    /// closes the stream, failing on any data received or not processed, otherwise data ingress
402    /// is aborted, see [`ConnHandler::DRAIN`].
403    pub fn new<T, P, Fut>(
404        mut rx: T,
405        paths: impl IntoIterator<Item = P>,
406        drain: bool,
407        on_ingress: impl FnOnce(T, std::io::Result<()>) -> Fut + Send + 'static,
408    ) -> Self
409    where
410        T: AsyncBufRead + Unpin + Send + 'static,
411        P: AsRef<[Option<usize>]>,
412        Fut: Future<Output = ()> + Send,
413    {
414        let (rx_tx, rx_rx) = mpsc::channel(128);
415        let state = Arc::new(IngressState {
416            index: std::sync::Mutex::new(paths.into_iter().collect()),
417            err: std::sync::OnceLock::default(),
418            unprocessed: AtomicBool::default(),
419        });
420        let (dropped_tx, dropped_rx) = drain.then(oneshot::channel).unzip();
421        let span = Span::current();
422        let task = tokio::spawn({
423            let state = Arc::clone(&state);
424            async move {
425                let res = async {
426                    let mut io = pin!(async {
427                        let res = ingress(&mut rx, &state.index, &rx_tx).await;
428                        if let Err(err) = &res {
429                            _ = state.err.set(copy_io_error(err));
430                        }
431                        drop(rx_tx);
432                        trace!("shutting down index trie");
433                        state
434                            .index
435                            .lock()
436                            .unwrap_or_else(PoisonError::into_inner)
437                            .close_tx();
438                        res
439                    });
440                    let Some(mut dropped) = dropped_rx else {
441                        return io.await;
442                    };
443                    select! {
444                        biased;
445
446                        res = &mut io => {
447                            res?;
448                            _ = dropped.await;
449                            close_rx(&state)
450                        }
451                        _ = &mut dropped => {
452                            close_rx(&state)?;
453                            io.await
454                        }
455                    }
456                }
457                .await;
458                on_ingress(rx, res).await;
459            }
460            .instrument(span.clone())
461        });
462        Self {
463            rx: Some(StreamReader::new(ReceiverStream::new(rx_rx))),
464            path: Arc::from([]),
465            io: Arc::new(match dropped_tx {
466                // TODO: Bound draining by a deadline or a byte limit
467                Some(dropped) => IngressGuard::Drain { _dropped: dropped },
468                None => IngressGuard::Abort {
469                    _task: AbortOnDropHandle::new(task),
470                },
471            }),
472            state,
473        }
474    }
475
476    /// Marks data received by this handle as not processed, which fails ingress once all
477    /// handles are dropped
478    pub(crate) fn mark_unprocessed(&self) {
479        trace!(path = ?self.path, "dropping handle with unprocessed data");
480        self.state.unprocessed.store(true, Ordering::Relaxed);
481    }
482
483    /// Index the incoming stream using a structural `path`, returning a handle to the
484    /// multiplexed sub-stream addressed by it.
485    #[instrument(level = "trace", skip(self), fields(path = ?self.path))]
486    pub fn index(&self, path: &[usize]) -> std::io::Result<Self> {
487        if path.is_empty() {
488            return Err(std::io::Error::new(
489                std::io::ErrorKind::InvalidInput,
490                "path cannot be empty",
491            ));
492        }
493        let path = if self.path.is_empty() {
494            Arc::from(path)
495        } else {
496            Arc::from([self.path.as_ref(), path].concat())
497        };
498        trace!("locking index trie");
499        let mut index = self
500            .state
501            .index
502            .lock()
503            .map_err(|err| std::io::Error::other(err.to_string()))?;
504        trace!(?path, "taking index subscription");
505        let rx = index
506            .take_rx(&path)
507            .map(|rx| StreamReader::new(ReceiverStream::new(rx)));
508        Ok(Self {
509            rx,
510            path,
511            io: Arc::clone(&self.io),
512            state: Arc::clone(&self.state),
513        })
514    }
515}
516
517impl AsyncRead for Incoming {
518    #[instrument(level = "trace", skip_all, fields(path = ?self.path), ret(level = "trace"))]
519    fn poll_read(
520        mut self: Pin<&mut Self>,
521        cx: &mut Context<'_>,
522        buf: &mut tokio::io::ReadBuf<'_>,
523    ) -> Poll<std::io::Result<()>> {
524        if buf.remaining() == 0 {
525            return Poll::Ready(Ok(()));
526        }
527        trace!("reading");
528        let this = self.as_mut().project();
529        let Some(rx) = this.rx.as_pin_mut() else {
530            trace!("reader is closed");
531            return Poll::Ready(Ok(()));
532        };
533        ready!(rx.poll_read(cx, buf))?;
534        trace!(buf = ?buf.filled(), "read buffer");
535        if buf.filled().is_empty() {
536            self.rx.take();
537            if let Some(err) = self.state.err.get() {
538                return Poll::Ready(Err(copy_io_error(err)));
539            }
540        }
541        Poll::Ready(Ok(()))
542    }
543}
544
545struct IngressState {
546    index: std::sync::Mutex<IndexTrie>,
547    err: std::sync::OnceLock<std::io::Error>,
548    unprocessed: AtomicBool,
549}
550
551enum IngressGuard {
552    Abort { _task: AbortOnDropHandle<()> },
553    Drain { _dropped: oneshot::Sender<()> },
554}
555
556/// Outgoing framed stream
557pub struct Outgoing {
558    tx: PollSender<(Bytes, Bytes)>,
559    path: Arc<[usize]>,
560    path_buf: Bytes,
561}
562
563impl Outgoing {
564    /// Creates a new [Outgoing] given an [`AsyncWrite`].
565    pub fn new<T, Fut>(
566        mut tx: T,
567        on_egress: impl FnOnce(T, std::io::Result<()>) -> Fut + Send + 'static,
568    ) -> Self
569    where
570        T: AsyncWrite + Unpin + Send + 'static,
571        Fut: Future<Output = ()> + Send,
572    {
573        let span = Span::current();
574        let (tx_tx, tx_rx) = mpsc::channel(128);
575        tokio::spawn(
576            async {
577                let res = egress(&mut tx, tx_rx).await;
578                on_egress(tx, res).await;
579            }
580            .instrument(span.clone()),
581        );
582        Self {
583            tx: PollSender::new(tx_tx),
584            path: Arc::from([]),
585            path_buf: Bytes::from_static(&[0]),
586        }
587    }
588
589    /// Index the outgoing stream using a structural `path`, returning a handle that writes
590    /// to the multiplexed sub-stream addressed by it.
591    #[instrument(level = "trace", skip(self), fields(path = ?self.path))]
592    pub fn index(&self, path: &[usize]) -> std::io::Result<Self> {
593        if path.is_empty() {
594            return Err(std::io::Error::new(
595                std::io::ErrorKind::InvalidInput,
596                "path cannot be empty",
597            ));
598        }
599        let path: Arc<[usize]> = if self.path.is_empty() {
600            Arc::from(path)
601        } else {
602            Arc::from([self.path.as_ref(), path].concat())
603        };
604        let mut buf = BytesMut::with_capacity(path.len().saturating_add(5));
605        let n = u32::try_from(path.len())
606            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
607        trace!(n, "encoding path length");
608        Leb128Encoder.encode(n, &mut buf)?;
609        for p in path.as_ref() {
610            let p = u32::try_from(*p)
611                .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
612            trace!(p, "encoding path element");
613            Leb128Encoder.encode(p, &mut buf)?;
614        }
615        Ok(Self {
616            tx: self.tx.clone(),
617            path,
618            path_buf: buf.freeze(),
619        })
620    }
621
622    pub(crate) fn poll_reserve(&mut self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
623        self.tx
624            .poll_reserve(cx)
625            .map_err(|err| std::io::Error::new(std::io::ErrorKind::BrokenPipe, err))
626    }
627
628    pub(crate) async fn reserve(&mut self) -> std::io::Result<()> {
629        poll_fn(|cx| self.poll_reserve(cx)).await
630    }
631
632    pub(crate) fn send_reserved(&mut self, buf: Bytes) -> std::io::Result<()> {
633        trace!(path = ?self.path, ?buf, "sending outgoing frame");
634        self.tx
635            .send_item((self.path_buf.clone(), buf))
636            .map_err(|err| std::io::Error::new(std::io::ErrorKind::BrokenPipe, err))
637    }
638
639    /// Writes `buf` to the stream as a single frame, without copying it.
640    pub async fn write_bytes(&mut self, buf: Bytes) -> std::io::Result<()> {
641        if buf.is_empty() {
642            return Ok(());
643        }
644        self.reserve().await?;
645        self.send_reserved(buf)
646    }
647}
648
649impl AsyncWrite for Outgoing {
650    #[instrument(level = "trace", skip_all, fields(path = ?self.path, buf = format!("{buf:02x?}")), ret(level = "trace"))]
651    fn poll_write(
652        mut self: Pin<&mut Self>,
653        cx: &mut Context<'_>,
654        buf: &[u8],
655    ) -> Poll<std::io::Result<usize>> {
656        if buf.is_empty() {
657            return Poll::Ready(Ok(0));
658        }
659        trace!("writing outgoing chunk");
660        ready!(self.poll_reserve(cx))?;
661        self.send_reserved(Bytes::copy_from_slice(buf))?;
662        Poll::Ready(Ok(buf.len()))
663    }
664
665    #[instrument(level = "trace", skip_all, fields(path = ?self.path), ret(level = "trace"))]
666    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
667        Poll::Ready(Ok(()))
668    }
669
670    #[instrument(level = "trace", skip_all, fields(path = ?self.path), ret(level = "trace"))]
671    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
672        Poll::Ready(Ok(()))
673    }
674}
675
676#[instrument(level = "trace", skip_all, ret(level = "trace"))]
677async fn ingress(
678    mut rx: impl AsyncRead + Unpin,
679    index: &std::sync::Mutex<IndexTrie>,
680    param_tx: &mpsc::Sender<std::io::Result<Bytes>>,
681) -> std::io::Result<()> {
682    loop {
683        trace!("reading path length");
684        let b = match rx.read_u8().await {
685            Ok(b) => b,
686            Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(()),
687            Err(err) => return Err(err),
688        };
689        let n = AsyncReadExt::chain([b].as_slice(), &mut rx)
690            .read_u32_leb128()
691            .await?;
692        let n = usize::try_from(n)
693            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
694        trace!(n, "read path length");
695        let tx = if n == 0 {
696            param_tx
697        } else {
698            let mut path = Vec::with_capacity(n.min(MAX_INITIAL_PATH_CAPACITY));
699            for i in 0..n {
700                trace!(i, "reading path element");
701                let p = rx.read_u32_leb128().await?;
702                let p = usize::try_from(p)
703                    .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
704                path.push(p);
705            }
706            trace!(?path, "read path");
707
708            trace!("locking index trie");
709            let mut index = index
710                .lock()
711                .map_err(|err| std::io::Error::other(err.to_string()))?;
712            &index.get_tx(&path).ok_or_else(|| {
713                std::io::Error::new(
714                    std::io::ErrorKind::NotFound,
715                    format!("`{path:?}` subscription not found"),
716                )
717            })?
718        };
719        trace!("reading data length");
720        let n = rx.read_u32_leb128().await?;
721        trace!(n, "read data length");
722        let mut data = (&mut rx).take(n.into());
723        while data.limit() > 0 {
724            let k = usize::try_from(data.limit())
725                .unwrap_or(usize::MAX)
726                .min(DEFAULT_MAX_INITIAL_CAPACITY);
727            let mut buf = BytesMut::with_capacity(k);
728            trace!(len = k, "reading data chunk");
729            while buf.len() < k {
730                if data.read_buf(&mut buf).await? == 0 {
731                    return Err(std::io::ErrorKind::UnexpectedEof.into());
732                }
733            }
734            trace!(?buf, "read data chunk");
735            tx.send(Ok(buf.freeze()))
736                .await
737                .map_err(|_| error_receiver_closed())?;
738        }
739    }
740}
741
742/// Closes all receivers in the index trie once all [Incoming] handles were dropped,
743/// failing if any of them or any of the dropped handles held data, which was not processed
744fn close_rx(state: &IngressState) -> std::io::Result<()> {
745    trace!("all handles dropped, closing index trie receivers");
746    let mut index = state.index.lock().unwrap_or_else(PoisonError::into_inner);
747    if index.close_rx() || state.unprocessed.load(Ordering::Relaxed) {
748        return Err(error_unprocessed());
749    }
750    Ok(())
751}
752
753/// Returns the error for received data, which was not processed
754fn error_unprocessed() -> std::io::Error {
755    std::io::Error::new(
756        std::io::ErrorKind::BrokenPipe,
757        "received data was not processed",
758    )
759}
760
761#[instrument(level = "trace", skip_all)]
762async fn egress(
763    mut tx: impl AsyncWrite + Unpin,
764    mut rx: mpsc::Receiver<(Bytes, Bytes)>,
765) -> std::io::Result<()> {
766    let mut buf = BytesMut::with_capacity(5);
767    trace!("waiting for next frame");
768    let mut next = rx.recv().await;
769    while let Some((path, data)) = next {
770        let data_len = u32::try_from(data.len())
771            .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?;
772        buf.clear();
773        Leb128Encoder.encode(data_len, &mut buf)?;
774        let mut frame = path.chain(&mut buf).chain(data);
775        trace!(?frame, "writing egress frame");
776        tx.write_all_buf(&mut frame).await?;
777        next = match rx.try_recv() {
778            Ok(frame) => Some(frame),
779            Err(..) => {
780                tx.flush().await?;
781                trace!("waiting for next frame");
782                rx.recv().await
783            }
784        };
785    }
786    trace!("shutting down outgoing stream");
787    tx.shutdown().await
788}
789
790/// Connection handler defines the connection I/O behavior.
791/// It is mostly useful for transports that may require additional clean up not already covered
792/// by [`AsyncWriteExt::shutdown`](tokio::io::AsyncWriteExt::shutdown), for example.
793/// This API is experimental and may change in backwards-incompatible ways in the future.
794pub trait ConnHandler<Rx, Tx> {
795    /// Whether to keep reading once all [Incoming] handles are dropped until the peer closes
796    /// the stream, rather than aborting ingress and dropping the reader.
797    /// Any data received after all handles are dropped, or received for a registered path, which
798    /// was never indexed, fails ingress with [`std::io::ErrorKind::BrokenPipe`].
799    /// Draining is required for transports that reset the connection when closed with
800    /// unread data, like TCP, discarding data not yet received by the peer.
801    const DRAIN: bool = true;
802
803    /// Handle ingress completion
804    fn on_ingress(rx: Rx, res: std::io::Result<()>) -> impl Future<Output = ()> + Send {
805        _ = rx;
806        if let Err(err) = res {
807            error!(?err, "ingress failed");
808        } else {
809            debug!("ingress successfully complete");
810        }
811        async {}
812    }
813
814    /// Handle egress completion
815    fn on_egress(tx: Tx, res: std::io::Result<()>) -> impl Future<Output = ()> + Send {
816        _ = tx;
817        if let Err(err) = res {
818            error!(?err, "egress failed");
819        } else {
820            debug!("egress successfully complete");
821        }
822        async {}
823    }
824}
825
826impl<Rx, Tx> ConnHandler<Rx, Tx> for () {}
827
828#[cfg(test)]
829mod tests {
830    use core::time::Duration;
831
832    use std::io::ErrorKind;
833
834    use tokio::io::{BufReader, DuplexStream};
835
836    use super::*;
837
838    const FRAME: [u8; 4] = [0x01, 0x00, 0x01, 0x2a];
839
840    fn incoming(
841        drain: bool,
842    ) -> (
843        DuplexStream,
844        Incoming,
845        oneshot::Receiver<std::io::Result<()>>,
846    ) {
847        let (tx, rx) = tokio::io::duplex(16);
848        let (done_tx, done_rx) = oneshot::channel();
849        let rx = Incoming::new(BufReader::new(rx), [[Some(0)]], drain, |_, res| async {
850            _ = done_tx.send(res);
851        });
852        (tx, rx, done_rx)
853    }
854
855    async fn done(rx: oneshot::Receiver<std::io::Result<()>>) -> std::io::Result<()> {
856        tokio::time::timeout(Duration::from_secs(5), rx)
857            .await
858            .expect("ingress timed out")
859            .expect("ingress aborted")
860    }
861
862    async fn failed(rx: oneshot::Receiver<std::io::Result<()>>) -> ErrorKind {
863        done(rx).await.unwrap_err().kind()
864    }
865
866    #[test_log::test(tokio::test)]
867    async fn incoming_drop_aborts() {
868        let (_tx, rx, done_rx) = incoming(false);
869        drop(rx);
870        tokio::time::timeout(Duration::from_secs(5), done_rx)
871            .await
872            .unwrap()
873            .unwrap_err();
874    }
875
876    #[test_log::test(tokio::test)]
877    async fn incoming_drop_waits_for_eof() {
878        let (tx, rx, mut done_rx) = incoming(true);
879        drop(rx);
880        tokio::task::yield_now().await;
881        assert!(done_rx.try_recv().is_err());
882        drop(tx);
883        done(done_rx).await.unwrap();
884    }
885
886    #[test_log::test(tokio::test)]
887    async fn incoming_drop_accepts_empty_frame() {
888        let (mut tx, rx, done_rx) = incoming(true);
889        drop(rx);
890        tx.write_all(&[0x00, 0x00]).await.unwrap();
891        drop(tx);
892        done(done_rx).await.unwrap();
893    }
894
895    #[test_log::test(tokio::test)]
896    async fn incoming_drop_rejects_data() {
897        let (mut tx, rx, done_rx) = incoming(true);
898        drop(rx);
899        tx.write_all(&FRAME).await.unwrap();
900        assert_eq!(failed(done_rx).await, ErrorKind::BrokenPipe);
901    }
902
903    #[test_log::test(tokio::test)]
904    async fn incoming_drop_reports_ingress_error() {
905        let (mut tx, rx, done_rx) = incoming(true);
906        tx.write_all(&FRAME).await.unwrap();
907        tx.write_all(&[0x01, 0x01, 0x01, 0x2a]).await.unwrap();
908        drop(rx);
909        assert_eq!(failed(done_rx).await, ErrorKind::NotFound);
910    }
911
912    #[test_log::test(tokio::test)]
913    async fn incoming_drop_rejects_pending() {
914        let (mut tx, rx, done_rx) = incoming(true);
915        tx.write_all(&FRAME).await.unwrap();
916        tokio::task::yield_now().await;
917        drop(rx);
918        assert_eq!(failed(done_rx).await, ErrorKind::BrokenPipe);
919    }
920
921    #[test_log::test(tokio::test)]
922    async fn incoming_drop_rejects_unread() {
923        let (mut tx, rx, done_rx) = incoming(true);
924        let sub = rx.index(&[0]).unwrap();
925        tx.write_all(&FRAME).await.unwrap();
926        tokio::task::yield_now().await;
927        drop(sub);
928        drop(rx);
929        assert_eq!(failed(done_rx).await, ErrorKind::BrokenPipe);
930    }
931
932    #[test_log::test(tokio::test)]
933    async fn buffered_drop_rejects_unread() {
934        let (tx, rx, done_rx) = incoming(true);
935        drop(crate::BufferedIncoming {
936            buffer: BytesMut::from(&FRAME[..]),
937            inner: rx,
938        });
939        drop(tx);
940        assert_eq!(failed(done_rx).await, ErrorKind::BrokenPipe);
941    }
942
943    #[test_log::test(tokio::test)]
944    async fn incoming_eof_waits_for_drop() {
945        let (tx, mut rx, mut done_rx) = incoming(true);
946        drop(tx);
947        rx.read_to_end(&mut Vec::default()).await.unwrap();
948        assert!(done_rx.try_recv().is_err());
949        drop(rx);
950        done(done_rx).await.unwrap();
951    }
952
953    #[test_log::test(tokio::test)]
954    async fn incoming_eof_rejects_pending() {
955        let (mut tx, rx, done_rx) = incoming(true);
956        tx.write_all(&FRAME).await.unwrap();
957        drop(tx);
958        tokio::task::yield_now().await;
959        drop(rx);
960        assert_eq!(failed(done_rx).await, ErrorKind::BrokenPipe);
961    }
962
963    #[test_log::test(tokio::test)]
964    async fn ingress_truncated_path() {
965        let index = std::sync::Mutex::new(IndexTrie::Empty);
966        let (tx, _rx) = mpsc::channel(1);
967        let err = ingress([0x02, 0x00].as_slice(), &index, &tx)
968            .await
969            .unwrap_err();
970        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
971    }
972
973    #[test_log::test(tokio::test)]
974    async fn ingress_chunked_data() {
975        let index = std::sync::Mutex::new(IndexTrie::Empty);
976        let (tx, mut rx) = mpsc::channel(4);
977        let n = DEFAULT_MAX_INITIAL_CAPACITY + 1;
978        let mut frame = vec![0x00, 0x81, 0x80, 0x40];
979        frame.resize(frame.len() + n, 0x42);
980        ingress(frame.as_slice(), &index, &tx).await.unwrap();
981        drop(tx);
982        let mut chunks = vec![];
983        while let Some(chunk) = rx.recv().await {
984            chunks.push(chunk.unwrap().len());
985        }
986        assert_eq!(chunks, [DEFAULT_MAX_INITIAL_CAPACITY, 1]);
987    }
988
989    #[test_log::test(tokio::test)]
990    async fn ingress_truncated_data() {
991        let index = std::sync::Mutex::new(IndexTrie::Empty);
992        let (tx, _rx) = mpsc::channel(1);
993        let err = ingress(
994            [0x00, 0xff, 0xff, 0xff, 0xff, 0x0f, 0x42].as_slice(),
995            &index,
996            &tx,
997        )
998        .await
999        .unwrap_err();
1000        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
1001    }
1002}
wrpc_transport/frame/conn/mod.rs

wrpc_transport/frame/conn/
mod.rs
