klieo-bus-memory 0.41.2

//! In-process `KvStore` implementation.
//!
//! State: per-bucket `HashMap<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`.

use async_trait::async_trait;
use bytes::Bytes;
use klieo_core::bus::{KvEntry, KvStore, Lease, LeaseImpl, Revision};
use klieo_core::error::BusError;
use std::collections::HashMap;
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: HashMap<(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, _), _)| b == bucket)
            .map(|((_, k), e)| (k.clone(), e.value.clone()))
            .collect();
        Ok(entries)
    }

    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 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 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());
    }
}