liminal_sdk/remote/

mailbox.rs

//! The bounded mailbox between a background socket reader and its consumer.
//!
//! Every SDK stream with a reader thread (the push client and both
//! subscription streams) drains its socket into one of these and the consumer
//! takes from it. A consumer that stops receiving must not let the reader keep
//! pulling the server's bytes into memory, so a mailbox holds at most a frame
//! count and a payload byte count. When either bound is reached the reader
//! waits on the mailbox and stops reading the socket, so TCP flow control
//! carries the pressure back to the server; nothing already read is dropped.
//! A single item that could never fit is refused by name and ends the reader.
//! Why the reader stopped is recorded here as its last act, and the consumer
//! receives it once every item before it has been taken.

use alloc::collections::VecDeque;
use alloc::format;
use alloc::string::ToString;
use core::time::Duration;

use std::sync::{Condvar, Mutex, MutexGuard};
use std::time::Instant;

use crate::SdkError;

/// Default cap on frames a push client holds for a consumer that has not
/// received them.
pub const DEFAULT_PUSH_MAILBOX_FRAMES: usize = 1024;

/// Default cap on payload bytes a push client holds. Equal to the push
/// reader's single-frame buffering bound, so every frame it can decode fits
/// an empty mailbox.
pub const DEFAULT_PUSH_MAILBOX_BYTES: usize = 16 * 1024 * 1024;

/// Default cap on deliveries a subscription stream holds for a consumer that
/// has not received them.
pub const DEFAULT_SUBSCRIPTION_MAILBOX_FRAMES: usize = 1024;

/// Default cap on delivered payload bytes a subscription stream holds.
///
/// Equal to the TCP subscription reader's single-frame buffering bound, so
/// every delivery it can decode fits an empty mailbox; the WebSocket stream
/// shares it so both transports hold the same.
pub const DEFAULT_SUBSCRIPTION_MAILBOX_BYTES: usize = 64 * 1024 * 1024;

/// One item a mailbox can hold: it reports the payload bytes it keeps alive.
pub trait MailboxItem {
    fn payload_len(&self) -> usize;
}

/// The bounds one mailbox enforces.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MailboxLimits {
    pub frames: usize,
    pub bytes: usize,
}

impl MailboxLimits {
    pub const PUSH: Self = Self {
        frames: DEFAULT_PUSH_MAILBOX_FRAMES,
        bytes: DEFAULT_PUSH_MAILBOX_BYTES,
    };
    pub const SUBSCRIPTION: Self = Self {
        frames: DEFAULT_SUBSCRIPTION_MAILBOX_FRAMES,
        bytes: DEFAULT_SUBSCRIPTION_MAILBOX_BYTES,
    };
}

/// The words a mailbox's consumer-facing errors use, so each stream keeps
/// its own wording.
#[derive(Clone, Copy, Debug)]
pub struct MailboxWords {
    /// Prefix of every receive failure, e.g. "push receive failed".
    pub receive: &'static str,
    /// What the consumer waits for, e.g. "server push" or "delivery".
    pub item: &'static str,
    /// The reader's name, e.g. "push reader".
    pub reader: &'static str,
}

/// Counters describing one stream's mailbox, readable at any time from the
/// stream that owns it.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct MailboxStats {
    /// Items held now.
    pub held_frames: usize,
    /// Payload bytes held now.
    pub held_bytes: usize,
    /// Most items ever held at once.
    pub peak_frames: usize,
    /// Most payload bytes ever held at once.
    pub peak_bytes: usize,
    /// Times the reader found the mailbox full and stopped reading the socket
    /// until the consumer made room.
    pub backpressure_waits: u64,
    /// Items refused because they could never fit the mailbox.
    pub refused_frames: u64,
    /// Frames of a kind this stream does not carry, read and counted rather
    /// than surfaced, so a stray frame cannot stop the frames behind it.
    pub ignored_frames: u64,
    /// Reader errors (a server-ended subscription, a read or decode failure)
    /// that no consumer was left to receive.
    pub undelivered_errors: u64,
}

/// Why the reader stopped; surfaced to the consumer after every held item.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReaderEnd {
    /// The reader ended with this error: the connection closed, a read
    /// failed, a frame did not decode, the server ended the stream, or an item
    /// was refused.
    Failed(SdkError),
    /// The peer closed a complete stream; the error preserves timed reception.
    PeerClosed(SdkError),
    /// The reader stopped with no error of its own: its stream was dropped,
    /// or the link ended without one.
    Stopped,
}

