Skip to main content

runsync_transfer/transport/
mem.rs

1//! In-process transport over `tokio::io::duplex`.
2//!
3//! Exists so the protocol, codec, and I/O pipeline can be tested at full speed
4//! without a socket, a certificate, or a port. The duplex buffers are bounded,
5//! so backpressure behaves the same way it does over a real stream — which is
6//! the property the large-file tests actually depend on.
7
8use super::{BoxRecv, BoxSend, Transport};
9use crate::error::{Error, Result};
10use tokio::io::DuplexStream;
11use tokio::sync::{mpsc, Mutex};
12
13type BiPair = (DuplexStream, DuplexStream);
14
15pub struct MemTransport {
16    label: String,
17    buf: usize,
18    uni_out: mpsc::UnboundedSender<DuplexStream>,
19    uni_in: Mutex<mpsc::UnboundedReceiver<DuplexStream>>,
20    bi_out: mpsc::UnboundedSender<BiPair>,
21    bi_in: Mutex<mpsc::UnboundedReceiver<BiPair>>,
22}
23
24/// Create two connected endpoints. `buf` is the per-stream buffer in bytes and
25/// sets how much data may sit in flight before the writer blocks.
26pub fn pair(buf: usize) -> (MemTransport, MemTransport) {
27    let (a_uni_tx, a_uni_rx) = mpsc::unbounded_channel();
28    let (b_uni_tx, b_uni_rx) = mpsc::unbounded_channel();
29    let (a_bi_tx, a_bi_rx) = mpsc::unbounded_channel();
30    let (b_bi_tx, b_bi_rx) = mpsc::unbounded_channel();
31
32    let a = MemTransport {
33        label: "mem:a".into(),
34        buf,
35        uni_out: b_uni_tx,
36        uni_in: Mutex::new(a_uni_rx),
37        bi_out: b_bi_tx,
38        bi_in: Mutex::new(a_bi_rx),
39    };
40    let b = MemTransport {
41        label: "mem:b".into(),
42        buf,
43        uni_out: a_uni_tx,
44        uni_in: Mutex::new(b_uni_rx),
45        bi_out: a_bi_tx,
46        bi_in: Mutex::new(b_bi_rx),
47    };
48    (a, b)
49}
50
51#[async_trait::async_trait]
52impl Transport for MemTransport {
53    async fn open_uni(&self) -> Result<BoxSend> {
54        let (mine, theirs) = tokio::io::duplex(self.buf);
55        self.uni_out
56            .send(theirs)
57            .map_err(|_| Error::Closed("peer endpoint dropped".into()))?;
58        Ok(Box::new(mine))
59    }
60
61    async fn accept_uni(&self) -> Result<BoxRecv> {
62        let mut rx = self.uni_in.lock().await;
63        match rx.recv().await {
64            Some(s) => Ok(Box::new(s)),
65            None => Err(Error::Closed("peer endpoint dropped".into())),
66        }
67    }
68
69    async fn open_bi(&self) -> Result<(BoxSend, BoxRecv)> {
70        // Two independent duplexes so each direction has its own buffer and a
71        // stalled reader cannot deadlock the writer in the other direction.
72        let (my_w, their_r) = tokio::io::duplex(self.buf);
73        let (their_w, my_r) = tokio::io::duplex(self.buf);
74        self.bi_out
75            .send((their_w, their_r))
76            .map_err(|_| Error::Closed("peer endpoint dropped".into()))?;
77        Ok((Box::new(my_w), Box::new(my_r)))
78    }
79
80    async fn accept_bi(&self) -> Result<(BoxSend, BoxRecv)> {
81        let mut rx = self.bi_in.lock().await;
82        match rx.recv().await {
83            Some((w, r)) => Ok((Box::new(w), Box::new(r))),
84            None => Err(Error::Closed("peer endpoint dropped".into())),
85        }
86    }
87
88    fn close(&self, _code: u32, _reason: &[u8]) {}
89
90    fn peer_label(&self) -> String {
91        self.label.clone()
92    }
93}