98 lines
3.2 KiB
Rust
98 lines
3.2 KiB
Rust
use std::future::Future;
|
|
use std::io::{self, ErrorKind};
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
use std::time::Duration;
|
|
|
|
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
|
|
use tokio::sync::mpsc::{channel, Receiver, Sender};
|
|
use tokio::time::sleep;
|
|
|
|
#[tokio::main]
|
|
async fn main() {
|
|
let (sender, receiver) = channel::<StreamMessage>(512);
|
|
let stream = Stream { sender, receiver };
|
|
let (mut s, mut w) = tokio::io::split(stream);
|
|
|
|
let t = tokio::spawn(async move {
|
|
let mut buf = Vec::<u8>::with_capacity(1024);
|
|
loop {
|
|
match s.read_buf(&mut buf).await {
|
|
Ok(len) => println!("Len: {}, Buf: {:?}", len, &buf[..len]),
|
|
Err(e) => {
|
|
println!("Err: {:?}", e);
|
|
if e.to_string().contains("broken pipe") {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
sleep(Duration::from_secs(1)).await;
|
|
let _ = w.write_all(b"hello world").await;
|
|
sleep(Duration::from_secs(1)).await;
|
|
let _ = w.write_all(b"hello world").await;
|
|
sleep(Duration::from_secs(1)).await;
|
|
let _ = w.write_all(b"hello world").await;
|
|
let a = tokio::join!(w.shutdown());
|
|
let b = tokio::join!(t);
|
|
println!("{:?}, {:?}", a, b)
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
enum StreamMessage {
|
|
Data(Vec<u8>),
|
|
Flush,
|
|
Close,
|
|
}
|
|
|
|
struct Stream {
|
|
sender: Sender<StreamMessage>,
|
|
receiver: Receiver<StreamMessage>,
|
|
}
|
|
|
|
impl AsyncRead for Stream {
|
|
fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
|
|
let m = self.receiver.poll_recv(cx);
|
|
match m {
|
|
Poll::Pending => Poll::Pending,
|
|
Poll::Ready(m_opt) => match m_opt {
|
|
None => Poll::Ready(Err(io::Error::new(ErrorKind::BrokenPipe, "broken pipe"))),
|
|
Some(StreamMessage::Close) => {
|
|
self.receiver.close();
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
Some(StreamMessage::Flush) => Poll::Ready(Ok(())),
|
|
Some(StreamMessage::Data(mm)) => {
|
|
buf.put_slice(mm.as_slice());
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl AsyncWrite for Stream {
|
|
fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
|
|
pool_send(self, cx, StreamMessage::Data(Vec::from(buf)), buf.len())
|
|
}
|
|
|
|
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
pool_send(self, cx, StreamMessage::Flush, ())
|
|
}
|
|
|
|
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
pool_send(self, cx, StreamMessage::Close, ())
|
|
}
|
|
}
|
|
|
|
fn pool_send<R>(s: Pin<&mut Stream>, cx: &mut Context<'_>, stream_message: StreamMessage, r: R) -> Poll<io::Result<R>> {
|
|
let fut = s.sender.send(stream_message);
|
|
tokio::pin!( fut);
|
|
match fut.poll(cx) {
|
|
Poll::Pending => Poll::Pending,
|
|
Poll::Ready(Ok(_)) => Poll::Ready(Ok(r)),
|
|
Poll::Ready(Err(e)) => Poll::Ready(Err(io::Error::new(ErrorKind::BrokenPipe, e.to_string()))),
|
|
}
|
|
} |