/// An item larger than the mailbox could ever hold.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MailboxOverflow {
    /// The refused item's payload bytes.
    pub payload_bytes: usize,
    /// The mailbox's byte bound.
    pub max_bytes: usize,
}

impl core::fmt::Display for MailboxOverflow {
    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(
            formatter,
            "mailbox overflow: a {}-byte item exceeds the {}-byte mailbox",
            self.payload_bytes, self.max_bytes
        )
    }
}

impl From<MailboxOverflow> for SdkError {
    fn from(overflow: MailboxOverflow) -> Self {
        Self::Backpressure {
            reason: format!("{overflow}"),
        }
    }
}

#[derive(Debug)]
struct State<T> {
    queue: VecDeque<T>,
    stats: MailboxStats,
    end: Option<ReaderEnd>,
    abandoned: bool,
    receiver_claimed: bool,
    direct_receives: usize,
    #[cfg(test)]
    blocking_wait_observer: Option<std::sync::mpsc::Sender<()>>,
    #[cfg(test)]
    receive_wakes: u64,
    #[cfg(test)]
    receive_timers: u64,
}

/// The reader's view of one offer.
#[derive(Debug, PartialEq, Eq)]
pub enum Offer {
    /// Held for the consumer, or discarded because the stream was dropped.
    Accepted,
    /// The item could never fit; the reader must end with this refusal.
    Refused(MailboxOverflow),
}

/// One blocking receive result, retaining the terminal's typed provenance.
#[derive(Debug)]
pub enum BlockingTake<T> {
    /// The next queued item, moved rather than cloned.
    Item(T),
    /// The owning stream stopped accepting reception.
    Abandoned,
    /// The reader's retained terminal after its queue drained.
    Ended(ReaderEnd),
}

/// The shared mailbox. One lock, one condition variable: the reader waits on
/// it when full, the consumer when empty, teardown until the reader ends.
#[derive(Debug)]
pub struct Mailbox<T> {
    limits: MailboxLimits,
    words: MailboxWords,
    state: Mutex<State<T>>,
    changed: Condvar,
}

fn poisoned() -> SdkError {
    SdkError::Connection {
        description: "stream mailbox lock poisoned: a thread panicked while holding it".to_string(),
    }
}

impl<T: MailboxItem> Mailbox<T> {
    pub fn new(limits: MailboxLimits, words: MailboxWords) -> Self {
        Self {
            limits,
            words,
            state: Mutex::new(State {
                queue: VecDeque::new(),
                stats: MailboxStats::default(),
                end: None,
                abandoned: false,
                receiver_claimed: false,
                direct_receives: 0,
                #[cfg(test)]
                blocking_wait_observer: None,
                #[cfg(test)]
                receive_wakes: 0,
                #[cfg(test)]
                receive_timers: 0,
            }),
            changed: Condvar::new(),
        }
    }

