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 time::{Duration, SystemTime},
11};
12
13use fs2::FileExt;
14use thiserror::Error;
15
16#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
18pub struct BlobId(String);
19
20impl BlobId {
21 #[must_use]
23 pub fn as_str(&self) -> &str {
24 &self.0
25 }
26
27 #[must_use]
29 pub fn from_hex(text: &str) -> Option<Self> {
30 (text.len() == 64 && text.bytes().all(|byte| byte.is_ascii_hexdigit()))
31 .then(|| Self(text.to_owned()))
32 }
33}
34
35impl fmt::Display for BlobId {
36 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37 f.write_str(&self.0)
38 }
39}
40
41#[derive(Debug, Error)]
43pub enum StoreError {
44 #[error("store I/O failed: {0}")]
46 Io(#[from] io::Error),
47 #[error("root CAS failed: expected {expected:?}, actual {actual:?}")]
49 CasFailed {
50 expected: Option<Vec<u8>>,
52 actual: Option<Vec<u8>>,
54 },
55 #[error("blob content did not match digest {0}")]
57 CorruptBlob(BlobId),
58 #[error("reachable blob is missing: {0}")]
60 MissingBlob(BlobId),
61 #[error("invalid root name {0:?}")]
63 InvalidRootName(String),
64}
65
66pub trait BlobStore: Send + Sync {
68 fn put(&self, bytes: &[u8]) -> Result<BlobId, StoreError>;
74 fn get(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError>;
80 fn contains(&self, id: &BlobId) -> Result<bool, StoreError> {
86 Ok(self.get(id)?.is_some())
87 }
88 fn delete(&self, id: &BlobId) -> Result<(), StoreError>;
94 fn list(&self) -> Result<Vec<BlobId>, StoreError>;
100 fn modified_at(&self, _id: &BlobId) -> Result<Option<SystemTime>, StoreError> {
107 Ok(None)
108 }
109}
110
111pub trait RootStore: Send + Sync {
113 fn get_root(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError>;
119 fn cas_root(&self, name: &str, expected: Option<&[u8]>, new: &[u8]) -> Result<(), StoreError>;
125 fn delete_root(&self, name: &str) -> Result<(), StoreError>;
131 fn list_roots(&self, prefix: &str) -> Result<Vec<String>, StoreError>;
137}
138
139#[derive(Clone, Default)]
141pub struct MemoryStore {
142 inner: Arc<RwLock<MemoryInner>>,
143}
144#[derive(Default)]
145struct MemoryInner {
146 blobs: HashMap<BlobId, Vec<u8>>,
147 roots: BTreeMap<String, Vec<u8>>,
148}
149
150impl BlobStore for MemoryStore {
151 fn put(&self, bytes: &[u8]) -> Result<BlobId, StoreError> {
152 let id = digest(bytes);
153 self.inner
154 .write()
155 .expect("poisoned store lock")
156 .blobs
157 .insert(id.clone(), bytes.to_vec());
158 Ok(id)
159 }
160 fn get(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError> {
161 Ok(self
162 .inner
163 .read()
164 .expect("poisoned store lock")
165 .blobs
166 .get(id)
167 .cloned())
168 }
169 fn delete(&self, id: &BlobId) -> Result<(), StoreError> {
170 self.inner
171 .write()
172 .expect("poisoned store lock")
173 .blobs
174 .remove(id);
175 Ok(())
176 }
177 fn list(&self) -> Result<Vec<BlobId>, StoreError> {
178 Ok(self
179 .inner
180 .read()
181 .expect("poisoned store lock")
182 .blobs
183 .keys()
184 .cloned()
185 .collect())
186 }
187}
188impl RootStore for MemoryStore {
189 fn get_root(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError> {
190 Ok(self
191 .inner
192 .read()
193 .expect("poisoned store lock")
194 .roots
195 .get(name)
196 .cloned())
197 }
198 fn cas_root(&self, name: &str, expected: Option<&[u8]>, new: &[u8]) -> Result<(), StoreError> {
199 let mut inner = self.inner.write().expect("poisoned store lock");
200 let actual = inner.roots.get(name).cloned();
201 if actual.as_deref() != expected {
202 return Err(StoreError::CasFailed {
203 expected: expected.map(<[u8]>::to_vec),
204 actual,
205 });
206 }
207 inner.roots.insert(name.to_owned(), new.to_vec());
208 Ok(())
209 }
210 fn delete_root(&self, name: &str) -> Result<(), StoreError> {
211 self.inner
212 .write()
213 .expect("poisoned store lock")
214 .roots
215 .remove(name);
216 Ok(())
217 }
218 fn list_roots(&self, prefix: &str) -> Result<Vec<String>, StoreError> {
219 Ok(self
220 .inner
221 .read()
222 .expect("poisoned store lock")
223 .roots
224 .keys()
225 .filter(|name| name.starts_with(prefix))
226 .cloned()
227 .collect())
228 }
229}
230
231pub struct FsStore {
233 root: PathBuf,
234}
235impl FsStore {
236 pub fn open(root: impl AsRef<Path>) -> Result<Self, StoreError> {
242 let root = root.as_ref().to_path_buf();
243 fs::create_dir_all(root.join("blobs"))?;
244 fs::create_dir_all(root.join("roots"))?;
245 Ok(Self { root })
246 }
247 fn blob_path(&self, id: &BlobId) -> PathBuf {
248 self.root.join("blobs").join(id.as_str())
249 }
250 fn root_path(&self, name: &str) -> Result<PathBuf, StoreError> {
251 if name.is_empty()
252 || name == "."
253 || name == ".."
254 || name.contains('/')
255 || name.contains('\\')
256 {
257 return Err(StoreError::InvalidRootName(name.to_owned()));
258 }
259 Ok(self.root.join("roots").join(name))
260 }
261
262 fn root_lock(&self, name: &str) -> Result<RootLock, StoreError> {
263 let root_path = self.root_path(name)?;
264 RootLock::acquire(&root_path.with_extension("lock"))
265 }
266}
267impl BlobStore for FsStore {
268 fn put(&self, bytes: &[u8]) -> Result<BlobId, StoreError> {
269 let id = digest(bytes);
270 let path = self.blob_path(&id);
271 if !path.exists() {
272 let tmp = path.with_extension("tmp");
273 fs::write(&tmp, bytes)?;
274 fs::rename(tmp, path)?;
275 }
276 Ok(id)
277 }
278 fn get(&self, id: &BlobId) -> Result<Option<Vec<u8>>, StoreError> {
279 let path = self.blob_path(id);
280 if !path.exists() {
281 return Ok(None);
282 }
283 let bytes = fs::read(path)?;
284 if &digest(&bytes) != id {
285 return Err(StoreError::CorruptBlob(id.clone()));
286 }
287 Ok(Some(bytes))
288 }
289 fn contains(&self, id: &BlobId) -> Result<bool, StoreError> {
290 Ok(self.blob_path(id).is_file())
291 }
292 fn delete(&self, id: &BlobId) -> Result<(), StoreError> {
293 match fs::remove_file(self.blob_path(id)) {
294 Ok(()) => Ok(()),
295 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
296 Err(error) => Err(error.into()),
297 }
298 }
299 fn list(&self) -> Result<Vec<BlobId>, StoreError> {
300 let mut ids = Vec::new();
301 for entry in fs::read_dir(self.root.join("blobs"))? {
302 let entry = entry?;
303 if entry.file_type()?.is_file() {
304 if let Some(name) = entry.file_name().to_str() {
305 if name.len() == 64 && name.bytes().all(|byte| byte.is_ascii_hexdigit()) {
306 ids.push(BlobId(name.to_owned()));
307 }
308 }
309 }
310 }
311 Ok(ids)
312 }
313 fn modified_at(&self, id: &BlobId) -> Result<Option<SystemTime>, StoreError> {
314 match fs::metadata(self.blob_path(id)) {
315 Ok(metadata) => Ok(Some(metadata.modified()?)),
316 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None),
317 Err(error) => Err(error.into()),
318 }
319 }
320}
321impl RootStore for FsStore {
322 fn get_root(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError> {
323 match fs::read(self.root_path(name)?) {
324 Ok(v) => Ok(Some(v)),
325 Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(None),
326 Err(e) => Err(e.into()),
327 }
328 }
329 fn cas_root(&self, name: &str, expected: Option<&[u8]>, new: &[u8]) -> Result<(), StoreError> {
330 let _lock = self.root_lock(name)?;
331 let path = self.root_path(name)?;
332 let actual = match fs::read(&path) {
333 Ok(v) => Some(v),
334 Err(e) if e.kind() == io::ErrorKind::NotFound => None,
335 Err(e) => return Err(e.into()),
336 };
337 if actual.as_deref() != expected {
338 return Err(StoreError::CasFailed {
339 expected: expected.map(<[u8]>::to_vec),
340 actual,
341 });
342 }
343 let tmp = path.with_extension("tmp");
344 fs::write(&tmp, new)?;
345 fs::rename(tmp, path)?;
346 Ok(())
347 }
348 fn delete_root(&self, name: &str) -> Result<(), StoreError> {
349 let _lock = self.root_lock(name)?;
350 match fs::remove_file(self.root_path(name)?) {
351 Ok(()) => Ok(()),
352 Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
353 Err(error) => Err(error.into()),
354 }
355 }
356 fn list_roots(&self, prefix: &str) -> Result<Vec<String>, StoreError> {
357 let mut names = Vec::new();
358 for entry in fs::read_dir(self.root.join("roots"))? {
359 let entry = entry?;
360 if !entry.file_type()?.is_file() {
361 continue;
362 }
363 if let Some(name) = entry.file_name().to_str() {
364 let auxiliary = Path::new(name).extension().is_some_and(|ext| {
365 ext.eq_ignore_ascii_case("lock") || ext.eq_ignore_ascii_case("tmp")
366 });
367 if name.starts_with(prefix) && !auxiliary {
368 names.push(name.to_owned());
369 }
370 }
371 }
372 names.sort();
373 Ok(names)
374 }
375}
376
377fn digest(bytes: &[u8]) -> BlobId {
378 BlobId(blake3::hash(bytes).to_hex().to_string())
379}
380
381#[must_use]
383pub fn db_root_name(db: &str) -> String {
384 format!("db:{db}")
385}
386
387pub const FORMAT_VERSION: u32 = 1;
389
390#[derive(Clone, Debug, Eq, PartialEq)]
393pub struct DbRoot {
394 pub format_version: u32,
396 pub lease_version: u64,
398 pub index_basis_t: u64,
400 pub roots: Option<[BlobId; 4]>,
403}
404
405impl DbRoot {
406 #[must_use]
408 pub fn encode(&self) -> Vec<u8> {
409 let mut out = format!(
410 "corium-root-v{}\n{}\n{}\n",
411 self.format_version, self.lease_version, self.index_basis_t
412 );
413 match &self.roots {
414 Some(roots) => {
415 for root in roots {
416 out.push_str(root.as_str());
417 out.push('\n');
418 }
419 }
420 None => out.push_str("-\n-\n-\n-\n"),
421 }
422 out.into_bytes()
423 }
424
425 #[must_use]
427 pub fn decode(bytes: &[u8]) -> Option<Self> {
428 let text = std::str::from_utf8(bytes).ok()?;
429 let mut lines = text.lines();
430 let first = lines.next()?;
431 let (format_version, lease_version) =
432 if let Some(version) = first.strip_prefix("corium-root-v") {
433 let format_version = version.parse().ok()?;
434 let lease_version = lines.next()?.parse().ok()?;
435 (format_version, lease_version)
436 } else {
437 (FORMAT_VERSION, first.parse().ok()?)
440 };
441 let index_basis_t = lines.next()?.parse().ok()?;
442 let ids: Vec<&str> = lines.take(4).collect();
443 if ids.len() != 4 {
444 return None;
445 }
446 let roots = if ids.iter().all(|id| *id == "-") {
447 None
448 } else {
449 Some([
450 BlobId::from_hex(ids[0])?,
451 BlobId::from_hex(ids[1])?,
452 BlobId::from_hex(ids[2])?,
453 BlobId::from_hex(ids[3])?,
454 ])
455 };
456 Some(Self {
457 format_version,
458 lease_version,
459 index_basis_t,
460 roots,
461 })
462 }
463}
464
465#[derive(Clone, Copy, Debug, Eq, PartialEq)]
467pub struct GcReport {
468 pub marked: usize,
470 pub swept: usize,
472 pub retained: usize,
474}
475
476pub fn mark_and_sweep(
485 store: &dyn BlobStore,
486 live_roots: impl IntoIterator<Item = BlobId>,
487 mut children: impl FnMut(&BlobId, &[u8]) -> Result<Vec<BlobId>, StoreError>,
488) -> Result<GcReport, StoreError> {
489 mark_and_sweep_retained(
490 store,
491 live_roots,
492 &mut children,
493 Duration::ZERO,
494 SystemTime::now(),
495 )
496}
497
498pub fn mark_and_sweep_retained(
504 store: &dyn BlobStore,
505 live_roots: impl IntoIterator<Item = BlobId>,
506 mut children: impl FnMut(&BlobId, &[u8]) -> Result<Vec<BlobId>, StoreError>,
507 retention: Duration,
508 now: SystemTime,
509) -> Result<GcReport, StoreError> {
510 let mut marked = HashSet::new();
511 let mut pending = live_roots.into_iter().collect::<Vec<_>>();
512 while let Some(id) = pending.pop() {
513 if !marked.insert(id.clone()) {
514 continue;
515 }
516 let bytes = store
517 .get(&id)?
518 .ok_or_else(|| StoreError::MissingBlob(id.clone()))?;
519 pending.extend(children(&id, &bytes)?);
520 }
521
522 let mut swept = 0;
523 let mut retained = 0;
524 for id in store.list()? {
525 if !marked.contains(&id) {
526 let old_enough = retention.is_zero()
530 || store.modified_at(&id)?.is_some_and(|modified| {
531 now.duration_since(modified).unwrap_or_default() >= retention
532 });
533 if old_enough {
534 store.delete(&id)?;
535 swept += 1;
536 } else {
537 retained += 1;
538 }
539 }
540 }
541 Ok(GcReport {
542 marked: marked.len(),
543 swept,
544 retained,
545 })
546}
547
548struct RootLock {
549 file: File,
550}
551
552impl RootLock {
553 fn acquire(path: &Path) -> Result<Self, StoreError> {
554 let file = OpenOptions::new()
555 .read(true)
556 .write(true)
557 .create(true)
558 .truncate(false)
559 .open(path)?;
560 file.lock_exclusive()?;
561 Ok(Self { file })
562 }
563}
564
565impl Drop for RootLock {
566 fn drop(&mut self) {
567 let _ = FileExt::unlock(&self.file);
568 }
572}
573
574#[derive(Default)]
576pub struct SegmentCache {
577 entries: RwLock<HashMap<BlobId, Arc<[u8]>>>,
578}
579impl SegmentCache {
580 pub fn get_or_load(
590 &self,
591 store: &dyn BlobStore,
592 id: &BlobId,
593 ) -> Result<Option<Arc<[u8]>>, StoreError> {
594 if let Some(v) = self.entries.read().expect("poisoned cache lock").get(id) {
595 return Ok(Some(v.clone()));
596 }
597 let Some(bytes) = store.get(id)? else {
598 return Ok(None);
599 };
600 let bytes: Arc<[u8]> = bytes.into();
601 self.entries
602 .write()
603 .expect("poisoned cache lock")
604 .insert(id.clone(), bytes.clone());
605 Ok(Some(bytes))
606 }
607}