mirror of https://github.com/fafhrd91/actix-net
55 lines
1.2 KiB
Rust
55 lines
1.2 KiB
Rust
use std::{
|
|
sync::Arc,
|
|
task::{Context, Poll},
|
|
};
|
|
|
|
use futures_util::task::AtomicWaker;
|
|
use rtrb::{Consumer, Producer, PushError, RingBuffer};
|
|
|
|
pub fn channel<T>(cap: usize) -> (Sender<T>, Receiver<T>) {
|
|
let (tx, rx) = RingBuffer::new(cap).split();
|
|
let waker = Arc::new(AtomicWaker::new());
|
|
let sender = Sender {
|
|
tx,
|
|
waker: waker.clone(),
|
|
};
|
|
|
|
let receiver = Receiver { rx, waker };
|
|
|
|
(sender, receiver)
|
|
}
|
|
|
|
pub struct Sender<T> {
|
|
tx: Producer<T>,
|
|
waker: Arc<AtomicWaker>,
|
|
}
|
|
|
|
impl<T> Sender<T> {
|
|
pub fn send(&mut self, item: T) -> Result<(), T> {
|
|
match self.tx.push(item) {
|
|
Ok(_) => {
|
|
self.waker.wake();
|
|
Ok(())
|
|
}
|
|
Err(PushError::Full(item)) => Err(item),
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct Receiver<T> {
|
|
rx: Consumer<T>,
|
|
waker: Arc<AtomicWaker>,
|
|
}
|
|
|
|
impl<T> Receiver<T> {
|
|
pub fn poll_recv_unpin(&mut self, cx: &mut Context<'_>) -> Poll<T> {
|
|
match self.rx.pop() {
|
|
Ok(item) => Poll::Ready(item),
|
|
Err(_) => {
|
|
self.waker.register(cx.waker());
|
|
Poll::Pending
|
|
}
|
|
}
|
|
}
|
|
}
|