    fn lock(&self) -> Result<MutexGuard<'_, State<T>>, SdkError> {
        self.state.lock().map_err(|_| poisoned())
    }

    /// Offers one item from the reader, waiting while the mailbox is full.
    ///
    /// # Errors
    /// Returns [`SdkError::Connection`] when the lock is poisoned.
    pub fn offer(&self, item: T) -> Result<Offer, SdkError> {
        let size = item.payload_len();
        if size > self.limits.bytes {
            let mut state = self.lock()?;
            state.stats.refused_frames = state.stats.refused_frames.saturating_add(1);
            drop(state);
            return Ok(Offer::Refused(MailboxOverflow {
                payload_bytes: size,
                max_bytes: self.limits.bytes,
            }));
        }
        let mut state = self.lock()?;
        let mut counted = false;
        while !state.abandoned
            && (state.queue.len() >= self.limits.frames
                || state.stats.held_bytes.saturating_add(size) > self.limits.bytes)
        {
            if !counted {
                state.stats.backpressure_waits = state.stats.backpressure_waits.saturating_add(1);
                counted = true;
                self.changed.notify_all();
            }
            state = self.changed.wait(state).map_err(|_| poisoned())?;
        }
        if state.abandoned {
            drop(state);
            return Ok(Offer::Accepted);
        }
        state.queue.push_back(item);
        let tally = &mut state.stats;
        tally.held_frames = tally.held_frames.saturating_add(1);
        tally.held_bytes = tally.held_bytes.saturating_add(size);
        tally.peak_frames = tally.peak_frames.max(tally.held_frames);
        tally.peak_bytes = tally.peak_bytes.max(tally.held_bytes);
        drop(state);
        self.changed.notify_all();
        Ok(Offer::Accepted)
    }

    /// Records why the reader stopped and wakes every waiter. Called exactly
    /// once, as the reader's last act. An error with no consumer left to
    /// receive it is counted.
    pub fn close(&self, end: ReaderEnd) {
        if let Ok(mut state) = self.state.lock() {
            if state.abandoned && matches!(end, ReaderEnd::Failed(_) | ReaderEnd::PeerClosed(_)) {
                state.stats.undelivered_errors = state.stats.undelivered_errors.saturating_add(1);
            }
            state.end = Some(end);
        }
        // A poisoned lock leaves `end` unset; waiters still wake here and
        // observe the poison themselves on their next lock.
        self.changed.notify_all();
    }

    /// Takes the next item, waiting up to `timeout` for one.
    ///
    /// # Errors
    /// Returns the reader's end once every held item has been taken, a
    /// [`SdkError::Connection`] naming the timeout, or the poisoned lock.
    pub fn take(&self, timeout: Duration) -> Result<T, SdkError> {
        let deadline = Instant::now().checked_add(timeout);
        let mut state = self.lock()?;
        loop {
            if let Some(item) = state.queue.pop_front() {
                let tally = &mut state.stats;
                tally.held_frames = tally.held_frames.saturating_sub(1);
                tally.held_bytes = tally.held_bytes.saturating_sub(item.payload_len());
                drop(state);
                self.changed.notify_all();
                return Ok(item);
            }
            match &state.end {
                Some(ReaderEnd::Failed(error) | ReaderEnd::PeerClosed(error)) => {
                    return Err(error.clone());
                }
                Some(ReaderEnd::Stopped) => {
                    return Err(SdkError::Connection {
                        description: format!(
                            "{}: the {} stopped before a {} arrived",
                            self.words.receive, self.words.reader, self.words.item
                        ),
                    });
                }
                None => {}
            }
            let remaining = deadline.map_or(timeout, |deadline| {
                deadline.saturating_duration_since(Instant::now())
            });
            if remaining.is_zero() {
                return Err(SdkError::Connection {
                    description: format!(
                        "{}: no {} arrived within the timeout",
                        self.words.receive, self.words.item
                    ),
                });
            }
            #[cfg(test)]
            {
                state.receive_timers = state.receive_timers.saturating_add(1);
            }
            state = self
                .changed
                .wait_timeout(state, remaining)
                .map_err(|_| poisoned())?
                .0;
        }
    }

    /// Claims the single push receiver while no direct receive is active.
    ///
    /// # Errors
    /// Returns [`SdkError::ReceiverAlreadyClaimed`] when the receiver is
    /// claimed or a direct receive is active, and
    /// [`SdkError::Connection`] when the lock is poisoned.
    pub fn claim_receiver(&self) -> Result<(), SdkError> {
        let mut state = self.lock()?;
        let claimable = !state.receiver_claimed && state.direct_receives == 0;
        if claimable {
            state.receiver_claimed = true;
        }
        drop(state);
        if claimable {
            Ok(())
        } else {
            Err(SdkError::ReceiverAlreadyClaimed)
        }
    }

    fn begin_direct(state: &mut State<T>) -> Result<(), SdkError> {
        if state.receiver_claimed {
            return Err(SdkError::ReceiverAlreadyClaimed);
        }
        state.direct_receives =
            state
                .direct_receives
                .checked_add(1)
                .ok_or_else(|| SdkError::Protocol {
                    description: "push direct receive count exhausted".to_string(),
                })?;
        Ok(())
    }

    /// Preserves timed reception while fencing receiver-handle ownership.
    pub fn take_push_timeout(&self, timeout: Duration) -> Result<T, SdkError> {
        let mut state = self.lock()?;
        Self::begin_direct(&mut state)?;
        drop(state);
        let result = self.take(timeout);
        let mut state = self.lock()?;
        state.direct_receives -= 1;
        result
    }

    /// Waits on the locked arrival or terminal predicate without a timer.
    pub fn take_blocking(&self, direct: bool) -> Result<BlockingTake<T>, SdkError> {
        let mut state = self.lock()?;
        if direct {
            Self::begin_direct(&mut state)?;
        }
        let result = loop {
            if state.abandoned {
                break Ok(BlockingTake::Abandoned);
            }
            if let Some(item) = state.queue.pop_front() {
                state.stats.held_frames -= 1;
                state.stats.held_bytes -= item.payload_len();
                break Ok(BlockingTake::Item(item));
            }
            if let Some(end) = &state.end {
                break Ok(BlockingTake::Ended(end.clone()));
            }
            #[cfg(test)]
            if let Some(observer) = &state.blocking_wait_observer {
                if let Err(error) = observer.send(()) {
                    break Err(SdkError::Connection {
                        description: format!("blocking receive observer failed: {error}"),
                    });
                }
            }
            state = self.changed.wait(state).map_err(|_| poisoned())?;
            #[cfg(test)]
            {
                state.receive_wakes = state.receive_wakes.saturating_add(1);
            }
        };
        if direct {
            state.direct_receives -= 1;
        }
        let made_room = matches!(&result, Ok(BlockingTake::Item(_)));
        drop(state);
        if made_room {
            self.changed.notify_all();
        }
        result
    }

    #[cfg(test)]
    pub fn observe_blocking_wait(
        &self,
        observer: std::sync::mpsc::Sender<()>,
    ) -> Result<(), SdkError> {
        self.lock()?.blocking_wait_observer = Some(observer);
        Ok(())
    }

    #[cfg(test)]
    pub fn receive_wait_counts(&self) -> Result<(u64, u64), SdkError> {
        let state = self.lock()?;
        Ok((state.receive_wakes, state.receive_timers))
    }

    /// Discards everything held and every later item, and releases a reader
    /// waiting on a full mailbox: the consumer is going away.
    ///
    /// # Errors
    /// Returns [`SdkError::Connection`] when the lock is poisoned.
    pub fn abandon(&self) -> Result<(), SdkError> {
        let mut state = self.lock()?;
        state.abandoned = true;
        state.queue.clear();
        state.stats.held_frames = 0;
        state.stats.held_bytes = 0;
        drop(state);
        self.changed.notify_all();
        Ok(())
    }

    /// [`Self::abandon`] for teardown paths with no caller to report to. On a
    /// poisoned lock the waiters are still woken, and a reader waiting to
    /// offer then observes the poison itself and ends.
    pub fn release(&self) {
        if let Ok(mut state) = self.state.lock() {
            state.abandoned = true;
            state.queue.clear();
            state.stats.held_frames = 0;
            state.stats.held_bytes = 0;
        }
        self.changed.notify_all();
    }

    /// Waits until `done` holds for the counters and reader end, or `budget`
    /// passes; reports whether it held.
    ///
    /// # Errors
    /// Returns [`SdkError::Connection`] when the lock is poisoned.
    pub fn wait_until(
        &self,
        budget: Duration,
        done: impl Fn(&MailboxStats, Option<&ReaderEnd>) -> bool,
    ) -> Result<bool, SdkError> {
        let deadline = Instant::now().checked_add(budget);
        let mut state = self.lock()?;
        loop {
            if done(&state.stats, state.end.as_ref()) {
                return Ok(true);
            }
            let remaining = deadline.map_or(budget, |deadline| {
                deadline.saturating_duration_since(Instant::now())
            });
            if remaining.is_zero() {
                return Ok(false);
            }
            state = self
                .changed
                .wait_timeout(state, remaining)
                .map_err(|_| poisoned())?
                .0;
        }
    }

    /// Counts one frame the reader read and set aside.
    ///
    /// # Errors
    /// Returns [`SdkError::Connection`] when the lock is poisoned.
    pub fn record_ignored_frame(&self) -> Result<(), SdkError> {
        let mut state = self.lock()?;
        state.stats.ignored_frames = state.stats.ignored_frames.saturating_add(1);
        drop(state);
        Ok(())
    }

    /// The counters as they stand.
    ///
    /// # Errors
    /// Returns [`SdkError::Connection`] when the lock is poisoned.
    pub fn stats(&self) -> Result<MailboxStats, SdkError> {
        Ok(self.lock()?.stats)
    }
}

/// Offers `item`, mapping the outcome onto the reader's loop: `None` to keep
/// reading, or the end the reader must stop with.
pub fn offer_or_end<T: MailboxItem>(mailbox: &Mailbox<T>, item: T) -> Option<ReaderEnd> {
    match mailbox.offer(item) {
        Ok(Offer::Accepted) => None,
        Ok(Offer::Refused(overflow)) => Some(ReaderEnd::Failed(SdkError::from(overflow))),
        Err(error) => Some(ReaderEnd::Failed(error)),
    }
}