use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
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::admission::{default_capacity, defer_after_append};
use crate::channel::schema::SchemaId;
use crate::channel::wire::{decode_envelope, encode_envelope};
use crate::durability::bridge::block_on;
use crate::durability::{DurableStore, MessageEnvelope, replay_range};
use crate::envelope::{Envelope, PublisherId};
use crate::error::LiminalError;
use crate::pressure::{CapacityError, CapacityTracker, ConsumerCapacity, PressureSignal};
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>,
pub capacity: Option<ConsumerCapacity>,
}
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())
.field("capacity", &self.capacity)
.finish_non_exhaustive()
}
}
struct InboxState {
queue: VecDeque<(Envelope, usize)>,
budget: Option<Arc<ConnectionInboxBudget>>,
depth_cap: usize,
notifier: Option<InboxNotifier>,
closed: bool,
capacity: ConsumerCapacity,
lagging: bool,
shed_generation: u64,
next_replay_seq: u64,
replay_offered_seq: u64,
}
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, PartialEq, Eq)]
pub(crate) enum InboxAdmission {
Admitted(PressureSignal),
Deferred(PressureSignal),
Rejected(PressureSignal),
AlreadyOffered(PressureSignal),
BudgetExceeded(PressureSignal),
FairnessTripped(PressureSignal),
Closed,
}
impl InboxAdmission {
pub(crate) const fn is_queued(&self) -> bool {
matches!(
*self,
Self::Admitted(_) | Self::Deferred(_) | Self::AlreadyOffered(_)
)
}
pub(crate) const fn signal(&self) -> Option<&PressureSignal> {
match *self {
Self::Admitted(ref signal)
| Self::Deferred(ref signal)
| Self::Rejected(ref signal)
| Self::AlreadyOffered(ref signal) => Some(signal),
Self::BudgetExceeded(_) | Self::FairnessTripped(_) | Self::Closed => None,
}
}
pub(crate) const fn shed_signal(&self) -> Option<&PressureSignal> {
match *self {
Self::BudgetExceeded(ref signal) | Self::FairnessTripped(ref signal) => Some(signal),
Self::Admitted(_)
| Self::Deferred(_)
| Self::Rejected(_)
| Self::AlreadyOffered(_)
| Self::Closed => None,
}
}
}
const fn derived_signal(capacity: &ConsumerCapacity, queued: usize) -> PressureSignal {
CapacityTracker::derived(*capacity, queued).pressure_signal()
}
const fn shed_signal(capacity: &ConsumerCapacity, queued: usize) -> PressureSignal {
let tracker = CapacityTracker::derived(*capacity, queued);
PressureSignal::reject(
tracker.current_in_flight(),
capacity.max_in_flight,
tracker.current_buffer_depth(),
capacity.max_buffer_depth,
)
}
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,
capacity: default_capacity(),
lagging: false,
shed_generation: 0,
next_replay_seq: 0,
replay_offered_seq: 0,
}),
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_capacity(&self, capacity: ConsumerCapacity) -> Result<(), CapacityError> {
capacity.validate()?;
if let Ok(mut state) = self.state.lock() {
state.capacity = capacity;
}
Ok(())
}
pub(crate) fn seed_replay_cursor(&self, head: u64) {
if let Ok(mut state) = self.state.lock() {
state.next_replay_seq = head;
}
}
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 {
self.admit_at(envelope, None)
}
pub(crate) fn admit_at(
&self,
envelope: Envelope,
durable_position: Option<u64>,
) -> InboxAdmission {
self.admit_inner(envelope, durable_position, false)
}
pub(crate) fn admit_replayed(&self, envelope: Envelope, sequence: u64) -> bool {
self.admit_inner(envelope, Some(sequence), true).is_queued()
}
pub(crate) fn note_replayed_filtered(&self, sequence: u64) {
if let Ok(mut state) = self.state.lock()
&& !state.closed
{
let next = sequence.saturating_add(1);
state.next_replay_seq = state.next_replay_seq.max(next);
state.replay_offered_seq = state.replay_offered_seq.max(next);
}
}
fn admit_inner(
&self,
envelope: Envelope,
durable_position: Option<u64>,
replay: bool,
) -> InboxAdmission {
let bytes = encode_envelope(&envelope).len();
let signal;
let notifier = {
let Ok(mut state) = self.state.lock() else {
self.overflowed.store(true, Ordering::Release);
return InboxAdmission::BudgetExceeded(PressureSignal::reject(0, 0, 0, 0));
};
if state.closed {
return InboxAdmission::Closed;
}
if !replay
&& let Some(position) = durable_position
&& position < state.replay_offered_seq
{
let signal = defer_after_append(derived_signal(&state.capacity, state.queue.len()));
return InboxAdmission::AlreadyOffered(signal);
}
if state.lagging && !replay {
state.shed_generation = state.shed_generation.wrapping_add(1);
let capacity = state.capacity;
return InboxAdmission::Rejected(PressureSignal::reject(
capacity.max_in_flight,
capacity.max_in_flight,
capacity.max_buffer_depth,
capacity.max_buffer_depth,
));
}
signal = derived_signal(&state.capacity, state.queue.len());
if matches!(signal, PressureSignal::Reject { .. }) {
if durable_position.is_some() && !replay {
state.lagging = true;
state.shed_generation = state.shed_generation.wrapping_add(1);
}
return InboxAdmission::Rejected(signal);
}
if state.queue.len() >= state.depth_cap {
self.overflowed.store(true, Ordering::Release);
let shed = shed_signal(&state.capacity, state.queue.len());
let notifier = state.notifier.clone();
drop(state);
if let Some(notifier) = notifier {
notifier();
}
return InboxAdmission::FairnessTripped(shed);
}
let charged = match state.budget.as_ref() {
Some(budget) => {
if !budget.try_charge(bytes) {
self.overflowed.store(true, Ordering::Release);
let shed = shed_signal(&state.capacity, state.queue.len());
let notifier = state.notifier.clone();
drop(state);
if let Some(notifier) = notifier {
notifier();
}
return InboxAdmission::BudgetExceeded(shed);
}
bytes
}
None => 0,
};
state.queue.push_back((envelope, charged));
if let Some(position) = durable_position {
let next = position.saturating_add(1);
state.next_replay_seq = state.next_replay_seq.max(next);
if replay {
state.replay_offered_seq = state.replay_offered_seq.max(next);
}
}
state.notifier.clone()
};
if let Some(notifier) = notifier {
notifier();
}
if matches!(signal, PressureSignal::Defer { .. }) {
InboxAdmission::Deferred(signal)
} else {
InboxAdmission::Admitted(signal)
}
}
pub(crate) fn note_filtered(&self, durable_position: u64) {
if let Ok(mut state) = self.state.lock()
&& !state.lagging
&& !state.closed
{
state.next_replay_seq = state
.next_replay_seq
.max(durable_position.saturating_add(1));
}
}
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)
}
pub(crate) fn queued_len(&self) -> usize {
self.state.lock().map_or(0, |state| state.queue.len())
}
#[cfg(test)]
pub(crate) fn len(&self) -> usize {
self.queued_len()
}
pub(crate) fn capacity_bound(&self) -> usize {
self.state.lock().map_or(0, |state| {
state
.capacity
.max_in_flight
.saturating_add(state.capacity.max_buffer_depth)
})
}
pub(crate) fn is_lagging(&self) -> bool {
self.state.lock().is_ok_and(|state| state.lagging)
}
pub(crate) fn refill_plan(&self) -> Option<(u64, u64, usize)> {
let state = self.state.lock().ok()?;
if !state.lagging || state.closed {
return None;
}
let queued = state.queue.len();
let low_watermark = state.capacity.max_in_flight / 2;
if queued > low_watermark {
return None;
}
let bound = state
.capacity
.max_in_flight
.saturating_add(state.capacity.max_buffer_depth);
let budget = bound.saturating_sub(queued);
if budget == 0 {
return None;
}
Some((state.shed_generation, state.next_replay_seq, budget))
}
pub(crate) fn clear_lagging_if_unchanged(&self, generation: u64) -> bool {
let Ok(mut state) = self.state.lock() else {
return false;
};
if state.shed_generation != generation {
return false;
}
state.lagging = false;
true
}
}
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 BinaryRef::new(message).is_some() {
if let Ok(owned) = beamr::ets::copy_term_to_ets(message)
&& let Some(binary) = BinaryRef::new(owned.root())
{
self.accept_remote_frame(binary.as_bytes(owned.borrow_terms()));
}
}
}
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,
durable_position: Option<u64>,
) -> Option<InboxAdmission> {
if let Some(predicate) = self.predicate.as_ref() {
if !predicate(envelope) {
if let Some(position) = durable_position {
self.inbox.note_filtered(position);
}
return None;
}
}
Some(self.inbox.admit_at(envelope.clone(), durable_position))
}
pub(crate) fn occupancy(&self) -> (usize, usize) {
(self.inbox.queued_len(), self.inbox.capacity_bound())
}
}
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>,
refill: OnceLock<DurableRefill>,
refilling: Mutex<()>,
}
struct DurableRefill {
store: Arc<dyn DurableStore>,
stream_key: String,
schema_id: SchemaId,
predicate: Option<SubscriptionPredicate>,
}
impl std::fmt::Debug for DurableRefill {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("DurableRefill")
.field("stream_key", &self.stream_key)
.field("has_predicate", &self.predicate.is_some())
.finish_non_exhaustive()
}
}
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(capacity) = install.capacity {
inbox.install_capacity(capacity).map_err(|error| {
LiminalError::SubscriptionFailed {
message: format!("declared consumer capacity is invalid: {error}"),
}
})?;
}
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_trap_exit(factory).map_err(|error| {
LiminalError::SubscriptionFailed {
message: format!("failed to spawn subscriber process: {error:?}"),
}
})?;
let handle = Self {
inner: Arc::new(SubscriptionInner {
pid,
inbox: Arc::clone(&inbox),
scheduler: Arc::clone(scheduler),
refill: OnceLock::new(),
refilling: Mutex::new(()),
}),
};
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> {
let next = self.inner.inbox.pop();
self.refill_if_lagging();
Ok(next)
}
pub(crate) fn seed_replay_cursor(&self, head: u64) {
self.inner.inbox.seed_replay_cursor(head);
}
pub(crate) fn attach_durable_refill(
&self,
store: Arc<dyn DurableStore>,
stream_key: String,
schema_id: SchemaId,
predicate: Option<SubscriptionPredicate>,
) {
let attached = self.inner.refill.set(DurableRefill {
store,
stream_key,
schema_id,
predicate,
});
debug_assert!(
attached.is_ok(),
"the refill OnceLock is set exactly once per attach: attach_durable_refill \
was called twice on one subscription, which only the durable subscribe \
path may call and only before the handle is returned"
);
}
fn refill_if_lagging(&self) {
let Some(refill) = self.inner.refill.get() else {
return;
};
if !self.inner.inbox.is_lagging() {
return;
}
let Ok(_guard) = self.inner.refilling.try_lock() else {
return;
};
while let Some((generation, cursor, budget)) = self.inner.inbox.refill_plan() {
let Ok(Ok(batch)) = block_on(replay_range(
refill.store.as_ref(),
&refill.stream_key,
cursor,
budget,
)) else {
return;
};
if batch.is_empty() {
if self.inner.inbox.clear_lagging_if_unchanged(generation) {
return;
}
continue;
}
for (sequence, stored) in batch {
let envelope = refilled_envelope(&stored, refill.schema_id);
let matched = refill
.predicate
.as_ref()
.is_none_or(|predicate| predicate(&envelope));
if matched {
if !self.inner.inbox.admit_replayed(envelope, sequence) {
return;
}
} else {
self.inner.inbox.note_replayed_filtered(sequence);
}
}
}
}
#[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);
}
}
fn refilled_envelope(stored: &MessageEnvelope, schema_id: SchemaId) -> Envelope {
let millis = i64::try_from(stored.timestamp).unwrap_or(i64::MAX);
let timestamp = chrono::TimeZone::timestamp_millis_opt(&chrono::Utc, millis)
.single()
.unwrap_or_else(chrono::Utc::now);
Envelope::with_timestamp(
stored.payload.clone(),
None,
schema_id,
PublisherId::new(stored.publisher_id.clone()),
timestamp,
)
}
#[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};
use crate::pressure::ConsumerCapacity;
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_for_every_admitted_envelope() {
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!(inbox.admit(envelope(b"a")).is_queued());
assert_eq!(fires.load(Ordering::Relaxed), 1, "the first admit fires");
assert!(inbox.admit(envelope(b"b")).is_queued());
assert_eq!(
fires.load(Ordering::Relaxed),
2,
"an admit into a non-empty inbox still fires"
);
assert!(inbox.pop().is_some());
assert!(inbox.pop().is_some());
assert!(inbox.admit(envelope(b"c")).is_queued());
assert_eq!(
fires.load(Ordering::Relaxed),
3,
"one fire per admitted envelope, whatever the queue depth was"
);
}
#[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!(inbox_a.admit(one).is_queued());
assert_eq!(budget.used(), cap, "the shared budget is now fully spent");
assert!(matches!(
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!(inbox.admit(env.clone()).is_queued());
assert!(matches!(
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!(inbox.admit(envelope(b"1")).is_queued());
assert!(inbox.admit(envelope(b"2")).is_queued());
assert!(matches!(
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!(inbox.admit(a).is_queued());
assert!(inbox.admit(b).is_queued());
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!(inbox_a.admit(one.clone()).is_queued());
}
assert_eq!(budget.used(), unit * 256, "the backlog holds the budget");
assert!(matches!(
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!(
inbox_b.admit(one).is_queued(),
"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!(inbox.admit(env.clone()).is_queued());
assert!(inbox.admit(env).is_queued());
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!(inbox.admit(env).is_queued());
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!(inbox.admit(uncharged).is_queued());
let budget = ConnectionInboxBudget::new(1024 * 1024);
inbox.install_budget(Arc::clone(&budget), usize::MAX);
let unit = admitted_bytes(&charged);
assert!(inbox.admit(charged).is_queued());
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!(inbox.admit(envelope(b"pre-install")).is_queued());
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!(inbox.admit(envelope(b"second")).is_queued());
assert_eq!(fires.load(Ordering::Relaxed), 2);
}
#[test]
fn a_live_push_is_not_suppressed_by_a_higher_positioned_sibling() {
let inbox = SubscriptionInbox::new();
assert!(
inbox.admit_at(envelope(b"position-5"), Some(5)).is_queued(),
"the higher-positioned push is queued"
);
let behind = inbox.admit_at(envelope(b"position-3"), Some(3));
assert!(
!matches!(behind, InboxAdmission::AlreadyOffered(_)),
"position 3 was offered by NOBODY — its only sin is arriving behind \
position 5, and suppressing it loses it silently; got {behind:?}"
);
assert!(behind.is_queued(), "position 3 must reach the queue");
assert_eq!(
inbox.queued_len(),
2,
"both live pushes are in the queue, in arrival order"
);
}
#[test]
fn a_live_push_the_replay_door_already_offered_is_suppressed() {
let inbox = SubscriptionInbox::new();
assert!(
inbox.admit_replayed(envelope(b"replayed-5"), 5),
"the replay door offers position 5"
);
let live = inbox.admit_at(envelope(b"live-5"), Some(5));
assert!(
matches!(live, InboxAdmission::AlreadyOffered(_)),
"position 5 was already offered by the refill; got {live:?}"
);
assert_eq!(
inbox.queued_len(),
1,
"the suppressed push did not enter the queue a second time"
);
}
#[test]
fn the_refill_cursor_never_backsteps_below_a_delivered_position() {
let inbox = SubscriptionInbox::new();
inbox
.install_capacity(ConsumerCapacity::new(1, 1).expect("1/1 is a legal capacity"))
.expect("a legal capacity installs");
assert!(inbox.admit_at(envelope(b"position-5"), Some(5)).is_queued());
assert!(inbox.admit_at(envelope(b"position-3"), Some(3)).is_queued());
assert!(matches!(
inbox.admit_at(envelope(b"position-6"), Some(6)),
InboxAdmission::Rejected(_)
));
assert!(inbox.is_lagging(), "the shed opened a gap");
assert!(inbox.pop().is_some());
assert!(inbox.pop().is_some());
let (_generation, cursor, _budget) = inbox
.refill_plan()
.expect("a drained lagging inbox has a refill due");
assert_eq!(
cursor, 6,
"the missed range starts above every delivered position; a cursor of \
4 would send the refill back over positions 4 and 5, and position 5 \
is already in this subscriber's hands"
);
}
}