use std::collections::BTreeMap;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use lru::LruCache;
pub use crate::error::KvError;
#[derive(Debug, Clone)]
pub enum WriteOp {
Put(String, Vec<u8>),
Delete(String),
}
#[async_trait]
pub trait KvStore: Send + Sync {
async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, KvError>;
async fn put(&self, key: &str, value: Vec<u8>) -> Result<(), KvError>;
async fn delete(&self, key: &str) -> Result<(), KvError>;
async fn list_prefix(&self, prefix: &str) -> Result<Vec<String>, KvError>;
async fn flush(&self) -> Result<(), KvError> {
Ok(())
}
async fn write_batch(&self, ops: Vec<WriteOp>) -> Result<(), KvError> {
for op in ops {
match op {
WriteOp::Put(key, value) => self.put(&key, value).await?,
WriteOp::Delete(key) => self.delete(&key).await?,
}
}
Ok(())
}
fn invalidate_cache(&self) {}
fn invalidate_keys(&self, keys: &[String]) {
let _ = keys;
}
}
#[derive(Debug, Default, Clone)]
pub struct MemoryKv {
inner: Arc<Mutex<BTreeMap<String, Vec<u8>>>>,
}
impl MemoryKv {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl KvStore for MemoryKv {
async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, KvError> {
Ok(self.inner.lock().unwrap().get(key).cloned())
}
async fn put(&self, key: &str, value: Vec<u8>) -> Result<(), KvError> {
self.inner.lock().unwrap().insert(key.to_string(), value);
Ok(())
}
async fn delete(&self, key: &str) -> Result<(), KvError> {
self.inner.lock().unwrap().remove(key);
Ok(())
}
async fn list_prefix(&self, prefix: &str) -> Result<Vec<String>, KvError> {
Ok(self
.inner
.lock()
.unwrap()
.keys()
.filter(|key| key.starts_with(prefix))
.cloned()
.collect())
}
async fn write_batch(&self, ops: Vec<WriteOp>) -> Result<(), KvError> {
let mut map = self.inner.lock().unwrap();
for op in ops {
match op {
WriteOp::Put(key, value) => {
map.insert(key, value);
}
WriteOp::Delete(key) => {
map.remove(&key);
}
}
}
Ok(())
}
}
#[async_trait]
pub trait ChangePublisher: Send + Sync {
async fn publish(&self, keys: &[String]);
}
pub struct CachedKv {
inner: Arc<dyn KvStore>,
cache: Mutex<LruCache<String, Vec<u8>>>,
publisher: Option<Arc<dyn ChangePublisher>>,
}
impl CachedKv {
pub fn new(inner: Arc<dyn KvStore>, capacity: usize) -> Self {
let capacity = NonZeroUsize::new(capacity.max(1)).expect("capacity >= 1");
Self {
inner,
cache: Mutex::new(LruCache::new(capacity)),
publisher: None,
}
}
pub fn with_publisher(mut self, publisher: Arc<dyn ChangePublisher>) -> Self {
self.publisher = Some(publisher);
self
}
async fn announce(&self, keys: Vec<String>) {
if let Some(publisher) = &self.publisher {
publisher.publish(&keys).await;
}
}
}
#[async_trait]
impl KvStore for CachedKv {
async fn flush(&self) -> Result<(), KvError> {
self.inner.flush().await
}
async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, KvError> {
{
let mut cache = self.cache.lock().unwrap();
if let Some(value) = cache.get(key) {
return Ok(Some(value.clone()));
}
}
let value = self.inner.get(key).await?;
if let Some(bytes) = &value {
self.cache
.lock()
.unwrap()
.put(key.to_string(), bytes.clone());
}
Ok(value)
}
async fn put(&self, key: &str, value: Vec<u8>) -> Result<(), KvError> {
self.inner.put(key, value.clone()).await?;
self.cache.lock().unwrap().put(key.to_string(), value);
self.announce(vec![key.to_string()]).await;
Ok(())
}
async fn delete(&self, key: &str) -> Result<(), KvError> {
self.inner.delete(key).await?;
self.cache.lock().unwrap().pop(key);
self.announce(vec![key.to_string()]).await;
Ok(())
}
async fn list_prefix(&self, prefix: &str) -> Result<Vec<String>, KvError> {
self.inner.list_prefix(prefix).await
}
async fn write_batch(&self, ops: Vec<WriteOp>) -> Result<(), KvError> {
self.inner.write_batch(ops.clone()).await?;
let mut changed = Vec::with_capacity(ops.len());
{
let mut cache = self.cache.lock().unwrap();
for op in ops {
match op {
WriteOp::Put(key, value) => {
changed.push(key.clone());
cache.put(key, value);
}
WriteOp::Delete(key) => {
cache.pop(&key);
changed.push(key);
}
}
}
}
self.announce(changed).await;
Ok(())
}
fn invalidate_cache(&self) {
self.cache.lock().unwrap().clear();
}
fn invalidate_keys(&self, keys: &[String]) {
let mut cache = self.cache.lock().unwrap();
for key in keys {
cache.pop(key);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn cached_kv_round_trips_and_caches() {
let backing = Arc::new(MemoryKv::new());
let kv = CachedKv::new(backing.clone(), 8);
assert_eq!(kv.get("a").await.unwrap(), None);
kv.put("a", b"1".to_vec()).await.unwrap();
assert_eq!(kv.get("a").await.unwrap(), Some(b"1".to_vec()));
backing.put("a", b"2".to_vec()).await.unwrap();
assert_eq!(kv.get("a").await.unwrap(), Some(b"1".to_vec()));
kv.delete("a").await.unwrap();
assert_eq!(kv.get("a").await.unwrap(), None);
}
#[tokio::test]
async fn invalidate_cache_drops_stale_entries() {
let backing = Arc::new(MemoryKv::new());
backing.put("k", b"v1".to_vec()).await.unwrap();
let kv = CachedKv::new(backing.clone(), 8);
assert_eq!(kv.get("k").await.unwrap(), Some(b"v1".to_vec()));
backing.put("k", b"v2".to_vec()).await.unwrap();
assert_eq!(
kv.get("k").await.unwrap(),
Some(b"v1".to_vec()),
"still cached"
);
kv.invalidate_cache();
assert_eq!(kv.get("k").await.unwrap(), Some(b"v2".to_vec()));
}
#[tokio::test]
async fn write_batch_applies_puts_and_deletes() {
let backing = Arc::new(MemoryKv::new());
backing.put("old", b"gone".to_vec()).await.unwrap();
let kv = CachedKv::new(backing.clone(), 8);
assert_eq!(kv.get("old").await.unwrap(), Some(b"gone".to_vec()));
kv.write_batch(vec![
WriteOp::Put("a".into(), b"1".to_vec()),
WriteOp::Put("b".into(), b"2".to_vec()),
WriteOp::Delete("old".into()),
])
.await
.unwrap();
assert_eq!(kv.get("a").await.unwrap(), Some(b"1".to_vec()));
assert_eq!(kv.get("b").await.unwrap(), Some(b"2".to_vec()));
assert_eq!(kv.get("old").await.unwrap(), None);
assert_eq!(backing.get("a").await.unwrap(), Some(b"1".to_vec()));
assert_eq!(backing.get("old").await.unwrap(), None);
}
}