unb_runtime/
transport.rs

1use unb_core::Envelope;
2
3use crate::BodyStream;
4use bytes::Bytes;
5use futures_util::stream::{SplitSink, SplitStream};
6use futures_util::{SinkExt, StreamExt};
7use tokio::sync::mpsc;
8use unb_transport::BoxTransport;
9
10use crate::error::WsError;
11
12pub type SessionStreams = Option<std::sync::Arc<unb_transport::ConnectionStreams>>;
13
14pub enum Pipe {
15    Local {
16        rx: mpsc::Receiver<Envelope>,
17        tx: mpsc::Sender<Envelope>,
18        initiator: bool,
19    },
20    Piped {
21        pipe: BoxTransport,
22        initiator: bool,
23    },
24    PipedWithStreams {
25        pipe: BoxTransport,
26        initiator: bool,
27        streams: SessionStreams,
28    },
29}
30
31impl Pipe {
32    pub fn piped_with_streams(
33        pipe: BoxTransport,
34        initiator: bool,
35        streams: std::sync::Arc<unb_transport::ConnectionStreams>,
36    ) -> Pipe {
37        Pipe::PipedWithStreams {
38            pipe,
39            initiator,
40            streams: Some(streams),
41        }
42    }
43
44    pub(crate) fn initiator(&self) -> bool {
45        match self {
46            Pipe::Local { initiator, .. }
47            | Pipe::Piped { initiator, .. }
48            | Pipe::PipedWithStreams { initiator, .. } => *initiator,
49        }
50    }
51
52    pub(crate) fn streams(&self) -> SessionStreams {
53        match self {
54            Pipe::Local { .. } | Pipe::Piped { .. } => None,
55            Pipe::PipedWithStreams { streams, .. } => streams.clone(),
56        }
57    }
58
59    pub(crate) fn split(self) -> (PipeReader, PipeWriter) {
60        match self {
61            Pipe::Local { rx, tx, .. } => (PipeReader::Local(rx), PipeWriter::Local(tx)),
62            Pipe::Piped { pipe, .. } | Pipe::PipedWithStreams { pipe, .. } => {
63                let (sink, stream) = pipe.split();
64                (PipeReader::Piped(stream), PipeWriter::Piped(sink))
65            }
66        }
67    }
68}
69
70pub(crate) enum PipeReader {
71    Local(mpsc::Receiver<Envelope>),
72    Piped(SplitStream<BoxTransport>),
73}
74
75pub(crate) enum PipeWriter {
76    Local(mpsc::Sender<Envelope>),
77    Piped(SplitSink<BoxTransport, Bytes>),
78}
79
80impl PipeReader {
81    pub(crate) async fn recv(&mut self) -> Result<Option<(Envelope, Option<BodyStream>)>, WsError> {
82        match self {
83            PipeReader::Local(rx) => Ok(rx.recv().await.map(|envelope| (envelope, None))),
84            PipeReader::Piped(stream) => match stream.next().await {
85                None => Ok(None),
86                Some(Ok(frame)) => Ok(Some((Envelope::decode(frame)?, None))),
87                Some(Err(error)) => Err(error.into()),
88            },
89        }
90    }
91}
92
93impl PipeWriter {
94    pub(crate) async fn feed(&mut self, envelope: Envelope) -> Result<(), WsError> {
95        match self {
96            PipeWriter::Local(tx) => tx.send(envelope).await.map_err(|_| WsError::Gone),
97            PipeWriter::Piped(sink) => Ok(sink.feed(envelope.encode()).await?),
98        }
99    }
100
101    pub(crate) async fn flush(&mut self) -> Result<(), WsError> {
102        match self {
103            PipeWriter::Local(_) => Ok(()),
104            PipeWriter::Piped(sink) => Ok(sink.flush().await?),
105        }
106    }
107
108    pub(crate) async fn close(&mut self) -> Result<(), WsError> {
109        match self {
110            PipeWriter::Local(_) => Ok(()),
111            PipeWriter::Piped(sink) => Ok(sink.close().await?),
112        }
113    }
114}
unb_runtime/transport.rs

unb_runtime/
transport.rs
