alkstore/alkstore-mem/src/state.rs

//! The guarded per-mechanism state core — the store instance *is* the
//! state (engine-mem.md §Architecture): per-mechanism tables owned by
//! the engine instance behind the family's locking discipline
//! (short-held synchronous internal locks; tokio primitives for
//! cross-await coordination — [`crate::wakes`]' channels and the
//! scheduler's boundary waits arrive with their tasks).
//!
//! The idiom: each mechanism's committed state sits behind one
//! short-held [`std::sync::Mutex`] section, accessed through
//! [`Guard::guard`]. Guards are never held across an `.await`, never
//! across user-provided code (delivery closures, payload callbacks) —
//! staging happens engine-side, so the mutex cannot observe a panic
//! while held and poisoning is structurally unreachable; the
//! [`Guard::guard`] fold into the inner value keeps the type honest
//! anyway.
//!
//! Section shapes are the commit-merge targets ([`crate::tx`]'s
//! overlay stages effects of exactly these rows). Row-level operations
//! live with their mechanisms (`crate::queue`, `crate::stream`,
//! `crate::lock`, `crate::scheduler`) — modules must take section
//! guards in the canonical order **schedules → queues → streams**
//! (locks and wakes depend on no other section; the commit-merge takes
//! queues → streams, the scheduler tick schedules → queues) so no path
//! can deadlock.

use std::collections::{BTreeMap, HashMap};
use std::sync::Mutex;
use std::sync::MutexGuard;
use std::sync::atomic::{AtomicI64, Ordering};

/// A short-held synchronous internal lock: `guard()` is the one access
/// idiom (see the module docs for the discipline).
pub(crate) trait Guard<T> {
    fn guard(&self) -> MutexGuard<'_, T>;
}

impl<T> Guard<T> for Mutex<T> {
    fn guard(&self) -> MutexGuard<'_, T> {
        self.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
    }
}

/// A live job row (ADR-010's state machine, minus the dead state —
/// dead rows move to [`QueueState::dead`] with their diagnosis). The
/// row carries the resolved enqueue fields and the §3a stamps; the
/// `pending` flag is the `pending`/`processing` split (absence = acked
/// or cancelled).
#[derive(Debug, Clone)]
pub(crate) struct JobRow {
    pub(crate) id: i64,
    pub(crate) queue: String,
    pub(crate) payload: Vec<u8>,
    pub(crate) pending: bool,
    pub(crate) priority: i64,
    pub(crate) run_at: i64,
    pub(crate) worker_id: Option<String>,
    pub(crate) claim_expires_at: Option<i64>,
    pub(crate) claimed_at: Option<i64>,
    pub(crate) attempts: i64,
    pub(crate) max_attempts: i64,
    pub(crate) created_at: i64,
    pub(crate) expires_at: Option<i64>,
    pub(crate) visibility_timeout_s: i64,
    pub(crate) backoff_base_s: i64,
    pub(crate) dead_letter_retention_s: Option<i64>,
}

/// A dead-lettered job row (ADR-010 §4): the live row's durable fields
/// plus the diagnosis (`last_error`, `died_at`) — `get_job` sees dead
/// rows with the stamps still attached. The claim identity
/// (`worker_id`/`claim_expires_at`) never rides a dead row.
#[derive(Debug, Clone)]
pub(crate) struct DeadRow {
    pub(crate) id: i64,
    pub(crate) queue: String,
    pub(crate) payload: Vec<u8>,
    pub(crate) priority: i64,
    pub(crate) run_at: i64,
    pub(crate) attempts: i64,
    pub(crate) max_attempts: i64,
    pub(crate) created_at: i64,
    pub(crate) expires_at: Option<i64>,
    pub(crate) claimed_at: Option<i64>,
    pub(crate) visibility_timeout_s: i64,
    pub(crate) backoff_base_s: i64,
    pub(crate) dead_letter_retention_s: Option<i64>,
    pub(crate) last_error: String,
    pub(crate) died_at: i64,
}

/// A committed stream event (ADR-015 §3's shape, engine-internal) —
/// `offset` is the per-stream log key, immutable, never renumbered.
#[derive(Debug, Clone)]
pub(crate) struct EventRow {
    pub(crate) offset: i64,
    pub(crate) key: Option<String>,
    pub(crate) payload: Vec<u8>,
    pub(crate) created_at: i64,
}

/// A named TTL lock's registry entry (ADR-008 §7): owner identity +
/// deadline. Expiry is a clock check at the try/renew boundaries — no
/// ambient revoker exists.
#[derive(Debug, Clone)]
pub(crate) struct LockEntry {
    pub(crate) owner: String,
    pub(crate) expires_at: i64,
}

/// A registered schedule row (ADR-009): the collapsed scheduler's
/// table. The interval is parsed and the stamps are resolved at
/// register time (the fire path stamps rows from the registered
/// values, never the raw spec) — `next_fire_at` is the collapse
/// machinery's boundary state.
#[derive(Debug, Clone)]
pub(crate) struct ScheduleRow {
    pub(crate) queue: String,
    pub(crate) interval_s: i64,
    pub(crate) payload: Vec<u8>,
    pub(crate) priority: i64,
    pub(crate) expires_s: Option<i64>,
    pub(crate) max_attempts: i64,
    pub(crate) visibility_timeout_s: i64,
    pub(crate) backoff_base_s: i64,
    pub(crate) dead_letter_retention_s: Option<i64>,
    pub(crate) next_fire_at: i64,
}

