use loom::sync::atomic::{
AtomicBool, AtomicUsize,
Ordering::{Acquire, Relaxed, Release},
};
use loom::sync::{Arc, Mutex};
use std::collections::VecDeque;
pub struct WakeupGate {
registered: Mutex<Option<loom::thread::Thread>>,
}
impl WakeupGate {
pub fn new() -> Self {
Self {
registered: Mutex::new(None),
}
}
pub fn enable(&self) {
*self.registered.lock().unwrap() = Some(loom::thread::current());
}
pub fn notify_waiters(&self) {
if let Some(waiter) = self.registered.lock().unwrap().take() {
waiter.unpark();
}
}
pub fn wait(&self) {
while self.registered.lock().unwrap().is_some() {
loom::thread::park();
}
}
}
pub struct LoomMailbox {
queue: Mutex<VecDeque<usize>>,
permits: Mutex<usize>,
size: AtomicUsize,
backpressure: AtomicBool,
capacity: usize,
backpressure_count: usize,
notify: WakeupGate,
}
impl LoomMailbox {
pub fn with_capacity(capacity: usize) -> Self {
Self {
queue: Mutex::new(VecDeque::new()),
permits: Mutex::new(capacity),
size: AtomicUsize::new(0),
backpressure: AtomicBool::new(false),
capacity,
backpressure_count: (capacity as f32 * 0.8) as usize,
notify: WakeupGate::new(),
}
}
pub fn try_send(&self, message: usize) -> bool {
{
let mut permits = self.permits.lock().unwrap();
if *permits == 0 {
return false;
}
*permits -= 1;
}
let size = self.size.fetch_add(1, Relaxed) + 1;
self.queue.lock().unwrap().push_back(message);
if size >= self.backpressure_count {
self.backpressure.store(true, Relaxed);
}
self.notify.notify_waiters();
true
}
pub fn try_recv(&self) -> Option<usize> {
let message = self.queue.lock().unwrap().pop_front()?;
self.size.fetch_sub(1, Relaxed);
*self.permits.lock().unwrap() += 1;
self.check_backpressure();
Some(message)
}
pub fn recv(&self) -> usize {
loop {
self.notify.enable();
if let Some(message) = self.try_recv() {
return message;
}
self.notify.wait();
}
}
pub fn len(&self) -> usize {
self.size.load(Relaxed)
}
pub fn is_backpressured(&self) -> bool {
self.backpressure.load(Relaxed)
}
pub fn is_fully_drained(&self) -> bool {
*self.permits.lock().unwrap() == self.capacity
&& self.size.load(Acquire) == 0
&& self.queue.lock().unwrap().is_empty()
}
fn check_backpressure(&self) {
let size = self.size.load(Relaxed);
if size < self.backpressure_count / 2 {
self.backpressure.store(false, Relaxed);
}
}
}
pub fn model_capacity_gate_two_senders() {
loom::model(move || {
let mailbox = Arc::new(LoomMailbox::with_capacity(2));
let mut handles = Vec::new();
for sender in 0..2usize {
let mailbox = Arc::clone(&mailbox);
handles.push(loom::thread::spawn(move || {
let mut successes = 0;
for _ in 0..2 {
if mailbox.try_send(sender) {
successes += 1;
}
}
successes
}));
}
let mut total = 0;
for h in handles {
total += h.join().unwrap();
}
assert!(
total <= 2,
"capacity gate raced through: {total} sends succeeded into capacity 2"
);
assert_eq!(mailbox.len(), total, "size counter must equal deliveries");
assert_eq!(
mailbox.is_backpressured(),
total >= 1,
"backpressure latch must reflect the final size (threshold 1 at capacity 2)"
);
let mut drained = 0;
while mailbox.try_recv().is_some() {
drained += 1;
}
assert_eq!(drained, total);
assert!(
mailbox.is_fully_drained(),
"permit/size conservation broken"
);
assert!(
mailbox.try_send(9),
"post-drain send must succeed after full permit restoration"
);
});
}
pub fn model_permit_conservation_roundtrip() {
const ATTEMPTS: usize = 4;
loom::model(move || {
let mailbox = Arc::new(LoomMailbox::with_capacity(2));
let mut handles = Vec::new();
for tag in 1..=2usize {
let mailbox = Arc::clone(&mailbox);
handles.push(loom::thread::spawn(move || {
for _ in 0..ATTEMPTS {
if mailbox.try_send(tag) {
return;
}
}
}));
}
let received = {
let mailbox = Arc::clone(&mailbox);
loom::thread::spawn(move || {
let mut got = Vec::new();
for _ in 0..ATTEMPTS * 2 {
if got.len() == 2 {
break;
}
if let Some(message) = mailbox.try_recv() {
got.push(message);
}
}
got
})
};
for h in handles {
h.join().unwrap();
}
let mut got = received.join().unwrap();
while let Some(message) = mailbox.try_recv() {
got.push(message);
}
got.sort_unstable();
assert!(
got == vec![1] || got == vec![2] || got == vec![1, 2],
"deliveries must be a duplicate-free subset of the sent tags"
);
assert!(
mailbox.is_fully_drained(),
"permit/size conservation broken"
);
});
}
pub fn model_recv_no_lost_wakeup() {
loom::model(move || {
let mailbox = Arc::new(LoomMailbox::with_capacity(2));
let sender = {
let mailbox = Arc::clone(&mailbox);
loom::thread::spawn(move || {
assert!(mailbox.try_send(1), "first send into a fresh mailbox fits");
})
};
let receiver = {
let mailbox = Arc::clone(&mailbox);
loom::thread::spawn(move || mailbox.recv())
};
sender.join().unwrap();
let received = receiver.join().unwrap();
assert_eq!(
received, 1,
"receiver must wake to exactly the sent message"
);
assert!(mailbox.is_fully_drained());
});
}
pub fn model_two_senders_receiver_drains_via_handshake() {
loom::model(move || {
let mailbox = Arc::new(LoomMailbox::with_capacity(2));
let mut senders = Vec::new();
for tag in 1..=2usize {
let mailbox = Arc::clone(&mailbox);
senders.push(loom::thread::spawn(move || {
assert!(mailbox.try_send(tag), "single send per sender always fits");
}));
}
let receiver = {
let mailbox = Arc::clone(&mailbox);
loom::thread::spawn(move || vec![mailbox.recv(), mailbox.recv()])
};
for h in senders {
h.join().unwrap();
}
let mut received = receiver.join().unwrap();
received.sort_unstable();
assert_eq!(received, vec![1, 2], "each message delivered exactly once");
assert!(mailbox.is_fully_drained());
});
}
pub fn model_broken_recv_loses_wakeup() {
loom::model(move || {
let mailbox = Arc::new(LoomMailbox::with_capacity(2));
let sender = {
let mailbox = Arc::clone(&mailbox);
loom::thread::spawn(move || {
assert!(mailbox.try_send(1));
})
};
let receiver = {
let mailbox = Arc::clone(&mailbox);
loom::thread::spawn(move || {
loop {
if let Some(message) = mailbox.try_recv() {
return message;
}
mailbox.notify.enable();
mailbox.notify.wait();
}
})
};
sender.join().unwrap();
let received = receiver.join().unwrap();
assert_eq!(received, 1);
});
}
#[cfg(all(test, loom))]
mod loom_tests {
use super::*;
#[test]
fn loom_capacity_gate_two_senders() {
model_capacity_gate_two_senders();
}
#[test]
fn loom_permit_conservation_roundtrip() {
model_permit_conservation_roundtrip();
}
#[test]
fn loom_recv_no_lost_wakeup() {
model_recv_no_lost_wakeup();
}
#[test]
fn loom_two_senders_receiver_drains_via_handshake() {
model_two_senders_receiver_drains_via_handshake();
}
}