armdb 0.11.0

#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
use std::collections::{HashSet, VecDeque};
use std::sync::Arc;
#[cfg(all(
    debug_assertions,
    feature = "queue-test-support",
    feature = "rapira-codec",
    feature = "var-collections"
))]
use std::sync::Barrier;
#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
use std::sync::Weak;
#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
use std::sync::atomic::AtomicBool;
#[cfg(all(
    debug_assertions,
    feature = "queue-test-support",
    feature = "rapira-codec",
    feature = "var-collections"
))]
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::{AtomicU64, Ordering};
#[cfg(all(debug_assertions, feature = "queue-test-support"))]
use std::sync::mpsc::{Receiver, Sender, channel};

#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
use armour_core::GetType;

#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
use crate::DbError;
use crate::hook::TypedWriteHook;
use crate::sync::{Condvar, Mutex, cv_wait, lock};
use crate::{DbResult, QueueKey, QueueSpec};

#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
use super::var::{QueuePayloadTree, QueueVarHeaderTree};

#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
const RETRY_CAPACITY: usize = 1024;

#[derive(Default)]
struct PublishState {
    producers: usize,
    maintenance: bool,
}

pub(crate) struct PublishGate {
    state: Mutex<PublishState>,
    changed: Condvar,
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    maintenance_waiting: Mutex<Option<Sender<()>>>,
}
impl PublishGate {
    pub(crate) fn new() -> Arc<Self> {
        Arc::new(Self {
            state: Mutex::new(PublishState::default()),
            changed: Condvar::new(),
            #[cfg(all(debug_assertions, feature = "queue-test-support"))]
            maintenance_waiting: Mutex::new(None),
        })
    }
    pub(crate) fn producer(self: &Arc<Self>) -> ProducerPermit {
        let mut state = lock(&self.state);
        while state.maintenance {
            state = cv_wait(&self.changed, state);
        }
        state.producers += 1;
        ProducerPermit(Arc::clone(self))
    }
    fn try_maintenance(self: &Arc<Self>) -> Option<MaintenancePermit> {
        let mut state = lock(&self.state);
        if state.maintenance || state.producers != 0 {
            #[cfg(all(debug_assertions, feature = "queue-test-support"))]
            if state.producers != 0
                && let Some(sender) = lock(&self.maintenance_waiting).take()
            {
                let _ = sender.send(());
            }
            return None;
        }
        state.maintenance = true;
        Some(MaintenancePermit(Arc::clone(self)))
    }
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    fn notify_maintenance_wait(&self) -> Receiver<()> {
        let (sender, receiver) = channel();
        *lock(&self.maintenance_waiting) = Some(sender);
        receiver
    }
}
#[doc(hidden)]
pub struct ProducerPermit(Arc<PublishGate>);
impl Drop for ProducerPermit {
    fn drop(&mut self) {
        let mut state = lock(&self.0.state);
        state.producers -= 1;
        self.0.changed.notify_all();
    }
}
struct MaintenancePermit(Arc<PublishGate>);
impl Drop for MaintenancePermit {
    fn drop(&mut self) {
        let mut state = lock(&self.0.state);
        state.maintenance = false;
        self.0.changed.notify_all();
    }
}

pub(crate) trait QueueRetentionPass: Send + Sync {
    fn run_retention(&self) -> DbResult<()>;
}

pub(crate) struct HeaderRetention {
    dead_count: Arc<AtomicU64>,
}

impl HeaderRetention {
    pub(crate) fn new() -> Arc<Self> {
        Arc::new(Self {
            dead_count: Arc::new(AtomicU64::new(0)),
        })
    }
    pub(crate) fn dead_count(&self) -> Arc<AtomicU64> {
        Arc::clone(&self.dead_count)
    }
    pub(crate) fn set_dead_count(&self, count: u64) {
        self.dead_count.store(count, Ordering::Release);
    }
    pub(crate) fn deleted_dead_header(&self, key: &QueueKey) {
        if key.class() == crate::DEAD_CLASS {
            let _ = self
                .dead_count
                .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
                    Some(count.saturating_sub(1))
                });
        }
    }
}

#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
pub(crate) struct VarRetention<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    headers: Mutex<Weak<QueueVarHeaderTree<N>>>,
    payloads: Mutex<Weak<QueuePayloadTree<N>>>,
    retries: Mutex<VecDeque<u64>>,
    overflow: AtomicBool,
    gate: Arc<PublishGate>,
    pub(crate) header: Arc<HeaderRetention>,
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    fail_payload_deletes: AtomicUsize,
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    push_pause: Mutex<Option<(Arc<Barrier>, Arc<Barrier>)>>,
}

