use std::sync::Arc; use std::task::{Context, Poll}; use atomic_waker::AtomicWaker; use crossbeam_queue::{ArrayQueue, SegQueue}; pub struct HybridChannel { ring: ArrayQueue, overflow: SegQueue, } impl HybridChannel { pub fn new(capacity: usize) -> Self { Self { ring: ArrayQueue::new(capacity), overflow: SegQueue::new(), } } pub fn push(&self, value: T) { if !self.overflow.is_empty() { self.overflow.push(value); return; } if let Err(v) = self.ring.push(value) { self.overflow.push(v); } } pub fn pop(&self) -> Option { self.ring.pop().or_else(|| self.overflow.pop()) } pub fn is_empty(&self) -> bool { self.ring.is_empty() && self.overflow.is_empty() } } struct AsyncInboxQueue { queue: HybridChannel, waker: AtomicWaker, } impl AsyncInboxQueue { fn new(capacity: usize) -> Self { Self { queue: HybridChannel::new(capacity), waker: AtomicWaker::new(), } } fn push(&self, value: T) { self.queue.push(value); self.waker.wake(); } fn try_pop(&self) -> Option { self.queue.pop() } fn poll_pop(&self, cx: &mut Context<'_>) -> Poll { if let Some(value) = self.try_pop() { return Poll::Ready(value); } self.waker.register(cx.waker()); match self.try_pop() { Some(value) => Poll::Ready(value), None => Poll::Pending, } } } /// Receiver used only by process-external inboxes. Actor and worker channels /// retain the non-waking [`Receiver`] fast path below. pub(crate) struct AsyncReceiver { queue: Arc>, } impl AsyncReceiver { pub fn new(capacity: usize) -> Self { Self { queue: Arc::new(AsyncInboxQueue::new(capacity)), } } pub fn try_recv(&self) -> Option { self.queue.try_pop() } pub fn poll_recv(&self, cx: &mut Context<'_>) -> Poll { self.queue.poll_pop(cx) } pub fn new_sender(&self) -> AsyncSender { AsyncSender { queue: self.queue.clone(), } } } pub(crate) struct AsyncSender { queue: Arc>, } impl AsyncSender { pub fn send(&self, value: T) { self.queue.push(value); } } pub(crate) struct Receiver { queue: Arc>, } impl Receiver { pub fn new(capacity: usize) -> Self { let queue = Arc::new(HybridChannel::new(capacity)); Self { queue } } pub fn try_recv(&self) -> Option { self.queue.pop() } pub fn is_empty(&self) -> bool { self.queue.is_empty() } pub fn new_sender(&self) -> Sender { Sender { queue: self.queue.clone(), } } } pub(crate) struct Sender { queue: Arc>, } impl Clone for Sender { fn clone(&self) -> Self { Self { queue: self.queue.clone(), } } } impl Sender { pub fn send(&self, value: T) { self.queue.push(value) } }