klieo-bus-memory 3.3.0

//! In-process `KvStore` implementation.
//!
//! State: a `BTreeMap<(bucket, key), Entry>` guarded by a single
//! `tokio::sync::Mutex`. Per-entry monotonically-increasing revision.
//! Leases share the same map; heartbeating extends the lease's
//! `expires_at`. The map is key-ordered so `keys_paginated` ranges from a
//! cursor in `O(log n + page)` without rescanning the bucket.

use async_trait::async_trait;
use bytes::Bytes;
use klieo_core::bus::{KeyPage, KvEntry, KvStore, Lease, LeaseImpl, Revision};
use klieo_core::error::BusError;
use std::collections::BTreeMap;
use std::ops::Bound;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;

#[derive(Debug, Clone)]
struct Entry {
    value: Bytes,
    revision: Revision,
    /// `Some(deadline)` — entry is leased until `deadline`. `None` — no
    /// active lease.
    leased_until: Option<Instant>,
}

#[derive(Default)]
struct State {
    /// `(bucket, key) -> Entry`.
    map: BTreeMap<(String, String), Entry>,
}

/// In-process `KvStore` impl.
pub struct MemoryKv {
    state: Arc<Mutex<State>>,
}

impl MemoryKv {
    /// Build an empty store.
    pub fn new() -> Self {
        Self {
            state: Arc::new(Mutex::new(State::default())),
        }
    }
}

impl Default for MemoryKv {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl KvStore for MemoryKv {
    async fn get(&self, bucket: &str, key: &str) -> Result<Option<KvEntry>, BusError> {
        let g = self.state.lock().await;
        Ok(g.map.get(&(bucket.into(), key.into())).map(|e| KvEntry {
            value: e.value.clone(),
            revision: e.revision,
        }))
    }

    async fn put(&self, bucket: &str, key: &str, value: Bytes) -> Result<Revision, BusError> {
        let mut g = self.state.lock().await;
        let key = (bucket.into(), key.into());
        let next_rev = g.map.get(&key).map(|e| e.revision + 1).unwrap_or(1);
        g.map.insert(
            key,
            Entry {
                value,
                revision: next_rev,
                leased_until: None,
            },
        );
        Ok(next_rev)
    }

    async fn cas(
        &self,
        bucket: &str,
        key: &str,
        value: Bytes,
        expected: Option<Revision>,
    ) -> Result<Revision, BusError> {
        let mut g = self.state.lock().await;
        let map_key = (bucket.into(), key.into());
        let existing = g.map.get(&map_key);
        match (expected, existing) {
            (None, None) => {
                g.map.insert(
                    map_key,
                    Entry {
                        value,
                        revision: 1,
                        leased_until: None,
                    },
                );
                Ok(1)
            }
            (None, Some(e)) => Err(BusError::CasConflict {
                expected: 0,
                actual: e.revision,
            }),
            (Some(want), Some(e)) if e.revision == want => {
                let next = e.revision + 1;
                g.map.insert(
                    map_key,
                    Entry {
                        value,
                        revision: next,
                        leased_until: None,
                    },
                );
                Ok(next)
            }
            (Some(want), Some(e)) => Err(BusError::CasConflict {
                expected: want,
                actual: e.revision,
            }),
            (Some(want), None) => Err(BusError::CasConflict {
                expected: want,
                actual: 0,
            }),
        }
    }

    async fn delete(&self, bucket: &str, key: &str) -> Result<(), BusError> {
        let mut g = self.state.lock().await;
        g.map.remove(&(bucket.into(), key.into()));
        Ok(())
    }

    async fn keys(&self, bucket: &str) -> Result<Vec<String>, BusError> {
        let g = self.state.lock().await;
        Ok(g.map
            .keys()
            .filter(|(b, _)| b == bucket)
            .map(|(_, k)| k.clone())
            .collect())
    }

    async fn scan_bucket(&self, bucket: &str) -> Result<Vec<(String, Bytes)>, BusError> {
        let g = self.state.lock().await;
        let entries = g
            .map
            .iter()
            .filter(|((b, k), _)| b == bucket && !klieo_core::bus::is_causer_key(k))
            .map(|((_, k), e)| (k.clone(), e.value.clone()))
            .collect();
        Ok(entries)
    }

    /// Ranges the key-ordered map from the cursor, reading only `limit`(+1) keys
    /// — no full-bucket rescan, unlike the trait default. `limit` is floored to
    /// 1 so a zero never stalls a page-walk.
    async fn keys_paginated(
        &self,
        bucket: &str,
        cursor: Option<String>,
        limit: usize,
    ) -> Result<KeyPage, BusError> {
        let limit = limit.max(1);
        let start = match &cursor {
            Some(c) => Bound::Excluded((bucket.to_string(), c.clone())),
            None => Bound::Included((bucket.to_string(), String::new())),
        };
        let g = self.state.lock().await;
        let mut after_cursor = g
            .map
            .range::<(String, String), _>((start, Bound::Unbounded))
            .map(|((b, k), _)| (b.as_str(), k))
            .take_while(|(b, _)| *b == bucket)
            .map(|(_, k)| k.clone());
        let keys: Vec<String> = after_cursor.by_ref().take(limit).collect();
        let next = if after_cursor.next().is_some() {
            keys.last().cloned()
        } else {
            None
        };
        Ok(KeyPage::new(keys, next))
    }