#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
impl<N> VarRetention<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    pub(crate) fn new(header: Arc<HeaderRetention>) -> Arc<Self> {
        Arc::new(Self {
            headers: Mutex::new(Weak::new()),
            payloads: Mutex::new(Weak::new()),
            retries: Mutex::new(VecDeque::new()),
            overflow: AtomicBool::new(false),
            gate: PublishGate::new(),
            header,
            #[cfg(all(debug_assertions, feature = "queue-test-support"))]
            fail_payload_deletes: AtomicUsize::new(0),
            #[cfg(all(debug_assertions, feature = "queue-test-support"))]
            push_pause: Mutex::new(None),
        })
    }
    pub(crate) fn bind(
        &self,
        headers: &Arc<QueueVarHeaderTree<N>>,
        payloads: &Arc<QueuePayloadTree<N>>,
    ) {
        *lock(&self.headers) = Arc::downgrade(headers);
        *lock(&self.payloads) = Arc::downgrade(payloads);
    }
    pub(crate) fn producer_permit(&self) -> ProducerPermit {
        self.gate.producer()
    }
    #[cfg(test)]
    pub(crate) fn producer_count(&self) -> usize {
        lock(&self.gate.state).producers
    }
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    pub(crate) fn notify_maintenance_wait(&self) -> Receiver<()> {
        self.gate.notify_maintenance_wait()
    }
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    pub(crate) fn pause_after_payload(&self) -> (Arc<Barrier>, Arc<Barrier>) {
        let pause = (Arc::new(Barrier::new(2)), Arc::new(Barrier::new(2)));
        *lock(&self.push_pause) = Some((Arc::clone(&pause.0), Arc::clone(&pause.1)));
        pause
    }
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    pub(crate) fn wait_after_payload(&self) {
        if let Some((written, release)) = lock(&self.push_pause).clone() {
            written.wait();
            release.wait();
        }
    }
    fn payloads(&self) -> DbResult<Arc<QueuePayloadTree<N>>> {
        lock(&self.payloads)
            .upgrade()
            .ok_or(DbError::Internal("queue payload retention is not bound"))
    }
    fn delete_payload(&self, id: u64) -> DbResult<()> {
        #[cfg(all(debug_assertions, feature = "queue-test-support"))]
        if self
            .fail_payload_deletes
            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| {
                remaining.checked_sub(1)
            })
            .is_ok()
        {
            return Err(DbError::Internal("debug queue payload delete failure"));
        }
        self.payloads()?
            .delete(&rend::u64_be::from_native(id))
            .map(|_| ())
    }
    fn retry_later(&self, id: u64) {
        let mut retries = lock(&self.retries);
        if retries.len() < RETRY_CAPACITY {
            retries.push_back(id);
        } else {
            self.overflow.store(true, Ordering::Release);
        }
    }
    pub(crate) fn deleted_dead_header(&self, key: &QueueKey) {
        self.header.deleted_dead_header(key);
        if key.class() == crate::DEAD_CLASS && self.delete_payload(key.id()).is_err() {
            self.retry_later(key.id());
        }
    }
    fn drain_retries(&self) -> DbResult<()> {
        let retries = std::mem::take(&mut *lock(&self.retries));
        for id in retries {
            if let Err(error) = self.delete_payload(id) {
                self.retry_later(id);
                // The temporary drain buffer still owns the unprocessed tail.
                // An orphan pass is the bounded, idempotent recovery path for it.
                self.overflow.store(true, Ordering::Release);
                return Err(error);
            }
        }
        Ok(())
    }
    pub(crate) fn orphan_pass_pair(
        headers: &Arc<QueueVarHeaderTree<N>>,
        payloads: &Arc<QueuePayloadTree<N>>,
    ) -> DbResult<()> {
        let ids: HashSet<_> = headers.iter().map(|(key, _)| key.id()).collect();
        let orphaned: Vec<_> = payloads
            .iter()
            .map(|(key, _)| key)
            .filter(|key| !ids.contains(&key.to_native()))
            .collect();
        for key in orphaned {
            payloads.delete(&key)?;
        }
        Ok(())
    }
    fn orphan_pass(&self) -> DbResult<()> {
        let headers = lock(&self.headers)
            .upgrade()
            .ok_or(DbError::Internal("queue header retention is not bound"))?;
        Self::orphan_pass_pair(&headers, &self.payloads()?)
    }
    /// Reclaim what is reclaimable before a queue operation takes its permit.
    ///
    /// Deliberately infallible. Retention frees bytes belonging to rows that
    /// are already gone, so its failure is never urgent — propagating it would
    /// make one transient payload-delete error reject every `push`, `lease` and
    /// `nack` on the queue until the delete happened to succeed. `run_retention`
    /// re-arms `overflow` on every failure path, so the work is only deferred,
    /// and `Db::sweep_expired_now` still surfaces the error to the sweeper.
    pub(crate) fn process_before_queue_operation(&self) {
        if let Err(error) = self.run_retention() {
            tracing::warn!(
                queue = N::NAME,
                %error,
                "queue payload retention pass failed; retrying on a later pass"
            );
        }
    }
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    pub(crate) fn fail_next_payload_delete(&self) {
        self.fail_payload_deletes.store(1, Ordering::Release);
    }
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    pub(crate) fn fail_payload_deletes(&self, count: usize) {
        self.fail_payload_deletes.store(count, Ordering::Release);
    }
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    pub(crate) fn force_full_orphan_pass(&self) -> DbResult<()> {
        self.overflow.store(true, Ordering::Release);
        self.run_retention()
    }
    #[cfg(all(debug_assertions, feature = "queue-test-support"))]
    pub(crate) fn mark_full_orphan_pass_pending(&self) {
        self.overflow.store(true, Ordering::Release);
    }
}
#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
impl<N> QueueRetentionPass for VarRetention<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    fn run_retention(&self) -> DbResult<()> {
        if lock(&self.retries).is_empty() && !self.overflow.load(Ordering::Acquire) {
            return Ok(());
        }
        let Some(_maintenance) = self.gate.try_maintenance() else {
            // A producer may have written a payload that is not visible through
            // its header yet. Retaining the overflow flag schedules a safe
            // retry after that producer (including a post-log recovery) drops.
            self.overflow.store(true, Ordering::Release);
            return Ok(());
        };
        self.drain_retries()?;
        if self.overflow.swap(false, Ordering::AcqRel)
            && let Err(error) = self.orphan_pass()
        {
            self.overflow.store(true, Ordering::Release);
            return Err(error);
        }
        Ok(())
    }
}

