arkhe_kernel/state/

scheduler.rs

//! BTreeMap-based scheduler with immediate-remove cancellation.
//!
//! Three indexed tables maintained in lockstep:
//! - `ready: BTreeMap<SchedKey, ScheduledEntry>` — primary execution queue,
//!   ordered by `(at, seq, id)`.
//! - `by_id: BTreeMap<ScheduledActionId, SchedKey>` — O(log n) cancel lookup.
//! - `by_actor: BTreeMap<EntityId, BTreeSet<ScheduledActionId>>` — O(k log n)
//!   actor-scoped cancel.
//!
//! No tombstones — `cancel` immediately removes from all three tables.
//! Determinism: BTreeMap iteration order is total over `SchedKey`, and
//! `seq` is monotonic per kernel lifetime, so identical schedule sequences
//! produce identical pop_due streams (deterministic).

use core::num::NonZeroU64;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};

use crate::abi::{EntityId, Principal, Tick, TypeCode};

/// Sentinel-free scheduled-action handle (A6).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ScheduledActionId(pub NonZeroU64);

impl ScheduledActionId {
    /// Returns `Some(_)` iff `v != 0`.
    #[inline]
    pub const fn new(v: u64) -> Option<Self> {
        match NonZeroU64::new(v) {
            Some(n) => Some(Self(n)),
            None => None,
        }
    }

    /// Underlying non-zero `u64`.
    #[inline]
    pub const fn get(self) -> u64 {
        self.0.get()
    }
}

