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;
pub const DEFAULT_PUSH_MAILBOX_FRAMES: usize = 1024;
pub const DEFAULT_PUSH_MAILBOX_BYTES: usize = 16 * 1024 * 1024;
pub const DEFAULT_SUBSCRIPTION_MAILBOX_FRAMES: usize = 1024;
pub const DEFAULT_SUBSCRIPTION_MAILBOX_BYTES: usize = 64 * 1024 * 1024;
pub trait MailboxItem {
fn payload_len(&self) -> usize;
}
#[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,
};
}
#[derive(Clone, Copy, Debug)]
pub struct MailboxWords {
pub receive: &'static str,
pub item: &'static str,
pub reader: &'static str,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct MailboxStats {
pub held_frames: usize,
pub held_bytes: usize,
pub peak_frames: usize,
pub peak_bytes: usize,
pub backpressure_waits: u64,
pub refused_frames: u64,
pub ignored_frames: u64,
pub undelivered_errors: u64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReaderEnd {
Failed(SdkError),
Stopped,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MailboxOverflow {
pub payload_bytes: usize,
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,
}
#[derive(Debug, PartialEq, Eq)]
pub enum Offer {
Accepted,
Refused(MailboxOverflow),
}
#[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,
}),
changed: Condvar::new(),
}
}
fn lock(&self) -> Result<MutexGuard<'_, State<T>>, SdkError> {
self.state.lock().map_err(|_| 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)
}
pub fn close(&self, end: ReaderEnd) {
if let Ok(mut state) = self.state.lock() {
if state.abandoned && matches!(end, ReaderEnd::Failed(_)) {
state.stats.undelivered_errors = state.stats.undelivered_errors.saturating_add(1);
}
state.end = Some(end);
}
self.changed.notify_all();
}
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)) => 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
),
});
}
state = self
.changed
.wait_timeout(state, remaining)
.map_err(|_| poisoned())?
.0;
}
}
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(())
}
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();
}
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;
}
}
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(())
}
pub fn stats(&self) -> Result<MailboxStats, SdkError> {
Ok(self.lock()?.stats)
}
}
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)),
}
}