1use std::{
4 collections::{BTreeMap, HashMap, HashSet},
5 fmt,
6 fs::{self, File, OpenOptions},
7 io,
8 path::{Path, PathBuf},
9 sync::{Arc, RwLock},
10};
11
12use fs2::FileExt;
13use thiserror::Error;
14
15#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
17pub struct BlobId(String);
18
19impl BlobId {
20 #[must_use]
22 pub fn as_str(&self) -> &str {
23 &self.0
24 }
25
26 #[must_use]
28 pub fn from_hex(text: &str) -> Option<Self> {
29 (text.len() == 64 && text.bytes().all(|byte| byte.is_ascii_hexdigit()))
30 .then(|| Self(text.to_owned()))
31 }
32}
33
34impl fmt::Display for BlobId {
35 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36 f.write_str(&self.0)
37 }
38}
39
40#[derive(Debug, Error)]
42pub enum StoreError {
43 #[error("store I/O failed: {0}")]
45 Io(#[from] io::Error),
46 #[error("root CAS failed: expected {expected:?}, actual {actual:?}")]
48 CasFailed {
49 expected: Option<Vec<u8>>,
51 actual: Option<Vec<u8>>,
53 },
54 #[error("blob content did not match digest {0}")]
56 CorruptBlob(BlobId),
57 #[error("reachable blob is missing: {0}")]
59 MissingBlob(BlobId),
60 #[error("invalid root name {0:?}")]
62 InvalidRootName(String),
63}
64
65pub trait BlobStore: Send + Sync {
67 fn put(&self, bytes: &[u8]) -> Result<BlobId, StoreError>;
73 fn get(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError>;
79 fn contains(&self, id: &BlobId) -> Result<bool, StoreError> {
85 Ok(self.get(id)?.is_some())
86 }
87 fn delete(&self, id: &BlobId) -> Result<(), StoreError>;
93 fn list(&self) -> Result<Vec<BlobId>, StoreError>;
99}
100
101pub trait RootStore: Send + Sync {
103 fn get_root(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError>;
109 fn cas_root(&self, name: &str, expected: Option<&[u8]>, new: &[u8]) -> Result<(), StoreError>;
115 fn delete_root(&self, name: &str) -> Result<(), StoreError>;
121 fn list_roots(&self, prefix: &str) -> Result<Vec<String>, StoreError>;
127}
128
129#[derive(Clone, Default)]
131pub struct MemoryStore {
132 inner: Arc<RwLock<MemoryInner>>,
133}
134#[derive(Default)]
135struct MemoryInner {
136 blobs: HashMap<BlobId, Vec<u8>>,
137 roots: BTreeMap<String, Vec<u8>>,
138}
139
140impl BlobStore for MemoryStore {
141 fn put(&self, bytes: &[u8]) -> Result<BlobId, StoreError> {
142 let id = digest(bytes);
143 self.inner
144 .write()
145 .expect("poisoned store lock")
146 .blobs
147 .insert(id.clone(), bytes.to_vec());
148 Ok(id)
149 }
150 fn get(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError> {
151 Ok(self
152 .inner
153 .read()
154 .expect("poisoned store lock")
155 .blobs
156 .get(id)
157 .cloned())
158 }
159 fn delete(&self, id: &BlobId) -> Result<(), StoreError> {
160 self.inner
161 .write()
162 .expect("poisoned store lock")
163 .blobs
164 .remove(id);
165 Ok(())
166 }
167 fn list(&self) -> Result<Vec<BlobId>, StoreError> {
168 Ok(self
169 .inner
170 .read()
171 .expect("poisoned store lock")
172 .blobs
173 .keys()
174 .cloned()
175 .collect())
176 }
177}
178impl RootStore for MemoryStore {
179 fn get_root(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError> {
180 Ok(self
181 .inner
182 .read()
183 .expect("poisoned store lock")
184 .roots
185 .get(name)
186 .cloned())
187 }
188 fn cas_root(&self, name: &str, expected: Option<&[u8]>, new: &[u8]) -> Result<(), StoreError> {
189 let mut inner = self.inner.write().expect("poisoned store lock");
190 let actual = inner.roots.get(name).cloned();
191 if actual.as_deref() != expected {
192 return Err(StoreError::CasFailed {
193 expected: expected.map(<[u8]>::to_vec),
194 actual,
195 });
196 }
197 inner.roots.insert(name.to_owned(), new.to_vec());
198 Ok(())
199 }
200 fn delete_root(&self, name: &str) -> Result<(), StoreError> {
201 self.inner
202 .write()
203 .expect("poisoned store lock")
204 .roots
205 .remove(name);
206 Ok(())
207 }
208 fn list_roots(&self, prefix: &str) -> Result<Vec<String>, StoreError> {
209 Ok(self
210 .inner
211 .read()
212 .expect("poisoned store lock")
213 .roots
214 .keys()
215 .filter(|name| name.starts_with(prefix))
216 .cloned()
217 .collect())
218 }
219}
220
221pub struct FsStore {
223 root: PathBuf,
224}
225impl FsStore {
226 pub fn open(root: impl AsRef<Path>) -> Result<Self, StoreError> {
232 let root = root.as_ref().to_path_buf();
233 fs::create_dir_all(root.join("blobs"))?;
234 fs::create_dir_all(root.join("roots"))?;
235 Ok(Self { root })
236 }
237 fn blob_path(&self, id: &BlobId) -> PathBuf {
238 self.root.join("blobs").join(id.as_str())
239 }
240 fn root_path(&self, name: &str) -> Result<PathBuf, StoreError> {
241 if name.is_empty()
242 || name == "."
243 || name == ".."
244 || name.contains('/')
245 || name.contains('\\')
246 {
247 return Err(StoreError::InvalidRootName(name.to_owned()));
248 }
249 Ok(self.root.join("roots").join(name))
250 }
251
252 fn root_lock(&self, name: &str) -> Result<RootLock, StoreError> {
253 let root_path = self.root_path(name)?;
254 RootLock::acquire(&root_path.with_extension("lock"))
255 }
256}
257impl BlobStore for FsStore {
258 fn put(&self, bytes: &[u8]) -> Result<BlobId, StoreError> {
259 let id = digest(bytes);
260 let path = self.blob_path(&id);
261 if !path.exists() {
262 let tmp = path.with_extension("tmp");
263 fs::write(&tmp, bytes)?;
264 fs::rename(tmp, path)?;
265 }
266 Ok(id)
267 }
268 fn get(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError> {
269 let path = self.blob_path(id);
270 if !path.exists() {
271 return Ok(None);
272 }
273 let bytes = fs::read(path)?;
274 if &digest(&bytes) != id {
275 return Err(StoreError::CorruptBlob(id.clone()));
276 }
277 Ok(Some(bytes))
278 }
279 fn contains(&self, id: &BlobId) -> Result<bool, StoreError> {
280 Ok(self.blob_path(id).is_file())
281 }
282 fn delete(&self, id: &BlobId) -> Result<(), StoreError> {
283 match fs::remove_file(self.blob_path(id)) {
284 Ok(()) => Ok(()),
285 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
286 Err(error) => Err(error.into()),
287 }
288 }
289 fn list(&self) -> Result<Vec<BlobId>, StoreError> {
290 let mut ids = Vec::new();
291 for entry in fs::read_dir(self.root.join("blobs"))? {
292 let entry = entry?;
293 if entry.file_type()?.is_file() {
294 if let Some(name) = entry.file_name().to_str() {
295 if name.len() == 64 && name.bytes().all(|byte| byte.is_ascii_hexdigit()) {
296 ids.push(BlobId(name.to_owned()));
297 }
298 }
299 }
300 }
301 Ok(ids)
302 }
303}
304impl RootStore for FsStore {
305 fn get_root(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError> {
306 match fs::read(self.root_path(name)?) {
307 Ok(v) => Ok(Some(v)),
308 Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
309 Err(e) => Err(e.into()),
310 }
311 }
312 fn cas_root(&self, name: &str, expected: Option<&[u8]>, new: &[u8]) -> Result<(), StoreError> {
313 let _lock = self.root_lock(name)?;
314 let path = self.root_path(name)?;
315 let actual = match fs::read(&path) {
316 Ok(v) => Some(v),
317 Err(e) if e.kind() == io::ErrorKind::NotFound => None,
318 Err(e) => return Err(e.into()),
319 };
320 if actual.as_deref() != expected {
321 return Err(StoreError::CasFailed {
322 expected: expected.map(<[u8]>::to_vec),
323 actual,
324 });
325 }
326 let tmp = path.with_extension("tmp");
327 fs::write(&tmp, new)?;
328 fs::rename(tmp, path)?;
329 Ok(())
330 }
331 fn delete_root(&self, name: &str) -> Result<(), StoreError> {
332 let _lock = self.root_lock(name)?;
333 match fs::remove_file(self.root_path(name)?) {
334 Ok(()) => Ok(()),
335 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
336 Err(error) => Err(error.into()),
337 }
338 }
339 fn list_roots(&self, prefix: &str) -> Result<Vec<String>, StoreError> {
340 let mut names = Vec::new();
341 for entry in fs::read_dir(self.root.join("roots"))? {
342 let entry = entry?;
343 if !entry.file_type()?.is_file() {
344 continue;
345 }
346 if let Some(name) = entry.file_name().to_str() {
347 let auxiliary = Path::new(name).extension().is_some_and(|ext| {
348 ext.eq_ignore_ascii_case("lock") || ext.eq_ignore_ascii_case("tmp")
349 });
350 if name.starts_with(prefix) && !auxiliary {
351 names.push(name.to_owned());
352 }
353 }
354 }
355 names.sort();
356 Ok(names)
357 }
358}
359
360fn digest(bytes: &[u8]) -> BlobId {
361 BlobId(blake3::hash(bytes).to_hex().to_string())
362}
363
364#[must_use]
366pub fn db_root_name(db: &str) -> String {
367 format!("db:{db}")
368}
369
370#[derive(Clone, Debug, Eq, PartialEq)]
373pub struct DbRoot {
374 pub lease_version: u64,
376 pub index_basis_t: u64,
378 pub roots: Option<[BlobId; 4]>,
381}
382
383impl DbRoot {
384 #[must_use]
386 pub fn encode(&self) -> Vec<u8> {
387 let mut out = format!("{}\n{}\n", self.lease_version, self.index_basis_t);
388 match &self.roots {
389 Some(roots) => {
390 for root in roots {
391 out.push_str(root.as_str());
392 out.push('\n');
393 }
394 }
395 None => out.push_str("-\n-\n-\n-\n"),
396 }
397 out.into_bytes()
398 }
399
400 #[must_use]
402 pub fn decode(bytes: &[u8]) -> Option<Self> {
403 let text = std::str::from_utf8(bytes).ok()?;
404 let mut lines = text.lines();
405 let lease_version = lines.next()?.parse().ok()?;
406 let index_basis_t = lines.next()?.parse().ok()?;
407 let ids: Vec<&str> = lines.take(4).collect();
408 if ids.len() != 4 {
409 return None;
410 }
411 let roots = if ids.iter().all(|id| *id == "-") {
412 None
413 } else {
414 Some([
415 BlobId::from_hex(ids[0])?,
416 BlobId::from_hex(ids[1])?,
417 BlobId::from_hex(ids[2])?,
418 BlobId::from_hex(ids[3])?,
419 ])
420 };
421 Some(Self {
422 lease_version,
423 index_basis_t,
424 roots,
425 })
426 }
427}
428
429#[derive(Clone, Copy, Debug, Eq, PartialEq)]
431pub struct GcReport {
432 pub marked: usize,
434 pub swept: usize,
436}
437
438pub fn mark_and_sweep(
447 store: &dyn BlobStore,
448 live_roots: impl IntoIterator<Item = BlobId>,
449 mut children: impl FnMut(&BlobId, &[u8]) -> Result<Vec<BlobId>, StoreError>,
450) -> Result<GcReport, StoreError> {
451 let mut marked = HashSet::new();
452 let mut pending = live_roots.into_iter().collect::<Vec<_>>();
453 while let Some(id) = pending.pop() {
454 if !marked.insert(id.clone()) {
455 continue;
456 }
457 let bytes = store
458 .get(&id)?
459 .ok_or_else(|| StoreError::MissingBlob(id.clone()))?;
460 pending.extend(children(&id, &bytes)?);
461 }
462
463 let mut swept = 0;
464 for id in store.list()? {
465 if !marked.contains(&id) {
466 store.delete(&id)?;
467 swept += 1;
468 }
469 }
470 Ok(GcReport {
471 marked: marked.len(),
472 swept,
473 })
474}
475
476struct RootLock {
477 file: File,
478}
479
480impl RootLock {
481 fn acquire(path: &Path) -> Result<Self, StoreError> {
482 let file = OpenOptions::new()
483 .read(true)
484 .write(true)
485 .create(true)
486 .truncate(false)
487 .open(path)?;
488 file.lock_exclusive()?;
489 Ok(Self { file })
490 }
491}
492
493impl Drop for RootLock {
494 fn drop(&mut self) {
495 let _ = FileExt::unlock(&self.file);
496 }
500}
501
502#[derive(Default)]
504pub struct SegmentCache {
505 entries: RwLock<HashMap<BlobId, Arc<[u8]>>>,
506}
507impl SegmentCache {
508 pub fn get_or_load(
518 &self,
519 store: &dyn BlobStore,
520 id: &BlobId,
521 ) -> Result<Option<Arc<[u8]>>, StoreError> {
522 if let Some(v) = self.entries.read().expect("poisoned cache lock").get(id) {
523 return Ok(Some(v.clone()));
524 }
525 let Some(bytes) = store.get(id)? else {
526 return Ok(None);
527 };
528 let bytes: Arc<[u8]> = bytes.into();
529 self.entries
530 .write()
531 .expect("poisoned cache lock")
532 .insert(id.clone(), bytes.clone());
533 Ok(Some(bytes))
534 }
535}