klieo-bus-memory 0.2.0

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