use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use beamr::native::native_process::{NativeContext, NativeHandler, NativeOutcome};
use beamr::process::ExitReason;
use beamr::scheduler::Scheduler;
use beamr::term::binary_ref::BinaryRef;
use crate::channel::wire::{decode_envelope, encode_envelope};
use crate::envelope::Envelope;
use crate::error::LiminalError;
pub(crate) type SubscriberInbox = Arc<SubscriptionInbox>;
pub type InboxNotifier = Arc<dyn Fn() + Send + Sync>;
#[derive(Debug)]
pub struct ConnectionInboxBudget {
used: AtomicUsize,
cap: usize,
}
impl ConnectionInboxBudget {
#[must_use]
pub fn new(cap: usize) -> Arc<Self> {
Arc::new(Self {
used: AtomicUsize::new(0),
cap,
})
}
fn try_charge(&self, bytes: usize) -> bool {
let mut current = self.used.load(Ordering::Acquire);
loop {
let Some(projected) = current.checked_add(bytes) else {
return false;
};
if projected > self.cap {
return false;
}
match self.used.compare_exchange_weak(
current,
projected,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => return true,
Err(observed) => current = observed,
}
}
}
fn release(&self, bytes: usize) {
let mut current = self.used.load(Ordering::Acquire);
loop {
let next = current.saturating_sub(bytes);
match self.used.compare_exchange_weak(
current,
next,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => return,
Err(observed) => current = observed,
}
}
}
#[cfg(test)]
pub(crate) fn used(&self) -> usize {
self.used.load(Ordering::Acquire)
}
}
pub struct InboxInstall {
pub budget: Arc<ConnectionInboxBudget>,
pub depth_cap: usize,
pub notifier: Option<InboxNotifier>,
}
impl std::fmt::Debug for InboxInstall {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("InboxInstall")
.field("depth_cap", &self.depth_cap)
.field("has_notifier", &self.notifier.is_some())
.finish_non_exhaustive()
}
}
struct InboxState {
queue: VecDeque<(Envelope, usize)>,
budget: Option<Arc<ConnectionInboxBudget>>,
depth_cap: usize,
notifier: Option<InboxNotifier>,
closed: bool,
}
pub(crate) struct SubscriptionInbox {
state: Mutex<InboxState>,
overflowed: AtomicBool,
}
impl std::fmt::Debug for SubscriptionInbox {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("SubscriptionInbox")
.field("overflowed", &self.overflowed.load(Ordering::Acquire))
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum InboxAdmission {
Admitted,
BudgetExceeded,
FairnessTripped,
Closed,
}
impl SubscriptionInbox {
pub(crate) fn new() -> Arc<Self> {
Arc::new(Self {
state: Mutex::new(InboxState {
queue: VecDeque::new(),
budget: None,
depth_cap: usize::MAX,
notifier: None,
closed: false,
}),
overflowed: AtomicBool::new(false),
})
}
pub(crate) fn install_budget(&self, budget: Arc<ConnectionInboxBudget>, depth_cap: usize) {
if let Ok(mut state) = self.state.lock() {
state.budget = Some(budget);
state.depth_cap = depth_cap;
}
}
pub(crate) fn install_notifier(&self, notifier: InboxNotifier) {
let fire = {
let Ok(mut state) = self.state.lock() else {
return;
};
let pending = !state.queue.is_empty();
let handle = notifier.clone();
state.notifier = Some(notifier);
pending.then_some(handle)
};
if let Some(notifier) = fire {
notifier();
}
}
pub(crate) fn admit(&self, envelope: Envelope) -> InboxAdmission {
let bytes = encode_envelope(&envelope).len();
let notifier = {
let Ok(mut state) = self.state.lock() else {
self.overflowed.store(true, Ordering::Release);
return InboxAdmission::BudgetExceeded;
};
if state.closed {
return InboxAdmission::Closed;
}
if state.queue.len() >= state.depth_cap {
self.overflowed.store(true, Ordering::Release);
let notifier = state.notifier.clone();
drop(state);
if let Some(notifier) = notifier {
notifier();
}
return InboxAdmission::FairnessTripped;
}
let charged = match state.budget.as_ref() {
Some(budget) => {
if !budget.try_charge(bytes) {
self.overflowed.store(true, Ordering::Release);
let notifier = state.notifier.clone();
drop(state);
if let Some(notifier) = notifier {
notifier();
}
return InboxAdmission::BudgetExceeded;
}
bytes
}
None => 0,
};
let was_empty = state.queue.is_empty();
state.queue.push_back((envelope, charged));
if was_empty {
state.notifier.clone()
} else {
None
}
};
if let Some(notifier) = notifier {
notifier();
}
InboxAdmission::Admitted
}
pub(crate) fn pop(&self) -> Option<Envelope> {
let (envelope, charged, budget) = {
let mut state = self.state.lock().ok()?;
let (envelope, charged) = state.queue.pop_front()?;
(envelope, charged, state.budget.clone())
};
if let Some(budget) = budget {
budget.release(charged);
}
Some(envelope)
}
pub(crate) fn has_pending(&self) -> bool {
self.state.lock().is_ok_and(|state| !state.queue.is_empty())
}
pub(crate) fn close(&self) {
let (released, budget) = {
let Ok(mut state) = self.state.lock() else {
return;
};
if state.closed {
return;
}
state.closed = true;
let released: usize = state
.queue
.drain(..)
.map(|(_envelope, charged)| charged)
.sum();
state.notifier = None;
(released, state.budget.take())
};
if let Some(budget) = budget {
budget.release(released);
}
}
pub(crate) fn is_overflowed(&self) -> bool {
self.overflowed.load(Ordering::Acquire)
}
#[cfg(test)]
pub(crate) fn len(&self) -> usize {
self.state.lock().map_or(0, |state| state.queue.len())
}
}
impl Drop for SubscriptionInbox {
fn drop(&mut self) {
self.close();
}
}
pub(crate) type SubscriptionPredicate = Arc<dyn Fn(&Envelope) -> bool + Send + Sync>;
struct SubscriberProcess {
inbox: SubscriberInbox,
}
impl NativeHandler for SubscriberProcess {
fn handle(&mut self, ctx: &mut NativeContext<'_>) -> NativeOutcome {
ctx.set_trap_exit(true);
while let Some(message) = ctx.recv() {
if let Some(binary) = BinaryRef::new(message) {
self.accept_remote_frame(binary.as_bytes());
}
}
NativeOutcome::Wait
}
}
impl SubscriberProcess {
fn accept_remote_frame(&self, bytes: &[u8]) {
let Ok(envelope) = decode_envelope(bytes) else {
return;
};
self.inbox.admit(envelope);
}
}
pub(crate) struct SubscriberRegistration {
pid: u64,
inbox: SubscriberInbox,
predicate: Option<SubscriptionPredicate>,
}
impl SubscriberRegistration {
pub(crate) const fn pid(&self) -> u64 {
self.pid
}
pub(crate) fn deliver(&self, envelope: &Envelope) -> bool {
if let Some(predicate) = self.predicate.as_ref() {
if !predicate(envelope) {
return false;
}
}
matches!(self.inbox.admit(envelope.clone()), InboxAdmission::Admitted)
}
}
impl std::fmt::Debug for SubscriberRegistration {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("SubscriberRegistration")
.field("pid", &self.pid)
.field("has_predicate", &self.predicate.is_some())
.finish_non_exhaustive()
}
}
#[derive(Clone)]
pub struct SubscriptionHandle {
inner: Arc<SubscriptionInner>,
}
struct SubscriptionInner {
pid: u64,
inbox: SubscriberInbox,
scheduler: Arc<Scheduler>,
}
impl SubscriptionHandle {
pub(crate) fn spawn(
scheduler: &Arc<Scheduler>,
predicate: Option<SubscriptionPredicate>,
install: Option<InboxInstall>,
) -> Result<(Self, SubscriberRegistration), LiminalError> {
let inbox: SubscriberInbox = SubscriptionInbox::new();
if let Some(install) = install {
inbox.install_budget(install.budget, install.depth_cap);
if let Some(notifier) = install.notifier {
inbox.install_notifier(notifier);
}
}
let process_inbox = Arc::clone(&inbox);
let factory = Box::new(move || {
Box::new(SubscriberProcess {
inbox: Arc::clone(&process_inbox),
}) as Box<dyn NativeHandler>
});
let pid =
scheduler
.spawn_native(factory)
.map_err(|error| LiminalError::SubscriptionFailed {
message: format!("failed to spawn subscriber process: {error:?}"),
})?;
scheduler
.set_trap_exit(pid, true)
.map_err(|error| LiminalError::SubscriptionFailed {
message: format!("failed to set trap_exit on subscriber process {pid}: {error:?}"),
})?;
let handle = Self {
inner: Arc::new(SubscriptionInner {
pid,
inbox: Arc::clone(&inbox),
scheduler: Arc::clone(scheduler),
}),
};
let registration = SubscriberRegistration {
pid,
inbox,
predicate,
};
Ok((handle, registration))
}
#[must_use]
pub(crate) fn pid(&self) -> u64 {
self.inner.pid
}
pub fn try_next(&self) -> Result<Option<Envelope>, LiminalError> {
Ok(self.inner.inbox.pop())
}
#[must_use]
pub fn has_pending(&self) -> bool {
self.inner.inbox.has_pending()
}
#[must_use]
pub fn is_overflowed(&self) -> bool {
self.inner.inbox.is_overflowed()
}
}
impl std::fmt::Debug for SubscriptionHandle {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("SubscriptionHandle")
.field("pid", &self.inner.pid)
.finish_non_exhaustive()
}
}
impl Drop for SubscriptionInner {
fn drop(&mut self) {
self.inbox.close();
self.scheduler
.terminate_process(self.pid, ExitReason::Normal);
}
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
mod cooperative_smoke {
use std::cell::RefCell;
use std::rc::Rc;
use std::sync::Arc;
use beamr::atom::AtomTable;
use beamr::ets::copy_term_to_ets;
use beamr::module::ModuleRegistry;
use beamr::native::BifRegistryImpl;
use beamr::process::heap::Heap;
use beamr::scheduler::WasmScheduler;
use beamr::term::shared_binary::{SharedBinary, write_proc_bin};
use super::{SubscriberInbox, SubscriberProcess, SubscriptionInbox};
use crate::channel::SchemaId;
use crate::channel::wire::encode_envelope;
use crate::envelope::{Envelope, PublisherId};
fn cooperative_scheduler() -> Rc<RefCell<WasmScheduler>> {
let atom_table = Arc::new(AtomTable::with_common_atoms());
let modules = Arc::new(ModuleRegistry::new());
let bifs = Arc::new(BifRegistryImpl::new());
Rc::new(RefCell::new(WasmScheduler::new(atom_table, modules, bifs)))
}
fn frame_as_owned_binary(envelope: &Envelope) -> beamr::ets::OwnedTerm {
let bytes = encode_envelope(envelope);
let shared = SharedBinary::new(bytes);
let mut scratch = Heap::new(8);
let words = scratch
.alloc_slice(3)
.expect("scratch heap holds a proc-bin reference");
let term = write_proc_bin(words, &shared).expect("proc-bin term writes");
copy_term_to_ets(term).expect("frame copies into an owned binary")
}
fn sample_envelope() -> Envelope {
let timestamp = chrono::TimeZone::timestamp_millis_opt(&chrono::Utc, 1_700_000_000_123)
.single()
.expect("valid fixed millisecond timestamp");
Envelope::with_timestamp(
b"{\"value\":42}".to_vec(),
None,
SchemaId::new(),
PublisherId::from("publisher-cooperative"),
timestamp,
)
}
#[test]
fn real_subscriber_process_delivers_a_published_envelope_cooperatively() {
let scheduler = cooperative_scheduler();
let inbox: SubscriberInbox = SubscriptionInbox::new();
let process_inbox = Arc::clone(&inbox);
let pid = scheduler.borrow_mut().spawn_native_root(Box::new(move || {
Box::new(SubscriberProcess {
inbox: Arc::clone(&process_inbox),
}) as Box<dyn beamr::native::native_process::NativeHandler>
}));
scheduler.borrow_mut().run_until_idle();
assert_eq!(
inbox.len(),
0,
"no envelope is delivered before one is published"
);
let published = sample_envelope();
let frame = frame_as_owned_binary(&published);
scheduler
.borrow_mut()
.send_owned(pid, &frame)
.expect("frame is delivered to the live subscriber pid");
let mut delivered = None;
for _ in 0..8 {
scheduler.borrow_mut().run_until_idle();
let next = inbox.pop();
if let Some(envelope) = next {
delivered = Some(envelope);
break;
}
}
assert_eq!(
delivered.as_ref(),
Some(&published),
"the real subscriber decoded and delivered the published envelope"
);
}
}
#[cfg(test)]
#[allow(clippy::expect_used)]
mod inbox_bounding {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use super::{ConnectionInboxBudget, InboxAdmission, SubscriptionInbox};
use crate::channel::SchemaId;
use crate::channel::wire::encode_envelope;
use crate::envelope::{Envelope, PublisherId};
fn envelope(payload: &[u8]) -> Envelope {
Envelope::new(
payload.to_vec(),
None,
SchemaId::new(),
PublisherId::from("inbox-bounding-test"),
)
}
fn admitted_bytes(env: &Envelope) -> usize {
encode_envelope(env).len()
}
#[test]
fn notifier_fires_only_on_empty_to_non_empty_transition() {
let inbox = SubscriptionInbox::new();
let fires = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&fires);
inbox.install_notifier(Arc::new(move || {
counter.fetch_add(1, Ordering::Relaxed);
}));
assert_eq!(inbox.admit(envelope(b"a")), InboxAdmission::Admitted);
assert_eq!(
fires.load(Ordering::Relaxed),
1,
"empty→non-empty fires once"
);
assert_eq!(inbox.admit(envelope(b"b")), InboxAdmission::Admitted);
assert_eq!(
fires.load(Ordering::Relaxed),
1,
"no re-fire while the inbox stays non-empty"
);
assert!(inbox.pop().is_some());
assert!(inbox.pop().is_some());
assert_eq!(inbox.admit(envelope(b"c")), InboxAdmission::Admitted);
assert_eq!(
fires.load(Ordering::Relaxed),
2,
"a fresh empty→non-empty edge fires again"
);
}
#[test]
fn shared_budget_is_spent_across_all_a_connections_inboxes() {
let one = envelope(b"payload-one");
let two = envelope(b"payload-two");
let cap = admitted_bytes(&one);
let budget = ConnectionInboxBudget::new(cap);
let inbox_a = SubscriptionInbox::new();
let inbox_b = SubscriptionInbox::new();
inbox_a.install_budget(Arc::clone(&budget), usize::MAX);
inbox_b.install_budget(Arc::clone(&budget), usize::MAX);
assert_eq!(inbox_a.admit(one), InboxAdmission::Admitted);
assert_eq!(budget.used(), cap, "the shared budget is now fully spent");
assert_eq!(inbox_b.admit(two), InboxAdmission::BudgetExceeded);
assert!(
inbox_b.is_overflowed(),
"the sibling that overflowed the shared budget is shed"
);
assert!(!inbox_a.is_overflowed(), "the inbox that fit is not shed");
assert!(inbox_a.pop().is_some());
assert_eq!(
budget.used(),
0,
"release returns bytes to the shared budget"
);
}
#[test]
fn overflow_sheds_and_does_not_grow_memory() {
let env = envelope(b"x");
let budget = ConnectionInboxBudget::new(admitted_bytes(&env)); let inbox = SubscriptionInbox::new();
inbox.install_budget(budget, usize::MAX);
assert_eq!(inbox.admit(env.clone()), InboxAdmission::Admitted);
assert_eq!(inbox.admit(env), InboxAdmission::BudgetExceeded);
assert!(inbox.is_overflowed());
assert_eq!(
inbox.len(),
1,
"the overflowed envelope is dropped, not queued"
);
}
#[test]
fn per_inbox_fairness_trip_stops_one_inbox_starving_siblings() {
let budget = ConnectionInboxBudget::new(usize::MAX);
let inbox = SubscriptionInbox::new();
inbox.install_budget(budget, 2);
assert_eq!(inbox.admit(envelope(b"1")), InboxAdmission::Admitted);
assert_eq!(inbox.admit(envelope(b"2")), InboxAdmission::Admitted);
assert_eq!(inbox.admit(envelope(b"3")), InboxAdmission::FairnessTripped);
assert!(inbox.is_overflowed());
assert_eq!(
inbox.len(),
2,
"the fairness trip holds the inbox at its cap"
);
}
#[test]
fn charge_and_release_are_exact() {
let budget = ConnectionInboxBudget::new(1024 * 1024);
let inbox = SubscriptionInbox::new();
inbox.install_budget(Arc::clone(&budget), usize::MAX);
let a = envelope(b"first-envelope");
let b = envelope(b"second-longer-envelope-payload");
let charge = admitted_bytes(&a) + admitted_bytes(&b);
assert_eq!(inbox.admit(a), InboxAdmission::Admitted);
assert_eq!(inbox.admit(b), InboxAdmission::Admitted);
assert_eq!(budget.used(), charge, "used == sum of admitted bytes");
assert!(inbox.pop().is_some());
assert!(inbox.pop().is_some());
assert_eq!(
budget.used(),
0,
"every admitted byte is released on dequeue — exact symmetry"
);
}
#[test]
fn close_releases_queued_charges_so_siblings_recover() {
let one = envelope(b"backlog-envelope-payload");
let unit = admitted_bytes(&one);
let budget = ConnectionInboxBudget::new(unit * 256);
let inbox_a = SubscriptionInbox::new();
let inbox_b = SubscriptionInbox::new();
inbox_a.install_budget(Arc::clone(&budget), usize::MAX);
inbox_b.install_budget(Arc::clone(&budget), usize::MAX);
for _ in 0..256 {
assert_eq!(inbox_a.admit(one.clone()), InboxAdmission::Admitted);
}
assert_eq!(budget.used(), unit * 256, "the backlog holds the budget");
assert_eq!(inbox_b.admit(one.clone()), InboxAdmission::BudgetExceeded);
inbox_a.close();
assert_eq!(
budget.used(),
0,
"close releases the entire queued backlog back to the shared budget"
);
assert_eq!(
inbox_b.admit(one),
InboxAdmission::Admitted,
"a sibling admits again after the other inbox is shed"
);
}
#[test]
fn closed_inbox_refuses_without_charging() {
let env = envelope(b"post-close");
let budget = ConnectionInboxBudget::new(1024 * 1024);
let inbox = SubscriptionInbox::new();
inbox.install_budget(Arc::clone(&budget), usize::MAX);
inbox.close();
assert_eq!(inbox.admit(env), InboxAdmission::Closed);
assert_eq!(budget.used(), 0, "a closed inbox never charges the budget");
assert_eq!(inbox.len(), 0, "a closed inbox never queues");
}
#[test]
fn drop_backstop_releases_queued_charges() {
let env = envelope(b"dropped-while-queued");
let unit = admitted_bytes(&env);
let budget = ConnectionInboxBudget::new(1024 * 1024);
{
let inbox = SubscriptionInbox::new();
inbox.install_budget(Arc::clone(&budget), usize::MAX);
assert_eq!(inbox.admit(env.clone()), InboxAdmission::Admitted);
assert_eq!(inbox.admit(env), InboxAdmission::Admitted);
assert_eq!(budget.used(), unit * 2);
}
assert_eq!(
budget.used(),
0,
"dropping the last inbox handle releases every queued charge"
);
}
#[test]
fn close_is_idempotent_and_drains_the_queue() {
let env = envelope(b"x");
let budget = ConnectionInboxBudget::new(1024 * 1024);
let inbox = SubscriptionInbox::new();
inbox.install_budget(Arc::clone(&budget), usize::MAX);
assert_eq!(inbox.admit(env), InboxAdmission::Admitted);
inbox.close();
inbox.close(); assert_eq!(budget.used(), 0);
assert!(inbox.pop().is_none(), "a closed inbox holds nothing");
}
#[test]
fn per_entry_charge_ownership_survives_budget_install() {
let uncharged = envelope(b"admitted-before-budget-install");
let charged = envelope(b"admitted-after-budget-install");
let inbox = SubscriptionInbox::new();
assert_eq!(inbox.admit(uncharged), InboxAdmission::Admitted);
let budget = ConnectionInboxBudget::new(1024 * 1024);
inbox.install_budget(Arc::clone(&budget), usize::MAX);
let unit = admitted_bytes(&charged);
assert_eq!(inbox.admit(charged), InboxAdmission::Admitted);
assert_eq!(budget.used(), unit, "only the post-install entry charged");
assert!(inbox.pop().is_some());
assert_eq!(budget.used(), unit, "the uncharged entry released nothing");
assert!(inbox.pop().is_some());
assert_eq!(
budget.used(),
0,
"the charged entry released its exact charge"
);
}
#[test]
fn notifier_install_onto_non_empty_inbox_fires_once() {
let inbox = SubscriptionInbox::new();
assert_eq!(
inbox.admit(envelope(b"pre-install")),
InboxAdmission::Admitted
);
let fires = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&fires);
inbox.install_notifier(Arc::new(move || {
counter.fetch_add(1, Ordering::Relaxed);
}));
assert_eq!(
fires.load(Ordering::Relaxed),
1,
"install onto a non-empty inbox regenerates exactly one wake"
);
assert_eq!(inbox.admit(envelope(b"second")), InboxAdmission::Admitted);
assert_eq!(fires.load(Ordering::Relaxed), 1);
}
}