1use std::collections::HashMap;
13use std::io;
14use std::sync::{Arc, RwLock};
15
16pub trait BlobStore: Send + Sync + 'static {
22 fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>>;
24
25 fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()>;
27
28 fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()>;
30
31 fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool>;
33
34 fn list(&self, index_name: &str) -> io::Result<Vec<String>>;
36
37 fn blob_len(&self, _index_name: &str, _file_name: &str) -> io::Result<Option<u64>> {
42 Ok(None)
43 }
44
45 fn load_range(
51 &self,
52 _index_name: &str,
53 _file_name: &str,
54 _range: std::ops::Range<u64>,
55 ) -> io::Result<Option<Vec<u8>>> {
56 Ok(None)
57 }
58}
59
60impl<T: BlobStore + ?Sized> BlobStore for std::sync::Arc<T> {
64 fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>> {
65 (**self).load(index_name, file_name)
66 }
67 fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()> {
68 (**self).save(index_name, file_name, data)
69 }
70 fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()> {
71 (**self).delete(index_name, file_name)
72 }
73 fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool> {
74 (**self).exists(index_name, file_name)
75 }
76 fn list(&self, index_name: &str) -> io::Result<Vec<String>> {
77 (**self).list(index_name)
78 }
79 fn blob_len(&self, index_name: &str, file_name: &str) -> io::Result<Option<u64>> {
80 (**self).blob_len(index_name, file_name)
81 }
82 fn load_range(
83 &self,
84 index_name: &str,
85 file_name: &str,
86 range: std::ops::Range<u64>,
87 ) -> io::Result<Option<Vec<u8>>> {
88 (**self).load_range(index_name, file_name, range)
89 }
90}
91
92#[derive(Debug, Clone)]
94pub struct MemBlobStore {
95 data: Arc<RwLock<HashMap<String, HashMap<String, Vec<u8>>>>>,
97}
98
99impl MemBlobStore {
100 pub fn new() -> Self {
101 Self {
102 data: Arc::new(RwLock::new(HashMap::new())),
103 }
104 }
105}
106
107impl Default for MemBlobStore {
108 fn default() -> Self {
109 Self::new()
110 }
111}
112
113impl BlobStore for MemBlobStore {
114 fn blob_len(&self, index_name: &str, file_name: &str) -> io::Result<Option<u64>> {
115 Ok(self.data.read().unwrap()
116 .get(index_name)
117 .and_then(|files| files.get(file_name))
118 .map(|b| b.len() as u64))
119 }
120
121 fn load_range(
122 &self,
123 index_name: &str,
124 file_name: &str,
125 range: std::ops::Range<u64>,
126 ) -> io::Result<Option<Vec<u8>>> {
127 Ok(self.data.read().unwrap()
128 .get(index_name)
129 .and_then(|files| files.get(file_name))
130 .map(|b| b[range.start as usize..(range.end as usize).min(b.len())].to_vec()))
131 }
132
133 fn load(&self, index_name: &str, file_name: &str) -> io::Result<Vec<u8>> {
134 let guard = self.data.read().map_err(|_| {
135 io::Error::new(io::ErrorKind::Other, "lock poisoned")
136 })?;
137 guard
138 .get(index_name)
139 .and_then(|files| files.get(file_name))
140 .cloned()
141 .ok_or_else(|| {
142 io::Error::new(
143 io::ErrorKind::NotFound,
144 format!("{index_name}/{file_name} not found"),
145 )
146 })
147 }
148
149 fn save(&self, index_name: &str, file_name: &str, data: &[u8]) -> io::Result<()> {
150 let mut guard = self.data.write().map_err(|_| {
151 io::Error::new(io::ErrorKind::Other, "lock poisoned")
152 })?;
153 guard
154 .entry(index_name.to_string())
155 .or_default()
156 .insert(file_name.to_string(), data.to_vec());
157 Ok(())
158 }
159
160 fn delete(&self, index_name: &str, file_name: &str) -> io::Result<()> {
161 let mut guard = self.data.write().map_err(|_| {
162 io::Error::new(io::ErrorKind::Other, "lock poisoned")
163 })?;
164 if let Some(files) = guard.get_mut(index_name) {
165 files.remove(file_name);
166 }
167 Ok(())
168 }
169
170 fn exists(&self, index_name: &str, file_name: &str) -> io::Result<bool> {
171 let guard = self.data.read().map_err(|_| {
172 io::Error::new(io::ErrorKind::Other, "lock poisoned")
173 })?;
174 Ok(guard
175 .get(index_name)
176 .map_or(false, |files| files.contains_key(file_name)))
177 }
178
179 fn list(&self, index_name: &str) -> io::Result<Vec<String>> {
180 let guard = self.data.read().map_err(|_| {
181 io::Error::new(io::ErrorKind::Other, "lock poisoned")
182 })?;
183 Ok(guard
184 .get(index_name)
185 .map(|files| files.keys().cloned().collect())
186 .unwrap_or_default())
187 }
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193
194 #[test]
195 fn test_mem_blob_store_roundtrip() {
196 let store = MemBlobStore::new();
197 store.save("idx1", "file.bin", b"hello world").unwrap();
198 assert!(store.exists("idx1", "file.bin").unwrap());
199 assert!(!store.exists("idx1", "other.bin").unwrap());
200 let loaded = store.load("idx1", "file.bin").unwrap();
201 assert_eq!(loaded, b"hello world");
202 store.delete("idx1", "file.bin").unwrap();
203 assert!(!store.exists("idx1", "file.bin").unwrap());
204 }
205
206 #[test]
207 fn test_mem_blob_store_multiple_indexes() {
208 let store = MemBlobStore::new();
209 store.save("idx1", "a.bin", b"aaa").unwrap();
210 store.save("idx2", "a.bin", b"xxx").unwrap();
211 assert_eq!(store.load("idx1", "a.bin").unwrap(), b"aaa");
212 assert_eq!(store.load("idx2", "a.bin").unwrap(), b"xxx");
213 }
214}