use std::collections::HashMap;
use std::io;
use std::sync::{Arc, RwLock};
pub trait BlobStore: Send + Sync + 'static {
fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>>;
fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()>;
fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()>;
fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool>;
fn list(&self, index_name: &str) -> io::Result<Vec<String>>;
fn blob_len(&self, _index_name: &str, _file_name: &str) -> io::Result<Option<u64>> {
Ok(None)
}
fn load_range(
&self,
_index_name: &str,
_file_name: &str,
_range: std::ops::Range<u64>,
) -> io::Result<Option<Vec<u8>>> {
Ok(None)
}
}
impl<T: BlobStore + ?Sized> BlobStore for std::sync::Arc<T> {
fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>> {
(**self).load(index_name, file_name)
}
fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()> {
(**self).save(index_name, file_name, data)
}
fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()> {
(**self).delete(index_name, file_name)
}
fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool> {
(**self).exists(index_name, file_name)
}
fn list(&self, index_name: &str) -> io::Result<Vec<String>> {
(**self).list(index_name)
}
fn blob_len(&self, index_name: &str, file_name: &str) -> io::Result<Option<u64>> {
(**self).blob_len(index_name, file_name)
}
fn load_range(
&self,
index_name: &str,
file_name: &str,
range: std::ops::Range<u64>,
) -> io::Result<Option<Vec<u8>>> {
(**self).load_range(index_name, file_name, range)
}
}
type MemFiles = Arc<RwLock<HashMap<String, HashMap<String, Vec<u8>>>>>;
#[derive(Debug, Clone)]
pub struct MemBlobStore {
data: MemFiles,
}
impl MemBlobStore {
pub fn new() -> Self {
Self {
data: Arc::new(RwLock::new(HashMap::new())),
}
}
}
impl Default for MemBlobStore {
fn default() -> Self {
Self::new()
}
}
impl BlobStore for MemBlobStore {
fn blob_len(&self, index_name: &str, file_name: &str) -> io::Result<Option<u64>> {
Ok(self.data.read().unwrap()
.get(index_name)
.and_then(|files| files.get(file_name))
.map(|b| b.len() as u64))
}
fn load_range(
&self,
index_name: &str,
file_name: &str,
range: std::ops::Range<u64>,
) -> io::Result<Option<Vec<u8>>> {
Ok(self.data.read().unwrap()
.get(index_name)
.and_then(|files| files.get(file_name))
.map(|b| b[range.start as usize..(range.end as usize).min(b.len())].to_vec()))
}
fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>> {
let guard = self.data.read().map_err(|_| io::Error::other("lock poisoned"))?;
guard
.get(index_name)
.and_then(|files| files.get(file_name))
.cloned()
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::NotFound,
format!("{index_name}/{file_name} not found"),
)
})
}
fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()> {
let mut guard = self.data.write().map_err(|_| io::Error::other("lock poisoned"))?;
guard
.entry(index_name.to_string())
.or_default()
.insert(file_name.to_string(), data.to_vec());
Ok(())
}
fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()> {
let mut guard = self.data.write().map_err(|_| io::Error::other("lock poisoned"))?;
if let Some(files) = guard.get_mut(index_name) {
files.remove(file_name);
}
Ok(())
}
fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool> {
let guard = self.data.read().map_err(|_| io::Error::other("lock poisoned"))?;
Ok(guard
.get(index_name)
.is_some_and(|files| files.contains_key(file_name)))
}
fn list(&self, index_name: &str) -> io::Result<Vec<String>> {
let guard = self.data.read().map_err(|_| io::Error::other("lock poisoned"))?;
Ok(guard
.get(index_name)
.map(|files| files.keys().cloned().collect())
.unwrap_or_default())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_mem_blob_store_roundtrip() {
let store = MemBlobStore::new();
store.save("idx1", "file.bin", b"hello world").unwrap();
assert!(store.exists("idx1", "file.bin").unwrap());
assert!(!store.exists("idx1", "other.bin").unwrap());
let loaded = store.load("idx1", "file.bin").unwrap();
assert_eq!(loaded, b"hello world");
store.delete("idx1", "file.bin").unwrap();
assert!(!store.exists("idx1", "file.bin").unwrap());
}
#[test]
fn test_mem_blob_store_multiple_indexes() {
let store = MemBlobStore::new();
store.save("idx1", "a.bin", b"aaa").unwrap();
store.save("idx2", "a.bin", b"xxx").unwrap();
assert_eq!(store.load("idx1", "a.bin").unwrap(), b"aaa");
assert_eq!(store.load("idx2", "a.bin").unwrap(), b"xxx");
}
}