use crate::message::Message;
use parking_lot::Mutex;
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::Notify;
struct MailboxInner {
queue: Mutex<VecDeque<Arc<Message>>>,
capacity: usize,
notify: Notify,
closed: Mutex<bool>,
}
impl std::fmt::Debug for MailboxInner {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MailboxInner")
.field("capacity", &self.capacity)
.field("len", &self.queue.lock().len())
.field("closed", &*self.closed.lock())
.finish_non_exhaustive()
}
}
impl MailboxInner {
fn new(capacity: usize) -> Self {
Self {
queue: Mutex::new(VecDeque::with_capacity(capacity)),
capacity,
notify: Notify::new(),
closed: Mutex::new(false),
}
}
}
pub fn mailbox(capacity: usize) -> (MailboxSender, MailboxReceiver) {
let inner = Arc::new(MailboxInner::new(capacity));
(
MailboxSender {
inner: Some(inner.clone()),
},
MailboxReceiver { inner },
)
}
#[derive(Clone, Debug)]
pub struct MailboxSender {
inner: Option<Arc<MailboxInner>>,
}
impl MailboxSender {
pub fn noop() -> Self {
Self { inner: None }
}
#[allow(clippy::result_unit_err)]
pub fn send(&self, msg: Arc<Message>) -> Result<(), ()> {
let inner = match &self.inner {
None => return Ok(()), Some(inner) => inner,
};
if *inner.closed.lock() {
return Err(());
}
let mut queue = inner.queue.lock();
if queue.len() >= inner.capacity {
return Err(()); }
queue.push_back(msg);
drop(queue);
inner.notify.notify_one();
Ok(())
}
}
pub struct MailboxReceiver {
inner: Arc<MailboxInner>,
}
impl MailboxReceiver {
pub fn try_recv(&mut self) -> Option<Arc<Message>> {
self.inner.queue.lock().pop_front()
}
pub async fn recv(&mut self) -> Option<Arc<Message>> {
if let Some(msg) = self.try_recv() {
return Some(msg);
}
self.inner.notify.notified().await;
self.try_recv()
}
pub fn try_recv_batch(&mut self, buf: &mut Vec<Arc<Message>>, max: usize) -> usize {
let mut queue = self.inner.queue.lock();
let count = queue.len().min(max);
for _ in 0..count {
buf.push(queue.pop_front().unwrap());
}
count
}
pub async fn recv_batch(&mut self, buf: &mut Vec<Arc<Message>>, max: usize) -> usize {
loop {
let count = self.try_recv_batch(buf, max);
if count > 0 {
return count;
}
self.inner.notify.notified().await;
}
}
pub fn close(&self) {
*self.inner.closed.lock() = true;
self.inner.notify.notify_one();
}
}
impl Drop for MailboxReceiver {
fn drop(&mut self) {
self.close();
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::actor::Addr;
use crate::message::Message;
fn make_hi() -> Arc<Message> {
Arc::new(Message::Hi {
from: Addr::noop(),
peer_id: "test".to_string(),
})
}
#[test]
fn test_send_recv_single() {
let (tx, mut rx) = mailbox(16);
tx.send(make_hi()).unwrap();
let msg = rx.try_recv().unwrap();
assert!(matches!(&*msg, Message::Hi { .. }));
}
#[test]
fn test_fifo_order() {
let (tx, mut rx) = mailbox(1024);
for i in 0..100 {
let msg = Arc::new(Message::Hi {
from: Addr::noop(),
peer_id: format!("peer_{i}"),
});
tx.send(msg).unwrap();
}
let mut batch = Vec::with_capacity(128);
let n = rx.try_recv_batch(&mut batch, 128);
assert_eq!(n, 100);
for (i, msg) in batch.iter().enumerate() {
match msg.as_ref() {
Message::Hi { peer_id, .. } => assert_eq!(peer_id, &format!("peer_{i}")),
_ => panic!("expected Hi message"),
}
}
}
#[test]
fn test_capacity_bound() {
let (tx, _rx) = mailbox(2);
tx.send(make_hi()).unwrap();
tx.send(make_hi()).unwrap();
assert!(tx.send(make_hi()).is_err());
}
#[test]
fn test_batch_drain_partial() {
let (tx, mut rx) = mailbox(1024);
for _ in 0..50 {
tx.send(make_hi()).unwrap();
}
let mut batch = Vec::with_capacity(64);
let n = rx.try_recv_batch(&mut batch, 32);
assert_eq!(n, 32);
assert_eq!(batch.len(), 32);
let n2 = rx.try_recv_batch(&mut batch, 32);
assert_eq!(n2, 18);
assert_eq!(batch.len(), 50);
}
#[test]
fn test_noop_sender() {
let tx = MailboxSender::noop();
assert!(tx.send(make_hi()).is_ok());
assert!(tx.send(make_hi()).is_ok());
}
#[test]
fn test_close() {
let (tx, rx) = mailbox(16);
tx.send(make_hi()).unwrap();
rx.close();
assert!(tx.send(make_hi()).is_err());
}
#[test]
fn test_drop_receiver_closes() {
let (tx, rx) = mailbox(16);
tx.send(make_hi()).unwrap();
drop(rx);
assert!(tx.send(make_hi()).is_err());
}
#[tokio::test]
async fn test_notify_wake() {
let (tx, mut rx) = mailbox(1024);
let mut batch = Vec::with_capacity(64);
let consumer = tokio::spawn(async move {
let n = rx.recv_batch(&mut batch, 64).await;
assert!(n > 0);
assert!(matches!(&*batch[0], Message::Hi { .. }));
});
crate::tokio_time::sleep(web_time::Duration::from_millis(50)).await;
tx.send(make_hi()).unwrap();
consumer.await.unwrap();
}
#[test]
fn test_clone_sender() {
let (tx, mut rx) = mailbox(16);
let tx2 = tx.clone();
tx.send(make_hi()).unwrap();
tx2.send(make_hi()).unwrap();
let mut batch = Vec::with_capacity(16);
let n = rx.try_recv_batch(&mut batch, 16);
assert_eq!(n, 2);
}
#[test]
fn test_empty_recv_returns_none() {
let (_tx, mut rx) = mailbox(16);
assert!(rx.try_recv().is_none());
let mut batch = Vec::with_capacity(16);
assert_eq!(rx.try_recv_batch(&mut batch, 16), 0);
}
}