81 lines
2.2 KiB
Rust
81 lines
2.2 KiB
Rust
use std::future;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::mpsc::{self, RecvError, SendError, TryRecvError};
|
|
use std::sync::Arc;
|
|
use std::task::Poll;
|
|
|
|
use super::AtomicWaker;
|
|
|
|
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
|
|
let (sender, receiver) = mpsc::channel();
|
|
let shared = Arc::new(Shared { closed: AtomicBool::new(false), waker: AtomicWaker::new() });
|
|
|
|
let sender = Sender { sender, shared: Arc::clone(&shared) };
|
|
let receiver = Receiver { receiver, shared };
|
|
|
|
(sender, receiver)
|
|
}
|
|
|
|
pub struct Sender<T> {
|
|
sender: mpsc::Sender<T>,
|
|
shared: Arc<Shared>,
|
|
}
|
|
|
|
impl<T> Sender<T> {
|
|
pub fn send(&self, event: T) -> Result<(), SendError<T>> {
|
|
self.sender.send(event)?;
|
|
self.shared.waker.wake();
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl<T> Drop for Sender<T> {
|
|
fn drop(&mut self) {
|
|
self.shared.closed.store(true, Ordering::Relaxed);
|
|
self.shared.waker.wake();
|
|
}
|
|
}
|
|
|
|
pub struct Receiver<T> {
|
|
receiver: mpsc::Receiver<T>,
|
|
shared: Arc<Shared>,
|
|
}
|
|
|
|
impl<T> Receiver<T> {
|
|
pub async fn next(&self) -> Result<T, RecvError> {
|
|
future::poll_fn(|cx| match self.receiver.try_recv() {
|
|
Ok(event) => Poll::Ready(Ok(event)),
|
|
Err(TryRecvError::Empty) => {
|
|
self.shared.waker.register(cx.waker());
|
|
|
|
match self.receiver.try_recv() {
|
|
Ok(event) => Poll::Ready(Ok(event)),
|
|
Err(TryRecvError::Empty) => {
|
|
if self.shared.closed.load(Ordering::Relaxed) {
|
|
Poll::Ready(Err(RecvError))
|
|
} else {
|
|
Poll::Pending
|
|
}
|
|
},
|
|
Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError)),
|
|
}
|
|
},
|
|
Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError)),
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub fn try_recv(&self) -> Result<Option<T>, RecvError> {
|
|
match self.receiver.try_recv() {
|
|
Ok(value) => Ok(Some(value)),
|
|
Err(TryRecvError::Empty) => Ok(None),
|
|
Err(TryRecvError::Disconnected) => Err(RecvError),
|
|
}
|
|
}
|
|
}
|
|
|
|
struct Shared {
|
|
closed: AtomicBool,
|
|
waker: AtomicWaker,
|
|
}
|