use std::{collections::HashMap, sync::RwLock};
use bytes::Bytes;
use super::StorageBackend;
use crate::error::{Error, Result};
#[derive(Debug, Default)]
pub struct MemoryBackend {
data: RwLock<HashMap<String, Bytes>>,
}
impl MemoryBackend {
pub fn new() -> Self {
Self {
data: RwLock::new(HashMap::new()),
}
}
pub fn with_data(data: HashMap<String, Bytes>) -> Self {
Self {
data: RwLock::new(data),
}
}
pub fn len(&self) -> usize {
self.data.read().map(|d| d.len()).unwrap_or(0)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn clear(&self) {
if let Ok(mut data) = self.data.write() {
data.clear();
}
}
}
impl StorageBackend for MemoryBackend {
fn list(&self, prefix: &str) -> Result<Vec<String>> {
let data = self
.data
.read()
.map_err(|_| Error::storage("Failed to acquire read lock"))?;
let keys: Vec<String> = data
.keys()
.filter(|k| k.starts_with(prefix))
.cloned()
.collect();
Ok(keys)
}
fn get(&self, key: &str) -> Result<Bytes> {
let data = self
.data
.read()
.map_err(|_| Error::storage("Failed to acquire read lock"))?;
data.get(key)
.cloned()
.ok_or_else(|| Error::storage(format!("Key not found: {}", key)))
}
fn put(&self, key: &str, data: Bytes) -> Result<()> {
let mut store = self
.data
.write()
.map_err(|_| Error::storage("Failed to acquire write lock"))?;
store.insert(key.to_string(), data);
Ok(())
}
fn delete(&self, key: &str) -> Result<()> {
let mut data = self
.data
.write()
.map_err(|_| Error::storage("Failed to acquire write lock"))?;
data.remove(key);
Ok(())
}
fn exists(&self, key: &str) -> Result<bool> {
let data = self
.data
.read()
.map_err(|_| Error::storage("Failed to acquire read lock"))?;
Ok(data.contains_key(key))
}
fn size(&self, key: &str) -> Result<u64> {
let data = self
.data
.read()
.map_err(|_| Error::storage("Failed to acquire read lock"))?;
data.get(key)
.map(|d| d.len() as u64)
.ok_or_else(|| Error::storage(format!("Key not found: {}", key)))
}
}
impl Clone for MemoryBackend {
fn clone(&self) -> Self {
let data = self.data.read().map(|d| d.clone()).unwrap_or_default();
Self::with_data(data)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new() {
let backend = MemoryBackend::new();
assert!(backend.is_empty());
}
#[test]
fn test_put_and_get() {
let backend = MemoryBackend::new();
let data = Bytes::from("hello world");
backend
.put("key", data.clone())
.ok()
.unwrap_or_else(|| panic!("Should put"));
let retrieved = backend
.get("key")
.ok()
.unwrap_or_else(|| panic!("Should get"));
assert_eq!(retrieved, data);
}
#[test]
fn test_exists() {
let backend = MemoryBackend::new();
assert!(!backend
.exists("key")
.ok()
.unwrap_or_else(|| panic!("Should check")));
backend
.put("key", Bytes::from("data"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
assert!(backend
.exists("key")
.ok()
.unwrap_or_else(|| panic!("Should check")));
}
#[test]
fn test_delete() {
let backend = MemoryBackend::new();
backend
.put("key", Bytes::from("data"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
assert!(backend
.exists("key")
.ok()
.unwrap_or_else(|| panic!("Should exist")));
backend
.delete("key")
.ok()
.unwrap_or_else(|| panic!("Should delete"));
assert!(!backend
.exists("key")
.ok()
.unwrap_or_else(|| panic!("Should not exist")));
}
#[test]
fn test_list() {
let backend = MemoryBackend::new();
backend
.put("foo/bar", Bytes::from("a"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
backend
.put("foo/baz", Bytes::from("b"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
backend
.put("other", Bytes::from("c"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
let foo_keys = backend
.list("foo/")
.ok()
.unwrap_or_else(|| panic!("Should list"));
assert_eq!(foo_keys.len(), 2);
let all_keys = backend
.list("")
.ok()
.unwrap_or_else(|| panic!("Should list"));
assert_eq!(all_keys.len(), 3);
}
#[test]
fn test_size() {
let backend = MemoryBackend::new();
let data = Bytes::from("1234567890"); backend
.put("key", data)
.ok()
.unwrap_or_else(|| panic!("Should put"));
let size = backend
.size("key")
.ok()
.unwrap_or_else(|| panic!("Should get size"));
assert_eq!(size, 10);
}
#[test]
fn test_get_nonexistent() {
let backend = MemoryBackend::new();
let result = backend.get("nonexistent");
assert!(result.is_err());
}
#[test]
fn test_len() {
let backend = MemoryBackend::new();
assert_eq!(backend.len(), 0);
backend
.put("a", Bytes::from("1"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
backend
.put("b", Bytes::from("2"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
assert_eq!(backend.len(), 2);
}
#[test]
fn test_clear() {
let backend = MemoryBackend::new();
backend
.put("a", Bytes::from("1"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
backend
.put("b", Bytes::from("2"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
assert_eq!(backend.len(), 2);
backend.clear();
assert!(backend.is_empty());
}
#[test]
fn test_with_data() {
let mut initial = HashMap::new();
initial.insert("key".to_string(), Bytes::from("value"));
let backend = MemoryBackend::with_data(initial);
assert_eq!(backend.len(), 1);
assert!(backend
.exists("key")
.ok()
.unwrap_or_else(|| panic!("Should check")));
}
#[test]
fn test_clone() {
let backend = MemoryBackend::new();
backend
.put("key", Bytes::from("value"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
let cloned = backend;
assert_eq!(cloned.len(), 1);
assert_eq!(
cloned
.get("key")
.ok()
.unwrap_or_else(|| panic!("Should get")),
Bytes::from("value")
);
}
#[test]
fn test_delete_nonexistent_is_ok() {
let backend = MemoryBackend::new();
let result = backend.delete("nonexistent");
assert!(result.is_ok());
}
#[test]
fn test_size_nonexistent() {
let backend = MemoryBackend::new();
let result = backend.size("nonexistent");
assert!(result.is_err());
}
#[test]
fn test_default() {
let backend = MemoryBackend::default();
assert!(backend.is_empty());
}
#[test]
fn test_debug() {
let backend = MemoryBackend::new();
let debug_str = format!("{:?}", backend);
assert!(debug_str.contains("MemoryBackend"));
}
#[test]
fn test_list_with_prefix_filter() {
let backend = MemoryBackend::new();
backend
.put("data/train.parquet", Bytes::from("a"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
backend
.put("data/test.parquet", Bytes::from("b"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
backend
.put("metadata/info.json", Bytes::from("c"))
.ok()
.unwrap_or_else(|| panic!("Should put"));
let data_keys = backend
.list("data/")
.ok()
.unwrap_or_else(|| panic!("Should list"));
assert_eq!(data_keys.len(), 2);
let metadata_keys = backend
.list("metadata/")
.ok()
.unwrap_or_else(|| panic!("Should list"));
assert_eq!(metadata_keys.len(), 1);
}
#[test]
fn test_put_overwrite() {
let backend = MemoryBackend::new();
backend
.put("key", Bytes::from("original"))
.ok()
.unwrap_or_else(|| panic!("put 1"));
backend
.put("key", Bytes::from("updated"))
.ok()
.unwrap_or_else(|| panic!("put 2"));
let content = backend.get("key").ok().unwrap_or_else(|| panic!("get"));
assert_eq!(content, Bytes::from("updated"));
assert_eq!(backend.len(), 1);
}
#[test]
fn test_clone_independence() {
let backend = MemoryBackend::new();
backend
.put("key", Bytes::from("value"))
.ok()
.unwrap_or_else(|| panic!("put"));
let cloned = backend.clone();
backend
.put("new_key", Bytes::from("new_value"))
.ok()
.unwrap_or_else(|| panic!("put new"));
assert_eq!(backend.len(), 2);
assert_eq!(cloned.len(), 1);
}
#[test]
fn test_list_empty_prefix() {
let backend = MemoryBackend::new();
backend
.put("a", Bytes::from("1"))
.ok()
.unwrap_or_else(|| panic!("put"));
backend
.put("b", Bytes::from("2"))
.ok()
.unwrap_or_else(|| panic!("put"));
backend
.put("c", Bytes::from("3"))
.ok()
.unwrap_or_else(|| panic!("put"));
let all = backend.list("").ok().unwrap_or_else(|| panic!("list"));
assert_eq!(all.len(), 3);
}
#[test]
fn test_list_no_matches() {
let backend = MemoryBackend::new();
backend
.put("data/file.txt", Bytes::from("content"))
.ok()
.unwrap_or_else(|| panic!("put"));
let matches = backend
.list("nonexistent/")
.ok()
.unwrap_or_else(|| panic!("list"));
assert!(matches.is_empty());
}
#[test]
fn test_size_empty_value() {
let backend = MemoryBackend::new();
backend
.put("empty", Bytes::new())
.ok()
.unwrap_or_else(|| panic!("put"));
let size = backend.size("empty").ok().unwrap_or_else(|| panic!("size"));
assert_eq!(size, 0);
}
#[test]
fn test_is_empty_after_operations() {
let backend = MemoryBackend::new();
assert!(backend.is_empty());
backend
.put("key", Bytes::from("value"))
.ok()
.unwrap_or_else(|| panic!("put"));
assert!(!backend.is_empty());
backend
.delete("key")
.ok()
.unwrap_or_else(|| panic!("delete"));
assert!(backend.is_empty());
}
#[test]
fn test_many_keys() {
let backend = MemoryBackend::new();
for i in 0..100 {
backend
.put(&format!("key_{}", i), Bytes::from(format!("value_{}", i)))
.ok()
.unwrap_or_else(|| panic!("put"));
}
assert_eq!(backend.len(), 100);
for i in 0..100 {
assert!(backend
.exists(&format!("key_{}", i))
.ok()
.unwrap_or_else(|| panic!("exists")));
}
}
#[test]
fn test_delete_then_reput() {
let backend = MemoryBackend::new();
backend
.put("key", Bytes::from("v1"))
.ok()
.unwrap_or_else(|| panic!("put 1"));
backend
.delete("key")
.ok()
.unwrap_or_else(|| panic!("delete"));
backend
.put("key", Bytes::from("v2"))
.ok()
.unwrap_or_else(|| panic!("put 2"));
let content = backend.get("key").ok().unwrap_or_else(|| panic!("get"));
assert_eq!(content, Bytes::from("v2"));
}
#[test]
fn test_with_data_multiple() {
let mut initial = HashMap::new();
initial.insert("key1".to_string(), Bytes::from("value1"));
initial.insert("key2".to_string(), Bytes::from("value2"));
initial.insert("key3".to_string(), Bytes::from("value3"));
let backend = MemoryBackend::with_data(initial);
assert_eq!(backend.len(), 3);
for i in 1..=3 {
let content = backend
.get(&format!("key{}", i))
.ok()
.unwrap_or_else(|| panic!("get"));
assert_eq!(content, Bytes::from(format!("value{}", i)));
}
}
#[test]
fn test_large_value() {
let backend = MemoryBackend::new();
let data: Vec<u8> = (0..1_000_000).map(|i| (i % 256) as u8).collect();
backend
.put("large", Bytes::from(data.clone()))
.ok()
.unwrap_or_else(|| panic!("put"));
let retrieved = backend.get("large").ok().unwrap_or_else(|| panic!("get"));
assert_eq!(retrieved.len(), 1_000_000);
assert_eq!(&retrieved[..], &data[..]);
}
#[test]
fn test_clear_and_reuse() {
let backend = MemoryBackend::new();
backend
.put("old", Bytes::from("data"))
.ok()
.unwrap_or_else(|| panic!("put"));
backend.clear();
assert!(backend.is_empty());
assert!(!backend
.exists("old")
.ok()
.unwrap_or_else(|| panic!("exists")));
backend
.put("new", Bytes::from("data"))
.ok()
.unwrap_or_else(|| panic!("put"));
assert_eq!(backend.len(), 1);
}
#[test]
fn test_list_with_various_prefixes() {
let backend = MemoryBackend::new();
backend
.put("data/train/file1.parquet", Bytes::from("1"))
.ok()
.unwrap();
backend
.put("data/train/file2.parquet", Bytes::from("2"))
.ok()
.unwrap();
backend
.put("data/test/file1.parquet", Bytes::from("3"))
.ok()
.unwrap();
backend
.put("metadata/schema.json", Bytes::from("4"))
.ok()
.unwrap();
assert_eq!(backend.list("data/train/").ok().unwrap().len(), 2);
assert_eq!(backend.list("data/test/").ok().unwrap().len(), 1);
assert_eq!(backend.list("data/").ok().unwrap().len(), 3);
assert_eq!(backend.list("metadata/").ok().unwrap().len(), 1);
assert_eq!(backend.list("").ok().unwrap().len(), 4);
}
#[test]
fn test_size_of_empty_and_non_empty() {
let backend = MemoryBackend::new();
backend.put("empty", Bytes::new()).ok().unwrap();
backend.put("small", Bytes::from("abc")).ok().unwrap();
backend
.put("medium", Bytes::from("0123456789"))
.ok()
.unwrap();
assert_eq!(backend.size("empty").ok().unwrap(), 0);
assert_eq!(backend.size("small").ok().unwrap(), 3);
assert_eq!(backend.size("medium").ok().unwrap(), 10);
}
#[test]
fn test_delete_idempotent() {
let backend = MemoryBackend::new();
backend.put("key", Bytes::from("value")).ok().unwrap();
assert!(backend.delete("key").is_ok());
assert!(backend.delete("key").is_ok());
assert!(backend.delete("key").is_ok());
}
#[test]
fn test_exists_after_operations() {
let backend = MemoryBackend::new();
assert!(!backend.exists("key").ok().unwrap());
backend.put("key", Bytes::from("value")).ok().unwrap();
assert!(backend.exists("key").ok().unwrap());
backend.delete("key").ok().unwrap();
assert!(!backend.exists("key").ok().unwrap());
backend.put("key", Bytes::from("new value")).ok().unwrap();
assert!(backend.exists("key").ok().unwrap());
}
#[test]
fn test_clone_deep_copy() {
let backend = MemoryBackend::new();
backend.put("key1", Bytes::from("value1")).ok().unwrap();
backend.put("key2", Bytes::from("value2")).ok().unwrap();
let cloned = backend.clone();
backend.put("key3", Bytes::from("value3")).ok().unwrap();
backend.delete("key1").ok().unwrap();
assert_eq!(cloned.len(), 2);
assert!(cloned.exists("key1").ok().unwrap());
assert!(!cloned.exists("key3").ok().unwrap());
}
#[test]
fn test_list_partial_prefix_match() {
let backend = MemoryBackend::new();
backend.put("prefix_a_1", Bytes::from("1")).ok().unwrap();
backend.put("prefix_a_2", Bytes::from("2")).ok().unwrap();
backend.put("prefix_b_1", Bytes::from("3")).ok().unwrap();
backend.put("other", Bytes::from("4")).ok().unwrap();
assert_eq!(backend.list("prefix_a").ok().unwrap().len(), 2);
assert_eq!(backend.list("prefix_b").ok().unwrap().len(), 1);
assert_eq!(backend.list("prefix_").ok().unwrap().len(), 3);
assert_eq!(backend.list("pre").ok().unwrap().len(), 3);
}
#[test]
fn test_get_error_message() {
let backend = MemoryBackend::new();
let result = backend.get("nonexistent_key");
assert!(result.is_err());
if let Err(e) = result {
let msg = format!("{:?}", e);
assert!(msg.contains("nonexistent_key") || msg.contains("not found"));
}
}
#[test]
fn test_size_error_message() {
let backend = MemoryBackend::new();
let result = backend.size("missing_file");
assert!(result.is_err());
if let Err(e) = result {
let msg = format!("{:?}", e);
assert!(msg.contains("missing_file") || msg.contains("not found"));
}
}
#[test]
fn test_with_data_preserves_all() {
let mut initial = HashMap::new();
initial.insert("a".to_string(), Bytes::from("1"));
initial.insert("b".to_string(), Bytes::from("2"));
initial.insert("c".to_string(), Bytes::from("3"));
let backend = MemoryBackend::with_data(initial);
assert_eq!(backend.len(), 3);
assert_eq!(backend.get("a").ok().unwrap(), Bytes::from("1"));
assert_eq!(backend.get("b").ok().unwrap(), Bytes::from("2"));
assert_eq!(backend.get("c").ok().unwrap(), Bytes::from("3"));
}
#[test]
fn test_binary_data_roundtrip() {
let backend = MemoryBackend::new();
let binary: Vec<u8> = (0..=255).collect();
backend
.put("binary", Bytes::from(binary.clone()))
.ok()
.unwrap();
let retrieved = backend.get("binary").ok().unwrap();
assert_eq!(retrieved.as_ref(), binary.as_slice());
}
#[test]
fn test_concurrent_access_simulation() {
let backend = MemoryBackend::new();
for i in 0..100 {
let key = format!("key_{}", i);
let value = format!("value_{}", i);
backend.put(&key, Bytes::from(value.clone())).ok().unwrap();
}
assert_eq!(backend.len(), 100);
for i in 0..100 {
let key = format!("key_{}", i);
let expected = format!("value_{}", i);
assert_eq!(backend.get(&key).ok().unwrap(), Bytes::from(expected));
}
}
#[test]
fn test_clear_multiple_times() {
let backend = MemoryBackend::new();
backend.put("key", Bytes::from("value")).ok().unwrap();
backend.clear();
backend.clear(); backend.clear();
assert!(backend.is_empty());
}
}