1use futures::channel::mpsc;
6use futures::sink::Sink;
7use futures::stream::Stream;
8use futures::{AsyncRead, AsyncWrite};
9use std::io::{self, Cursor, Read};
10use std::pin::Pin;
11use std::task::{Context, Poll};
12
13#[derive(Debug)]
15pub struct MpscTransport {
16 sender: mpsc::Sender<Vec<u8>>,
17 receiver: mpsc::Receiver<Vec<u8>>,
18 current_read_buf: Option<Cursor<Vec<u8>>>,
19}
20
21impl MpscTransport {
22 pub fn pair() -> (Self, Self) {
24 Self::pair_with_capacity(128)
25 }
26
27 pub fn pair_with_capacity(capacity: usize) -> (Self, Self) {
29 let (tx_a, rx_b) = mpsc::channel(capacity);
30 let (tx_b, rx_a) = mpsc::channel(capacity);
31 (
32 Self { sender: tx_a, receiver: rx_a, current_read_buf: None },
33 Self { sender: tx_b, receiver: rx_b, current_read_buf: None },
34 )
35 }
36
37 pub fn new(sender: mpsc::Sender<Vec<u8>>, receiver: mpsc::Receiver<Vec<u8>>) -> Self {
39 Self { sender, receiver, current_read_buf: None }
40 }
41}
42
43impl AsyncRead for MpscTransport {
44 fn poll_read(
45 mut self: Pin<&mut Self>,
46 cx: &mut Context<'_>,
47 buf: &mut [u8],
48 ) -> Poll<io::Result<usize>> {
49 if buf.is_empty() {
50 return Poll::Ready(Ok(0));
51 }
52
53 let this = &mut *self;
54 loop {
55 if let Some(cursor) = &mut this.current_read_buf {
56 let n = match cursor.read(buf) {
57 Ok(n) => n,
58 Err(e) => return Poll::Ready(Err(e)),
59 };
60 if n > 0 {
61 if cursor.position() >= cursor.get_ref().len() as u64 {
62 this.current_read_buf = None;
63 }
64 return Poll::Ready(Ok(n));
65 }
66 this.current_read_buf = None;
67 }
68
69 match Pin::new(&mut this.receiver).poll_next(cx) {
70 Poll::Ready(Some(chunk)) => {
71 if chunk.is_empty() {
72 continue;
73 }
74 this.current_read_buf = Some(Cursor::new(chunk));
75 }
76 Poll::Ready(None) => {
77 return Poll::Ready(Ok(0));
78 }
79 Poll::Pending => {
80 return Poll::Pending;
81 }
82 }
83 }
84 }
85}
86
87impl AsyncWrite for MpscTransport {
88 fn poll_write(
89 mut self: Pin<&mut Self>,
90 cx: &mut Context<'_>,
91 buf: &[u8],
92 ) -> Poll<io::Result<usize>> {
93 if buf.is_empty() {
94 return Poll::Ready(Ok(0));
95 }
96
97 let this = &mut *self;
98 match Pin::new(&mut this.sender).poll_ready(cx) {
99 Poll::Ready(Ok(())) => {
100 let len = buf.len();
101 match Pin::new(&mut this.sender).start_send(buf.to_vec()) {
102 Ok(()) => Poll::Ready(Ok(len)),
103 Err(e) if e.is_disconnected() => {
104 Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
105 }
106 Err(e) => Poll::Ready(Err(io::Error::other(e))),
107 }
108 }
109 Poll::Ready(Err(e)) => {
110 if e.is_disconnected() {
111 Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
112 } else {
113 Poll::Ready(Err(io::Error::other(e)))
114 }
115 }
116 Poll::Pending => Poll::Pending,
117 }
118 }
119
120 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
121 let this = &mut *self;
122 match Pin::new(&mut this.sender).poll_flush(cx) {
123 Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
124 Poll::Ready(Err(e)) => {
125 if e.is_disconnected() {
126 Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
127 } else {
128 Poll::Ready(Err(io::Error::other(e)))
129 }
130 }
131 Poll::Pending => Poll::Pending,
132 }
133 }
134
135 fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
136 let this = &mut *self;
137 match Pin::new(&mut this.sender).poll_close(cx) {
138 Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
139 Poll::Ready(Err(e)) => {
140 if e.is_disconnected() {
141 Poll::Ready(Ok(()))
142 } else {
143 Poll::Ready(Err(io::Error::other(e)))
144 }
145 }
146 Poll::Pending => Poll::Pending,
147 }
148 }
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154 use futures::{AsyncReadExt, AsyncWriteExt};
155
156 #[fuchsia::test]
157 async fn test_basic_transfer() {
158 let (mut a, mut b) = MpscTransport::pair();
159 a.write_all(b"hello world").await.unwrap();
160
161 let mut buf = [0u8; 11];
162 b.read_exact(&mut buf).await.unwrap();
163 assert_eq!(&buf, b"hello world");
164 }
165
166 #[fuchsia::test]
167 async fn test_fragmented_read() {
168 let (mut a, mut b) = MpscTransport::pair();
169 a.write_all(b"1234567890").await.unwrap();
170
171 let mut buf1 = [0u8; 4];
172 let mut buf2 = [0u8; 6];
173 b.read_exact(&mut buf1).await.unwrap();
174 b.read_exact(&mut buf2).await.unwrap();
175 assert_eq!(&buf1, b"1234");
176 assert_eq!(&buf2, b"567890");
177 }
178
179 #[fuchsia::test]
180 async fn test_eof_on_drop() {
181 let (a, mut b) = MpscTransport::pair();
182 drop(a);
183
184 let mut buf = [0u8; 10];
185 let n = b.read(&mut buf).await.unwrap();
186 assert_eq!(n, 0);
187 }
188
189 #[fuchsia::test]
190 async fn test_broken_pipe_on_write() {
191 let (mut a, b) = MpscTransport::pair();
192 drop(b);
193
194 let err = a.write_all(b"fail").await.unwrap_err();
195 assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);
196 }
197}