/// Total-ordered key — `(at, seq, id)`. Tick first; same-tick FIFO by `seq`;
/// final disambiguator by `id` (defensive — `seq` alone is unique).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub(crate) struct SchedKey {
    pub at: Tick,
    pub seq: u64,
    pub id: ScheduledActionId,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct ScheduledEntry {
    pub id: ScheduledActionId,
    pub at: Tick,
    pub actor: Option<EntityId>,
    pub principal: Principal,
    pub action_type_code: TypeCode,
    /// Canonical bytes (postcard); deserialization through `ActionRegistry`
    /// happens at dispatch time.
    pub action_bytes: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct Scheduler {
    ready: BTreeMap<SchedKey, ScheduledEntry>,
    by_id: BTreeMap<ScheduledActionId, SchedKey>,
    by_actor: BTreeMap<EntityId, BTreeSet<ScheduledActionId>>,
    /// Monotonic per kernel lifetime; same-tick FIFO discriminator.
    next_seq: u64,
    /// Monotonic ID counter (NonZeroU64 starts at 1).
    next_id: u64,
}

impl Scheduler {
    pub(crate) fn new() -> Self {
        Self {
            ready: BTreeMap::new(),
            by_id: BTreeMap::new(),
            by_actor: BTreeMap::new(),
            next_seq: 0,
            next_id: 0,
        }
    }

    /// Insert into `ready` + `by_id` + `by_actor` atomically.
    pub(crate) fn schedule(
        &mut self,
        at: Tick,
        actor: Option<EntityId>,
        principal: Principal,
        type_code: TypeCode,
        bytes: Vec<u8>,
    ) -> ScheduledActionId {
        // Saturating (never wraps to 0), matching the kernel-wide A12
        // panic-free arithmetic discipline; the NonZeroU64 below therefore
        // cannot fail even after astronomically many schedules.
        self.next_id = self.next_id.saturating_add(1);
        let id = ScheduledActionId(
            NonZeroU64::new(self.next_id).expect("next_id incremented before use; never zero"),
        );

        let seq = self.next_seq;
        self.next_seq = self.next_seq.saturating_add(1);

        let key = SchedKey { at, seq, id };
        let entry = ScheduledEntry {
            id,
            at,
            actor,
            principal,
            action_type_code: type_code,
            action_bytes: bytes,
        };

        self.ready.insert(key, entry);
        self.by_id.insert(id, key);
        if let Some(actor_id) = actor {
            self.by_actor.entry(actor_id).or_default().insert(id);
        }

        id
    }

    /// Schedule with a caller-provided `id` (e.g. when Kernel pre-allocates
    /// the ScheduledActionId so it can be returned from `submit`). Internal
    /// `next_id` is bumped so future auto-allocations stay monotonic.
    pub(crate) fn schedule_with_id(
        &mut self,
        id: ScheduledActionId,
        at: Tick,
        actor: Option<EntityId>,
        principal: Principal,
        type_code: TypeCode,
        bytes: Vec<u8>,
    ) {
        if id.get() > self.next_id {
            self.next_id = id.get();
        }

        let seq = self.next_seq;
        self.next_seq = self.next_seq.saturating_add(1);

        let key = SchedKey { at, seq, id };
        let entry = ScheduledEntry {
            id,
            at,
            actor,
            principal,
            action_type_code: type_code,
            action_bytes: bytes,
        };

        self.ready.insert(key, entry);
        self.by_id.insert(id, key);
        if let Some(actor_id) = actor {
            self.by_actor.entry(actor_id).or_default().insert(id);
        }
    }

    /// Immediate cancel — three-table consistent removal.
    /// Returns the removed entry, or `None` if `id` was not scheduled
    /// (collapsed `CancelMiss` semantics — never-scheduled,
    /// already-executed, already-cancelled all return `None`).
    pub(crate) fn cancel(&mut self, id: ScheduledActionId) -> Option<ScheduledEntry> {
        let key = self.by_id.remove(&id)?;
        let entry = self
            .ready
            .remove(&key)
            .expect("ready/by_id consistency violated");
        if let Some(actor_id) = entry.actor {
            if let Some(set) = self.by_actor.get_mut(&actor_id) {
                set.remove(&id);
                if set.is_empty() {
                    self.by_actor.remove(&actor_id);
                }
            }
        }
        Some(entry)
    }

    /// Bulk-cancel every entry owned by `actor`. Returns removed entries
    /// in scheduler order (BTreeSet iteration over IDs is ascending).
    /// Production wiring (entity-despawn cascade) lands with the
    /// per-entity ownership refinement (deferred).
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn cancel_by_actor(&mut self, actor: EntityId) -> Vec<ScheduledEntry> {
        let Some(ids) = self.by_actor.remove(&actor) else {
            return Vec::new();
        };
        let mut cancelled = Vec::with_capacity(ids.len());
        for id in ids {
            if let Some(key) = self.by_id.remove(&id) {
                if let Some(entry) = self.ready.remove(&key) {
                    cancelled.push(entry);
                }
            }
        }
        cancelled
    }

    /// Pop the earliest-due entry whose `at <= now`. Returns `None` if the
    /// queue is empty or the head is in the future.
    pub(crate) fn pop_due(&mut self, now: Tick) -> Option<ScheduledEntry> {
        let (&key, _) = self.ready.first_key_value()?;
        if key.at > now {
            return None;
        }
        let entry = self
            .ready
            .remove(&key)
            .expect("first_key_value just returned this key");
        self.by_id.remove(&entry.id);
        if let Some(actor_id) = entry.actor {
            if let Some(set) = self.by_actor.get_mut(&actor_id) {
                set.remove(&entry.id);
                if set.is_empty() {
                    self.by_actor.remove(&actor_id);
                }
            }
        }
        Some(entry)
    }

    /// Validate three-table index consistency. Run after a snapshot decode
    /// so a corrupt or tampered snapshot is rejected up front rather than
    /// triggering a latent panic in `cancel` / `pop_due` (both rely on
    /// `ready`/`by_id` agreement). Returns `Err(reason)` on any mismatch.
    pub(crate) fn validate(&self) -> Result<(), &'static str> {
        // `ready` and `by_id` must be exact inverses.
        if self.ready.len() != self.by_id.len() {
            return Err("scheduler ready/by_id size mismatch");
        }
        for (key, entry) in &self.ready {
            if entry.id != key.id {
                return Err("scheduler entry id does not match its key");
            }
            match self.by_id.get(&entry.id) {
                Some(k) if *k == *key => {}
                _ => return Err("scheduler by_id does not map back to the ready key"),
            }
        }
        // `by_actor` ids must exist in `by_id` and reference the right actor.
        for (actor, ids) in &self.by_actor {
            if ids.is_empty() {
                return Err("scheduler by_actor holds an empty actor set");
            }
            for id in ids {
                let Some(key) = self.by_id.get(id) else {
                    return Err("scheduler by_actor references an unknown id");
                };
                match self.ready.get(key) {
                    Some(entry) if entry.actor == Some(*actor) => {}
                    _ => return Err("scheduler by_actor actor mismatch"),
                }
            }
        }
        // Reverse direction: every ready entry that names an actor MUST be
        // indexed under by_actor[actor], else cancel_by_actor would silently
        // miss it (the forward check above alone leaves that gap).
        for entry in self.ready.values() {
            if let Some(actor) = entry.actor {
                match self.by_actor.get(&actor) {
                    Some(set) if set.contains(&entry.id) => {}
                    _ => return Err("scheduler ready entry missing from its by_actor index"),
                }
            }
        }
        // Monotonic counters must dominate every live key.
        for key in self.ready.keys() {
            if key.seq >= self.next_seq {
                return Err("scheduler next_seq not monotonic over live keys");
            }
            if key.id.get() > self.next_id {
                return Err("scheduler next_id not monotonic over live ids");
            }
        }
        Ok(())
    }

    // Test-only observability accessors. Production introspection wiring
    // lands with the future IntrospectHandle interface (deferred).
    #[cfg_attr(not(test), allow(dead_code))]
    #[inline]
    pub(crate) fn len(&self) -> usize {
        self.ready.len()
    }

    #[cfg_attr(not(test), allow(dead_code))]
    #[inline]
    pub(crate) fn is_empty(&self) -> bool {
        self.ready.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::abi::{EntityId, Tick, TypeCode};

    fn p() -> Principal {
        Principal::System
    }
    fn tc() -> TypeCode {
        TypeCode(1)
    }

    #[test]
    fn empty_state() {
        let s = Scheduler::new();
        assert_eq!(s.len(), 0);
        assert!(s.is_empty());
    }

    #[test]
    fn schedule_then_pop_due_single() {
        let mut s = Scheduler::new();
        let id = s.schedule(Tick(5), None, p(), tc(), vec![1, 2, 3]);
        assert_eq!(s.len(), 1);
        let entry = s.pop_due(Tick(5)).expect("entry due");
        assert_eq!(entry.id, id);
        assert_eq!(entry.at, Tick(5));
        assert_eq!(entry.action_bytes, vec![1, 2, 3]);
        assert!(s.is_empty());
    }

    #[test]
    fn pop_due_before_time_returns_none() {
        let mut s = Scheduler::new();
        s.schedule(Tick(10), None, p(), tc(), vec![]);
        assert!(s.pop_due(Tick(9)).is_none());
        assert_eq!(s.len(), 1);
    }

    #[test]
    fn pop_due_at_exact_tick_pops() {
        let mut s = Scheduler::new();
        let id = s.schedule(Tick(5), None, p(), tc(), vec![]);
        assert_eq!(s.pop_due(Tick(5)).unwrap().id, id);
    }

    #[test]
    fn pop_due_ordering_by_tick() {
        let mut s = Scheduler::new();
        let id_late = s.schedule(Tick(20), None, p(), tc(), vec![]);
        let id_early = s.schedule(Tick(5), None, p(), tc(), vec![]);
        let id_mid = s.schedule(Tick(10), None, p(), tc(), vec![]);
        assert_eq!(s.pop_due(Tick(100)).unwrap().id, id_early);
        assert_eq!(s.pop_due(Tick(100)).unwrap().id, id_mid);
        assert_eq!(s.pop_due(Tick(100)).unwrap().id, id_late);
    }

    #[test]
    fn pop_due_tiebreak_by_seq() {
        let mut s = Scheduler::new();
        let id1 = s.schedule(Tick(5), None, p(), tc(), vec![1]);
        let id2 = s.schedule(Tick(5), None, p(), tc(), vec![2]);
        let id3 = s.schedule(Tick(5), None, p(), tc(), vec![3]);
        assert_eq!(s.pop_due(Tick(5)).unwrap().id, id1);
        assert_eq!(s.pop_due(Tick(5)).unwrap().id, id2);
        assert_eq!(s.pop_due(Tick(5)).unwrap().id, id3);
    }

    #[test]
    fn cancel_removes_entry() {
        let mut s = Scheduler::new();
        let id = s.schedule(Tick(5), None, p(), tc(), vec![]);
        let cancelled = s.cancel(id).expect("found");
        assert_eq!(cancelled.id, id);
        assert!(s.is_empty());
        assert!(s.pop_due(Tick(100)).is_none());
    }

    #[test]
    fn cancel_unknown_returns_none() {
        let mut s = Scheduler::new();
        let bogus = ScheduledActionId::new(999).unwrap();
        assert!(s.cancel(bogus).is_none());
    }

    #[test]
    fn cancel_by_actor_removes_all() {
        let mut s = Scheduler::new();
        let actor = EntityId::new(1).unwrap();
        let other = EntityId::new(2).unwrap();
        let _ = s.schedule(Tick(5), Some(actor), p(), tc(), vec![]);
        let _ = s.schedule(Tick(10), Some(actor), p(), tc(), vec![]);
        let id_other = s.schedule(Tick(7), Some(other), p(), tc(), vec![]);
        assert_eq!(s.len(), 3);
        let cancelled = s.cancel_by_actor(actor);
        assert_eq!(cancelled.len(), 2);
        assert_eq!(s.len(), 1);
        assert_eq!(s.pop_due(Tick(100)).unwrap().id, id_other);
    }

    #[test]
    fn cancel_by_actor_unknown_returns_empty() {
        let mut s = Scheduler::new();
        let actor = EntityId::new(99).unwrap();
        assert!(s.cancel_by_actor(actor).is_empty());
    }

    #[test]
    fn schedule_id_monotonic() {
        let mut s = Scheduler::new();
        let id1 = s.schedule(Tick(0), None, p(), tc(), vec![]);
        let id2 = s.schedule(Tick(0), None, p(), tc(), vec![]);
        let id3 = s.schedule(Tick(0), None, p(), tc(), vec![]);
        assert!(id1 < id2);
        assert!(id2 < id3);
        assert_eq!(id1.get(), 1);
        assert_eq!(id3.get(), 3);
    }

    #[test]
    fn no_tombstones() {
        // After cancel, len decrements immediately — no lazy deletion.
        let mut s = Scheduler::new();
        let id1 = s.schedule(Tick(5), None, p(), tc(), vec![]);
        let _id2 = s.schedule(Tick(5), None, p(), tc(), vec![]);
        assert_eq!(s.len(), 2);
        s.cancel(id1);
        assert_eq!(s.len(), 1);
    }

    #[test]
    fn determinism_same_sequence() {
        fn run() -> Vec<u64> {
            let mut s = Scheduler::new();
            s.schedule(Tick(3), None, p(), tc(), vec![]);
            s.schedule(Tick(1), None, p(), tc(), vec![]);
            s.schedule(Tick(2), None, p(), tc(), vec![]);
            s.schedule(Tick(1), None, p(), tc(), vec![]);
            let mut out = Vec::new();
            while let Some(e) = s.pop_due(Tick(100)) {
                out.push(e.id.get());
            }
            out
        }
        assert_eq!(run(), run());
    }

    #[test]
    fn validate_accepts_consistent_scheduler() {
        let mut s = Scheduler::new();
        s.schedule(Tick(5), Some(EntityId::new(1).unwrap()), p(), tc(), vec![1]);
        s.schedule(Tick(3), None, p(), tc(), vec![2]);
        s.schedule(Tick(5), Some(EntityId::new(1).unwrap()), p(), tc(), vec![3]);
        assert!(s.validate().is_ok());
    }

    #[test]
    fn validate_rejects_broken_ready_by_id_bijection() {
        let mut s = Scheduler::new();
        s.schedule(Tick(5), None, p(), tc(), vec![1]);
        s.by_id.clear(); // ready holds an entry with no by_id mapping
        assert!(s.validate().is_err());
    }

    #[test]
    fn validate_rejects_missing_by_actor_index() {
        let mut s = Scheduler::new();
        s.schedule(Tick(5), Some(EntityId::new(1).unwrap()), p(), tc(), vec![1]);
        s.by_actor.clear(); // actor-owned entry no longer indexed
        assert!(s.validate().is_err());
    }
}