pub(crate) struct FixedHeaderHook<N: QueueSpec> {
    retention: Arc<HeaderRetention>,
    _queue: std::marker::PhantomData<fn() -> N>,
}
impl<N> FixedHeaderHook<N>
where
    N: QueueSpec,
    N::Payload: armour_core::GetType,
{
    pub(crate) fn new(retention: Arc<HeaderRetention>) -> Self {
        Self {
            retention,
            _queue: std::marker::PhantomData,
        }
    }
}
impl<N> TypedWriteHook<QueueKey, super::fixed::QueueFixedRow<N>> for FixedHeaderHook<N>
where
    N: QueueSpec,
    N::Payload: armour_core::GetType,
{
    const NEEDS_OLD_VALUE: bool = true;
    fn on_write(
        &self,
        key: &QueueKey,
        old: Option<&super::fixed::QueueFixedRow<N>>,
        new: Option<&super::fixed::QueueFixedRow<N>>,
    ) {
        if old.is_some() && new.is_none() {
            self.retention.deleted_dead_header(key);
        }
    }
}

#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
pub(crate) struct VarHeaderHook<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    retention: Arc<VarRetention<N>>,
}
#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
impl<N> VarHeaderHook<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    pub(crate) fn new(retention: Arc<VarRetention<N>>) -> Self {
        Self { retention }
    }
}
#[cfg(all(feature = "rapira-codec", feature = "var-collections"))]
impl<N> TypedWriteHook<QueueKey, super::var::QueueVarHeaderRow<N>> for VarHeaderHook<N>
where
    N: QueueSpec,
    N::Payload: GetType,
{
    const NEEDS_OLD_VALUE: bool = true;
    fn on_write(
        &self,
        key: &QueueKey,
        old: Option<&super::var::QueueVarHeaderRow<N>>,
        new: Option<&super::var::QueueVarHeaderRow<N>>,
    ) {
        if old.is_some() && new.is_none() {
            self.retention.deleted_dead_header(key);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn deleting_an_unaccounted_dead_header_does_not_wrap_the_count() {
        let retention = HeaderRetention::new();
        retention.deleted_dead_header(&QueueKey::dead(0, 1));
        assert_eq!(retention.dead_count().load(Ordering::Acquire), 0);
    }
}