use super::backend::{BatchOperations, StorageBackend};
use crate::error::{Error, Result};
use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};
pub struct MemoryBackend {
data: Arc<RwLock<BTreeMap<Vec<u8>, Vec<u8>>>>,
}
impl MemoryBackend {
#[allow(clippy::must_use_candidate)]
pub fn new() -> Self {
Self {
data: Arc::new(RwLock::new(BTreeMap::new())),
}
}
}
impl Default for MemoryBackend {
fn default() -> Self {
Self::new()
}
}
impl StorageBackend for MemoryBackend {
fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
let data = self
.data
.read()
.map_err(|e| Error::Internal(e.to_string()))?;
Ok(data.get(key).cloned())
}
fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> {
let mut data = self
.data
.write()
.map_err(|e| Error::Internal(e.to_string()))?;
data.insert(key.to_vec(), value.to_vec());
Ok(())
}
fn remove(&self, key: &[u8]) -> Result<()> {
let mut data = self
.data
.write()
.map_err(|e| Error::Internal(e.to_string()))?;
data.remove(key);
Ok(())
}
fn prefix_scan(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
let data = self
.data
.read()
.map_err(|e| Error::Internal(e.to_string()))?;
let results: Vec<_> = data
.range(prefix.to_vec()..)
.take_while(|(k, _)| k.starts_with(prefix))
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
Ok(results)
}
fn prefix_count(&self, prefix: &[u8]) -> Result<usize> {
let data = self
.data
.read()
.map_err(|e| Error::Internal(e.to_string()))?;
Ok(data
.range(prefix.to_vec()..)
.take_while(|(k, _)| k.starts_with(prefix))
.count())
}
fn prefix_scan_keys(&self, prefix: &[u8]) -> Result<Vec<Vec<u8>>> {
let data = self
.data
.read()
.map_err(|e| Error::Internal(e.to_string()))?;
Ok(data
.range(prefix.to_vec()..)
.take_while(|(k, _)| k.starts_with(prefix))
.map(|(k, _)| k.clone())
.collect())
}
fn prefix_scan_batch(
&self,
prefix: &[u8],
batch_size: usize,
after_key: Option<&[u8]>,
) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
let data = self
.data
.read()
.map_err(|e| Error::Internal(e.to_string()))?;
let start: Vec<u8> = if let Some(after) = after_key {
let mut s = after.to_vec();
s.push(0);
s
} else {
prefix.to_vec()
};
Ok(data
.range(start..)
.take_while(|(k, _)| k.starts_with(prefix))
.take(batch_size)
.map(|(k, v)| (k.clone(), v.clone()))
.collect())
}
fn range_scan(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
let data = self
.data
.read()
.map_err(|e| Error::Internal(e.to_string()))?;
let results: Vec<_> = data
.range(start.to_vec()..end.to_vec())
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
Ok(results)
}
fn batch(&self) -> Box<dyn BatchOperations> {
Box::new(MemoryBatch {
data: Arc::clone(&self.data),
operations: Vec::new(),
preconditions: Vec::new(),
})
}
fn flush(&self) -> Result<()> {
Ok(())
}
}
enum BatchOp {
Insert(Vec<u8>, Vec<u8>),
Remove(Vec<u8>),
}
struct Precondition {
key: Vec<u8>,
expected_value: Vec<u8>,
}
pub struct MemoryBatch {
data: Arc<RwLock<BTreeMap<Vec<u8>, Vec<u8>>>>,
operations: Vec<BatchOp>,
preconditions: Vec<Precondition>,
}
impl BatchOperations for MemoryBatch {
fn insert(&mut self, key: Vec<u8>, value: Vec<u8>) {
self.operations.push(BatchOp::Insert(key, value));
}
fn remove(&mut self, key: Vec<u8>) {
self.operations.push(BatchOp::Remove(key));
}
fn expect_value(&mut self, key: Vec<u8>, expected_value: Vec<u8>) {
self.preconditions.push(Precondition {
key,
expected_value,
});
}
fn commit(self: Box<Self>) -> Result<()> {
let mut data = self
.data
.write()
.map_err(|e| Error::Internal(e.to_string()))?;
for precondition in &self.preconditions {
let actual = data.get(&precondition.key);
match actual {
Some(val) if val.as_slice() == precondition.expected_value.as_slice() => {}
_ => {
return Err(Error::Conflict(
"optimistic lock failed: value was modified".into(),
));
}
}
}
for op in self.operations {
match op {
BatchOp::Insert(k, v) => {
data.insert(k, v);
}
BatchOp::Remove(k) => {
data.remove(&k);
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_basic_operations() {
let backend = MemoryBackend::new();
StorageBackend::insert(&backend, b"key1", b"value1").unwrap();
assert_eq!(
StorageBackend::get(&backend, b"key1").unwrap(),
Some(b"value1".to_vec())
);
StorageBackend::remove(&backend, b"key1").unwrap();
assert_eq!(StorageBackend::get(&backend, b"key1").unwrap(), None);
}
#[test]
fn test_prefix_scan() {
let backend = MemoryBackend::new();
StorageBackend::insert(&backend, b"users/1", b"alice").unwrap();
StorageBackend::insert(&backend, b"users/2", b"bob").unwrap();
StorageBackend::insert(&backend, b"posts/1", b"hello").unwrap();
let results = StorageBackend::prefix_scan(&backend, b"users/").unwrap();
assert_eq!(results.len(), 2);
}
#[test]
fn test_batch_commit() {
let backend = MemoryBackend::new();
let mut batch = StorageBackend::batch(&backend);
batch.insert(b"key1".to_vec(), b"value1".to_vec());
batch.insert(b"key2".to_vec(), b"value2".to_vec());
batch.commit().unwrap();
assert_eq!(
StorageBackend::get(&backend, b"key1").unwrap(),
Some(b"value1".to_vec())
);
assert_eq!(
StorageBackend::get(&backend, b"key2").unwrap(),
Some(b"value2".to_vec())
);
}
#[test]
fn test_optimistic_lock_success() {
let backend = MemoryBackend::new();
StorageBackend::insert(&backend, b"key1", b"value1").unwrap();
let mut batch = StorageBackend::batch(&backend);
batch.expect_value(b"key1".to_vec(), b"value1".to_vec());
batch.insert(b"key1".to_vec(), b"value2".to_vec());
batch.commit().unwrap();
assert_eq!(
StorageBackend::get(&backend, b"key1").unwrap(),
Some(b"value2".to_vec())
);
}
#[test]
fn test_optimistic_lock_failure() {
let backend = MemoryBackend::new();
StorageBackend::insert(&backend, b"key1", b"value1").unwrap();
let mut batch = StorageBackend::batch(&backend);
batch.expect_value(b"key1".to_vec(), b"wrong_value".to_vec());
batch.insert(b"key1".to_vec(), b"value2".to_vec());
let result = batch.commit();
assert!(result.is_err());
assert_eq!(
StorageBackend::get(&backend, b"key1").unwrap(),
Some(b"value1".to_vec())
);
}
}