    async fn lease(&self, bucket: &str, key: &str, ttl: Duration) -> Result<Lease, BusError> {
        let now = Instant::now();
        let mut g = self.state.lock().await;
        let map_key = (bucket.to_string(), key.to_string());
        let entry = g.map.entry(map_key.clone()).or_insert_with(|| Entry {
            value: Bytes::new(),
            revision: 0,
            leased_until: None,
        });
        if let Some(deadline) = entry.leased_until {
            if deadline > now {
                return Err(BusError::Permanent("lease already held".into()));
            }
        }
        entry.leased_until = Some(now + ttl);
        Ok(Lease::new(Box::new(MemoryLease {
            state: self.state.clone(),
            bucket: bucket.to_string(),
            key: key.to_string(),
            ttl,
        })))
    }
}

/// Lease handle. Heartbeat extends `leased_until` by `ttl` from now.
struct MemoryLease {
    state: Arc<Mutex<State>>,
    bucket: String,
    key: String,
    ttl: Duration,
}

#[async_trait]
impl LeaseImpl for MemoryLease {
    async fn heartbeat(&self) -> Result<(), BusError> {
        let mut g = self.state.lock().await;
        let key = (self.bucket.clone(), self.key.clone());
        match g.map.get_mut(&key) {
            Some(e) => {
                e.leased_until = Some(Instant::now() + self.ttl);
                Ok(())
            }
            None => Err(BusError::NotFound(format!("{}/{}", self.bucket, self.key))),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bytes::Bytes;
    use klieo_core::error::BusError;
    use klieo_core::KvStore;
    use std::time::Duration;

    #[tokio::test]
    async fn put_then_get_returns_value() {
        let kv = MemoryKv::new();
        let rev = kv.put("b", "k", Bytes::from_static(b"v1")).await.unwrap();
        assert!(rev > 0);
        let entry = kv.get("b", "k").await.unwrap().unwrap();
        assert_eq!(entry.value, Bytes::from_static(b"v1"));
        assert_eq!(entry.revision, rev);
    }

    #[tokio::test]
    async fn kv_causer_round_trips_and_leaves_value_key_intact() {
        use klieo_core::ids::RunId;
        let kv = MemoryKv::new();
        let run = RunId::new();
        kv.put("b", "k", Bytes::from_static(b"val")).await.unwrap();
        kv.put_causer("b", "k", run).await.unwrap();
        assert_eq!(kv.causer_of("b", "k").await.unwrap(), Some(run));
        // the value key is unaffected by the sidecar causer key
        assert_eq!(
            kv.get("b", "k").await.unwrap().unwrap().value,
            Bytes::from_static(b"val")
        );
        // an unstamped key has no causer
        kv.put("b", "k2", Bytes::from_static(b"v2")).await.unwrap();
        assert_eq!(kv.causer_of("b", "k2").await.unwrap(), None);

        // first-writer-wins: a later put_causer must NOT overwrite the causer.
        let later = RunId::new();
        kv.put_causer("b", "k", later).await.unwrap();
        assert_eq!(
            kv.causer_of("b", "k").await.unwrap(),
            Some(run),
            "the first writer's causer survives a later stamp"
        );

        // scan_bucket hides the reserved causer sidecar keys.
        let scanned = kv.scan_bucket("b").await.unwrap();
        assert!(
            scanned
                .iter()
                .all(|(k, _)| !klieo_core::bus::is_causer_key(k)),
            "scan_bucket must not surface causer sidecar keys; got {scanned:?}"
        );
        assert!(
            scanned.iter().any(|(k, _)| k == "k"),
            "the value key is still present in scan_bucket"
        );
    }

    #[tokio::test]
    async fn cas_create_if_absent_succeeds_when_key_missing() {
        let kv = MemoryKv::new();
        let rev = kv
            .cas("b", "k", Bytes::from_static(b"v"), None)
            .await
            .unwrap();
        assert_eq!(rev, 1);
    }

    #[tokio::test]
    async fn cas_create_if_absent_conflicts_when_key_present() {
        let kv = MemoryKv::new();
        kv.put("b", "k", Bytes::from_static(b"v0")).await.unwrap();
        let err = kv
            .cas("b", "k", Bytes::from_static(b"v1"), None)
            .await
            .unwrap_err();
        assert!(matches!(err, BusError::CasConflict { .. }));
    }

    #[tokio::test]
    async fn cas_with_expected_revision_succeeds_on_match() {
        let kv = MemoryKv::new();
        let r0 = kv.put("b", "k", Bytes::from_static(b"v0")).await.unwrap();
        let r1 = kv
            .cas("b", "k", Bytes::from_static(b"v1"), Some(r0))
            .await
            .unwrap();
        assert_eq!(r1, r0 + 1);
    }

    #[tokio::test]
    async fn cas_with_expected_revision_conflicts_on_stale() {
        let kv = MemoryKv::new();
        let r0 = kv.put("b", "k", Bytes::from_static(b"v0")).await.unwrap();
        kv.put("b", "k", Bytes::from_static(b"v1")).await.unwrap();
        let err = kv
            .cas("b", "k", Bytes::from_static(b"v2"), Some(r0))
            .await
            .unwrap_err();
        assert!(matches!(err, BusError::CasConflict { .. }));
    }

    #[tokio::test]
    async fn delete_removes_key() {
        let kv = MemoryKv::new();
        kv.put("b", "k", Bytes::from_static(b"v")).await.unwrap();
        kv.delete("b", "k").await.unwrap();
        assert!(kv.get("b", "k").await.unwrap().is_none());
    }

    #[tokio::test]
    async fn lease_can_heartbeat() {
        let kv = MemoryKv::new();
        let lease = kv.lease("b", "k", Duration::from_millis(50)).await.unwrap();
        lease.heartbeat().await.unwrap();
    }

    #[tokio::test]
    async fn keys_returns_keys_in_bucket() {
        let kv = MemoryKv::new();
        kv.put("b1", "alpha", Bytes::from_static(b"a"))
            .await
            .unwrap();
        kv.put("b1", "beta", Bytes::from_static(b"b"))
            .await
            .unwrap();
        kv.put("b2", "gamma", Bytes::from_static(b"c"))
            .await
            .unwrap();
        let mut keys = kv.keys("b1").await.unwrap();
        keys.sort();
        assert_eq!(keys, vec!["alpha".to_string(), "beta".to_string()]);
        let keys_other = kv.keys("b2").await.unwrap();
        assert_eq!(keys_other, vec!["gamma".to_string()]);
        let keys_empty = kv.keys("missing").await.unwrap();
        assert!(keys_empty.is_empty());
    }

    #[tokio::test]
    async fn keys_paginated_walks_bucket_in_order_across_pages() {
        let kv = MemoryKv::new();
        for k in ["e", "a", "d", "b", "c"] {
            kv.put("b1", k, Bytes::from_static(b"v")).await.unwrap();
        }
        kv.put("b2", "zzz", Bytes::from_static(b"v")).await.unwrap();

        let mut seen = Vec::new();
        let mut cursor = None;
        loop {
            let page = kv.keys_paginated("b1", cursor, 2).await.unwrap();
            assert!(page.keys.len() <= 2, "page respects the limit");
            seen.extend(page.keys.iter().cloned());
            match page.next {
                Some(c) => cursor = Some(c),
                None => break,
            }
        }
        assert_eq!(
            seen,
            vec!["a", "b", "c", "d", "e"],
            "every key once, ascending; the b2 key never leaks in"
        );
    }

    #[tokio::test]
    async fn keys_paginated_empty_bucket_yields_one_terminal_page() {
        let kv = MemoryKv::new();
        let page = kv.keys_paginated("missing", None, 8).await.unwrap();
        assert!(page.keys.is_empty());
        assert!(page.next.is_none());
    }

    #[tokio::test]
    async fn keys_paginated_limit_past_count_is_single_page() {
        let kv = MemoryKv::new();
        kv.put("b1", "k1", Bytes::from_static(b"v")).await.unwrap();
        kv.put("b1", "k2", Bytes::from_static(b"v")).await.unwrap();
        let page = kv.keys_paginated("b1", None, 16).await.unwrap();
        assert_eq!(page.keys, vec!["k1".to_string(), "k2".to_string()]);
        assert!(page.next.is_none());
    }

    #[tokio::test]
    async fn scan_bucket_returns_all_pairs_in_single_read() {
        let kv = MemoryKv::new();
        kv.put("b1", "k1", Bytes::from_static(b"v1")).await.unwrap();
        kv.put("b1", "k2", Bytes::from_static(b"v2")).await.unwrap();
        kv.put("b2", "k3", Bytes::from_static(b"v3")).await.unwrap();

        let mut pairs = kv.scan_bucket("b1").await.unwrap();
        pairs.sort_by(|a, b| a.0.cmp(&b.0));

        assert_eq!(pairs.len(), 2);
        assert_eq!(pairs[0], ("k1".to_string(), Bytes::from_static(b"v1")));
        assert_eq!(pairs[1], ("k2".to_string(), Bytes::from_static(b"v2")));

        let pairs_other = kv.scan_bucket("b2").await.unwrap();
        assert_eq!(pairs_other.len(), 1);
        assert_eq!(pairs_other[0].0, "k3");

        let pairs_empty = kv.scan_bucket("missing").await.unwrap();
        assert!(pairs_empty.is_empty());
    }
}