use super::{request::RequestOutput, sleep::TimerId, waiter::WaiterId, waker::Waker};
use crate::Error;
use commonware_utils::channel::oneshot;
use std::mem;
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
use loom::sync::{Mutex, MutexGuard};
} else {
use commonware_utils::sync::{Mutex, MutexGuard};
}
}
pub enum Message {
WakeRoot,
Forward(Forward),
Cancel(Cancel),
}
pub enum Forward {
Waiter(WaiterId, oneshot::Sender<Result<RequestOutput, Error>>),
Timer(TimerId, oneshot::Sender<Result<(), Error>>),
}
pub enum Cancel {
Waiter(WaiterId),
Timer(TimerId),
}
struct Inbox {
open: bool,
messages: Vec<Message>,
}
pub struct Mailbox {
pub waker: Waker,
inbox: Mutex<Inbox>,
}
impl Mailbox {
pub fn new() -> std::io::Result<Self> {
Ok(Self {
waker: Waker::new()?,
inbox: Mutex::new(Inbox {
open: true,
messages: Vec::new(),
}),
})
}
fn lock(&self) -> MutexGuard<'_, Inbox> {
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
let inbox = self.inbox.lock().unwrap();
} else {
let inbox = self.inbox.lock();
}
}
inbox
}
pub fn send(&self, message: Message) -> Result<(), Message> {
let signal = {
let mut inbox = self.lock();
if !inbox.open {
return Err(message);
}
let first = inbox.messages.is_empty();
inbox.messages.push(message);
first && self.waker.publish()
};
if signal {
self.waker.wake();
}
Ok(())
}
pub fn take(&self, scratch: &mut Vec<Message>) -> bool {
assert!(
scratch.is_empty(),
"mailbox scratch must be drained before transfer"
);
let mut inbox = self.lock();
if inbox.messages.is_empty() {
return false;
}
mem::swap(&mut inbox.messages, scratch);
true
}
pub fn close(&self) -> Vec<Message> {
let mut inbox = self.lock();
inbox.open = false;
mem::take(&mut inbox.messages)
}
#[cfg(test)]
pub fn is_open(&self) -> bool {
self.lock().open
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::iouring::{slab::Id, waker::tests::eventfd_count};
use futures::task::{ArcWake, waker};
use std::{
future::Future,
pin::Pin,
sync::{
Arc, Barrier, Weak,
atomic::{AtomicBool, Ordering},
},
task::Context,
thread,
};
struct Reentrant {
mailbox: Weak<Mailbox>,
woken: AtomicBool,
}
impl ArcWake for Reentrant {
fn wake_by_ref(this: &Arc<Self>) {
let mailbox = this.mailbox.upgrade().unwrap();
let inbox = mailbox
.inbox
.try_lock()
.expect("messages must be dropped outside the inbox lock");
assert!(!inbox.open);
this.woken.store(true, Ordering::Relaxed);
}
}
fn forward_message(
mailbox: &Arc<Mailbox>,
) -> (
Message,
oneshot::Receiver<Result<RequestOutput, Error>>,
Arc<Reentrant>,
) {
let (sender, mut receiver) = oneshot::channel();
let reentrant = Arc::new(Reentrant {
mailbox: Arc::downgrade(mailbox),
woken: AtomicBool::new(false),
});
let waker = waker(reentrant.clone());
let mut cx = Context::from_waker(&waker);
assert!(Pin::new(&mut receiver).poll(&mut cx).is_pending());
let id = WaiterId(Id {
index: 0,
generation: 0,
});
(
Message::Forward(Forward::Waiter(id, sender)),
receiver,
reentrant,
)
}
const fn cancel_message() -> Message {
Message::Cancel(Cancel::Waiter(WaiterId(Id {
index: 0,
generation: 0,
})))
}
#[test]
fn test_messages_publish_once_per_batch() {
let mailbox = Mailbox::new().unwrap();
let mut scratch = Vec::new();
assert!(mailbox.is_open());
assert!(!mailbox.take(&mut scratch));
assert!(!mailbox.waker.pending(0));
assert!(mailbox.send(Message::WakeRoot).is_ok());
assert!(mailbox.send(cancel_message()).is_ok());
assert!(mailbox.waker.pending(0));
assert!(mailbox.take(&mut scratch));
assert!(!mailbox.waker.pending(1));
assert!(matches!(
scratch.as_slice(),
[Message::WakeRoot, Message::Cancel(_)]
));
assert!(mailbox.send(Message::WakeRoot).is_ok());
assert!(mailbox.waker.pending(1));
scratch.clear();
assert!(mailbox.take(&mut scratch));
assert!(!mailbox.waker.pending(2));
assert!(matches!(scratch.as_slice(), [Message::WakeRoot]));
scratch.clear();
assert!(!mailbox.take(&mut scratch));
assert!(!mailbox.waker.pending(2));
}
#[test]
fn test_send_signals_armed_worker() {
let mailbox = Mailbox::new().unwrap();
let arm = mailbox.waker.arm(0);
assert!(arm.still_idle());
assert!(mailbox.send(Message::WakeRoot).is_ok());
assert!(mailbox.waker.pending(0));
assert_eq!(eventfd_count(&mailbox.waker), 1);
}
#[test]
#[should_panic(expected = "mailbox scratch must be drained before transfer")]
fn test_take_requires_empty_scratch() {
let mailbox = Mailbox::new().unwrap();
let mut scratch = vec![Message::WakeRoot];
mailbox.take(&mut scratch);
}
#[test]
fn test_close_returns_messages_and_rejects_new_ones() {
let mailbox = Arc::new(Mailbox::new().unwrap());
let (message, _receiver, reentrant) = forward_message(&mailbox);
assert!(mailbox.send(message).is_ok());
let queued = mailbox.close();
assert!(!mailbox.is_open());
assert!(matches!(queued.as_slice(), [Message::Forward(_)]));
assert!(!reentrant.woken.load(Ordering::Relaxed));
drop(queued);
assert!(reentrant.woken.load(Ordering::Relaxed));
let (message, _receiver, reentrant) = forward_message(&mailbox);
let rejected = mailbox.send(message);
assert!(matches!(rejected, Err(Message::Forward(_))));
assert!(!reentrant.woken.load(Ordering::Relaxed));
assert!(mailbox.waker.pending(0));
assert!(!mailbox.waker.pending(1));
drop(rejected);
assert!(reentrant.woken.load(Ordering::Relaxed));
let mut scratch = Vec::new();
assert!(!mailbox.take(&mut scratch));
assert!(mailbox.close().is_empty());
}
#[test]
fn test_send_racing_close_preserves_payload_ownership() {
let mailbox = Arc::new(Mailbox::new().unwrap());
let gate = Arc::new(Barrier::new(2));
let (message, _receiver, reentrant) = forward_message(&mailbox);
let producer = thread::spawn({
let mailbox = mailbox.clone();
let gate = gate.clone();
move || {
gate.wait();
mailbox.send(message)
}
});
gate.wait();
let queued = mailbox.close();
let result = producer.join().unwrap();
assert_eq!(queued.len(), usize::from(result.is_ok()));
assert!(!mailbox.is_open());
assert!(!reentrant.woken.load(Ordering::Relaxed));
drop(queued);
drop(result);
assert!(reentrant.woken.load(Ordering::Relaxed));
}
}