/// The queue mechanism's committed section: the jobs table and the
/// dead-letter storage. `live` keys by job id (the enqueue-order
/// claim tiebreak reads the id directly); `dead` keys the same id
/// space.
#[derive(Debug, Default)]
pub(crate) struct QueueState {
    pub(crate) live: BTreeMap<i64, JobRow>,
    pub(crate) dead: BTreeMap<i64, DeadRow>,
}

/// The stream mechanism's committed section: per-stream event logs
/// (keyed by offset, ordered) and the per-(stream, consumer) offset
/// checkpoints.
#[derive(Debug, Default)]
pub(crate) struct StreamState {
    pub(crate) logs: BTreeMap<String, BTreeMap<i64, EventRow>>,
    pub(crate) offsets: HashMap<(String, String), i64>,
}

/// The named-locks mechanism's committed section: the TTL registry.
#[derive(Debug, Default)]
pub(crate) struct LockState {
    pub(crate) entries: HashMap<String, LockEntry>,
}

/// The scheduler mechanism's committed section: the schedule table.
#[derive(Debug, Default)]
pub(crate) struct ScheduleState {
    pub(crate) rows: BTreeMap<String, ScheduleRow>,
}

/// The store's committed state: one guarded section per mechanism, the
/// engine-wide id/offset allocators, and the clock read that every
/// timestamp in every section defers to.
///
/// `job_ids`/`event_offsets` are the allocators the tx overlay and the
/// auto-commit paths share: ids and offsets are assigned monotone at
/// allocation time (stage time for a tx, op time for auto-commit),
/// which is the pg-shaped story — an uncommitted allocation gaps out
/// when the tx rolls back (legal: ids are claim-order tiebreaks and
/// offsets are immutable, never renumbered; gaps are ordinary).
pub(crate) struct State {
    pub(crate) queues: Mutex<QueueState>,
    pub(crate) streams: Mutex<StreamState>,
    pub(crate) locks: Mutex<LockState>,
    pub(crate) schedules: Mutex<ScheduleState>,
    pub(crate) job_ids: AtomicI64,
    pub(crate) event_offsets: AtomicI64,
}

impl State {
    pub(crate) fn new() -> Self {
        Self {
            queues: Mutex::new(QueueState::default()),
            streams: Mutex::new(StreamState::default()),
            locks: Mutex::new(LockState::default()),
            schedules: Mutex::new(ScheduleState::default()),
            job_ids: AtomicI64::new(0),
            event_offsets: AtomicI64::new(0),
        }
    }

    /// The next job id — monotone, never reused (a rolled-back tx's
    /// staged allocation gaps).
    pub(crate) fn allocate_job_id(&self) -> i64 {
        self.job_ids.fetch_add(1, Ordering::Relaxed) + 1
    }

    /// The next event offset — monotone, immutable, never renumbered.
    pub(crate) fn allocate_event_offset(&self) -> i64 {
        self.event_offsets.fetch_add(1, Ordering::Relaxed) + 1
    }
}

impl std::fmt::Debug for State {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("State").finish_non_exhaustive()
    }
}

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

    #[test]
    fn sections_are_per_mechanism_and_independent() {
        let state = State::new();
        let queued_ptr = {
            let mut queues = state.queues.guard();
            std::ptr::from_mut(&mut *queues)
        };
        let streams_ptr = {
            let mut streams = state.streams.guard();
            std::ptr::from_mut(&mut *streams)
        };
        let queues_ptr_again = {
            let mut queues = state.queues.guard();
            std::ptr::from_mut(&mut *queues)
        };
        assert_eq!(queued_ptr, queues_ptr_again, "the same section relocks");
        assert_ne!(
            queued_ptr.cast::<()>(),
            streams_ptr.cast::<()>(),
            "sections are distinct storage"
        );
    }

    #[test]
    fn sections_relock_after_guard_drop() {
        let state = std::sync::Arc::new(State::new());
        std::thread::scope(|scope| {
            let held = scope.spawn(|| {
                for _ in 0..100 {
                    let _guard = state.streams.guard();
                    std::thread::yield_now();
                }
            });
            for _ in 0..100 {
                let _guard = state.queues.guard();
                std::thread::yield_now();
            }
            held.join().expect("the lock churn task never panics");
        });
    }

    #[test]
    fn allocators_are_monotone_and_distinct() {
        let state = State::new();
        let ids: Vec<i64> = (0..8).map(|_| state.allocate_job_id()).collect();
        let offsets: Vec<i64> = (0..8).map(|_| state.allocate_event_offset()).collect();
        assert_eq!(ids, (1..=8).collect::<Vec<_>>(), "job ids count from 1");
        assert_eq!(offsets, (1..=8).collect::<Vec<_>>(), "offsets count from 1");
    }
}