use alloc::format;
use alloc::string::String;
use alloc::vec::Vec;
use core::time::Duration;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel};
use std::time::Instant;
use crate::SdkError;
pub const FLUSH_BUDGET: Duration = Duration::from_secs(5);
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PublishRejection {
reason_code: u16,
message: Option<String>,
}
impl PublishRejection {
pub(crate) const fn new(reason_code: u16, message: Option<String>) -> Self {
Self {
reason_code,
message,
}
}
#[must_use]
pub const fn reason_code(&self) -> u16 {
self.reason_code
}
#[must_use]
pub fn message(&self) -> Option<&str> {
self.message.as_deref()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FlushMode {
FlushedAndHalfClosed,
VerdictOnly,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FlushOutcome {
failures: Vec<PublishRejection>,
unresolved: usize,
mode: FlushMode,
}
impl FlushOutcome {
pub(crate) const fn new(
failures: Vec<PublishRejection>,
unresolved: usize,
mode: FlushMode,
) -> Self {
Self {
failures,
unresolved,
mode,
}
}
#[must_use]
pub fn failures(&self) -> &[PublishRejection] {
&self.failures
}
#[must_use]
pub fn into_failures(self) -> Vec<PublishRejection> {
self.failures
}
#[must_use]
pub const fn unresolved(&self) -> usize {
self.unresolved
}
#[must_use]
pub const fn mode(&self) -> FlushMode {
self.mode
}
#[must_use]
pub fn is_proven_accepted(&self) -> bool {
self.failures.is_empty() && self.unresolved == 0
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum PublishVerdict {
Accepted,
Rejected(PublishRejection),
}
struct VerdictInbox {
verdicts: Receiver<PublishVerdict>,
resolved: u64,
}
#[derive(Debug)]
pub struct FlushLedger {
written: AtomicU64,
arrived: AtomicU64,
inbox: Mutex<VerdictInbox>,
}
impl core::fmt::Debug for VerdictInbox {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
formatter
.debug_struct("VerdictInbox")
.field("resolved", &self.resolved)
.finish_non_exhaustive()
}
}
impl FlushLedger {
pub(crate) fn new() -> (Self, Sender<PublishVerdict>) {
let (sender, verdicts) = channel();
let ledger = Self {
written: AtomicU64::new(0),
arrived: AtomicU64::new(0),
inbox: Mutex::new(VerdictInbox {
verdicts,
resolved: 0,
}),
};
(ledger, sender)
}
pub(crate) fn record_written(&self) {
self.written.fetch_add(1, Ordering::SeqCst);
}
pub(crate) fn record_arrival(&self) {
self.arrived.fetch_add(1, Ordering::SeqCst);
}
pub(crate) fn drain(
&self,
budget: Duration,
) -> Result<(Vec<PublishRejection>, usize), SdkError> {
let written_at_call = self.written.load(Ordering::SeqCst);
let mut inbox = self.inbox.lock().map_err(|error| SdkError::Connection {
description: format!("flush guard poisoned: {error}"),
})?;
self.check_pairing_invariant()?;
let window = written_at_call.saturating_sub(inbox.resolved);
let deadline = Instant::now() + budget;
let mut failures = Vec::new();
let mut collected: u64 = 0;
while collected < window {
let now = Instant::now();
if now >= deadline {
break;
}
match inbox.verdicts.recv_timeout(deadline.duration_since(now)) {
Ok(verdict) => {
collected += 1;
inbox.resolved += 1;
if let PublishVerdict::Rejected(rejection) = verdict {
failures.push(rejection);
}
}
Err(RecvTimeoutError::Timeout | RecvTimeoutError::Disconnected) => break,
}
}
drop(inbox);
self.check_pairing_invariant()?;
let unresolved = usize::try_from(window.saturating_sub(collected)).map_err(|error| {
SdkError::Protocol {
description: format!("unresolved publish count overflowed usize: {error}"),
}
})?;
Ok((failures, unresolved))
}
fn check_pairing_invariant(&self) -> Result<(), SdkError> {
let arrived = self.arrived.load(Ordering::SeqCst);
let written = self.written.load(Ordering::SeqCst);
if arrived > written {
return Err(SdkError::Protocol {
description: format!(
"publish response-count mismatch: {arrived} publish responses arrived for \
{written} response-eliciting publishes; refusing to pair verdicts"
),
});
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::string::ToString;
use alloc::vec;
const SHORT_BUDGET: Duration = Duration::from_millis(50);
#[test]
fn budget_expiry_counts_unresolved_and_is_not_an_error() -> Result<(), SdkError> {
let (ledger, _sender) = FlushLedger::new();
ledger.record_written();
ledger.record_written();
ledger.record_written();
let (failures, unresolved) = ledger.drain(SHORT_BUDGET)?;
assert!(failures.is_empty());
assert_eq!(unresolved, 3);
let expired = FlushOutcome::new(failures, unresolved, FlushMode::VerdictOnly);
assert!(!expired.is_proven_accepted());
let clean = FlushOutcome::new(Vec::new(), 0, FlushMode::VerdictOnly);
assert!(clean.is_proven_accepted());
Ok(())
}
#[test]
fn surplus_verdict_is_a_typed_mechanism_error() {
let (ledger, sender) = FlushLedger::new();
assert!(sender.send(PublishVerdict::Accepted).is_ok());
ledger.record_arrival();
let result = ledger.drain(SHORT_BUDGET);
assert!(matches!(result, Err(SdkError::Protocol { .. })));
}
#[test]
fn rejections_surface_in_wire_order_verbatim() -> Result<(), SdkError> {
let (ledger, sender) = FlushLedger::new();
let first = PublishRejection::new(0xFFFF, Some("first".to_string()));
let second = PublishRejection::new(0xFFFF, None);
for verdict in [
PublishVerdict::Rejected(first.clone()),
PublishVerdict::Accepted,
PublishVerdict::Rejected(second.clone()),
] {
ledger.record_written();
ledger.record_arrival();
assert!(sender.send(verdict).is_ok());
}
let (failures, unresolved) = ledger.drain(SHORT_BUDGET)?;
assert_eq!(failures, vec![first, second]);
assert_eq!(unresolved, 0);
Ok(())
}
#[test]
fn concurrent_drains_serialize_without_misattribution() -> Result<(), SdkError> {
use std::sync::{Arc, Barrier};
let (ledger, sender) = FlushLedger::new();
let ledger = Arc::new(ledger);
let rejection = PublishRejection::new(0xFFFF, Some("boom".to_string()));
for verdict in [
PublishVerdict::Accepted,
PublishVerdict::Accepted,
PublishVerdict::Rejected(rejection.clone()),
PublishVerdict::Accepted,
] {
ledger.record_written();
ledger.record_arrival();
assert!(sender.send(verdict).is_ok());
}
let barrier = Arc::new(Barrier::new(2));
let spawn_drain = |ledger: Arc<FlushLedger>, barrier: Arc<Barrier>| {
std::thread::spawn(move || {
barrier.wait();
ledger.drain(FLUSH_BUDGET)
})
};
let first = spawn_drain(Arc::clone(&ledger), Arc::clone(&barrier));
let second = spawn_drain(Arc::clone(&ledger), barrier);
let joined =
|handle: std::thread::JoinHandle<Result<(Vec<PublishRejection>, usize), SdkError>>| {
handle.join().map_err(|_| SdkError::Protocol {
description: "drain thread panicked".to_string(),
})
};
let (failures_a, unresolved_a) = joined(first)??;
let (failures_b, unresolved_b) = joined(second)??;
let combined: Vec<PublishRejection> = failures_a.into_iter().chain(failures_b).collect();
assert_eq!(combined, vec![rejection]);
assert_eq!(unresolved_a, 0);
assert_eq!(unresolved_b, 0);
Ok(())
}
#[test]
fn second_drain_covers_only_its_own_boundary() -> Result<(), SdkError> {
let (ledger, sender) = FlushLedger::new();
ledger.record_written();
ledger.record_arrival();
assert!(sender.send(PublishVerdict::Accepted).is_ok());
let (_, unresolved) = ledger.drain(SHORT_BUDGET)?;
assert_eq!(unresolved, 0);
let started = Instant::now();
let (failures, unresolved) = ledger.drain(FLUSH_BUDGET)?;
assert!(failures.is_empty());
assert_eq!(unresolved, 0);
assert!(started.elapsed() < FLUSH_BUDGET);
Ok(())
}
}