1#[cfg(test)]
5use std::sync::atomic::AtomicUsize;
6use std::{
7 collections::{BTreeSet, HashMap, VecDeque},
8 hash::Hash,
9 path::{Path, PathBuf},
10 sync::{
11 Arc, Mutex, RwLock,
12 atomic::{AtomicBool, Ordering},
13 },
14};
15
16use heddle_format::compression::CompressionConfig;
17
18use super::{
19 fs_io::{AtomicWriteMode, write_atomic},
20 fs_paths::{actions_dir, blobs_dir, packs_dir, states_dir, tree_lineage_dir, trees_dir},
21 npk1::Npk1Manager,
22};
23use crate::{
24 fs_atomic::sync_directory,
25 object::{Blob, ContentHash, State, StateId, Tree},
26 store::{Result, SnapshotPackManager, pack::PackObjectId},
27};
28
29const RECENT_BLOB_CACHE_CAPACITY: usize = 2_048;
30const RECENT_TREE_CACHE_CAPACITY: usize = 1_024;
31const VERIFIED_LOOSE_BLOB_CACHE_CAPACITY: usize = 65_536;
38pub(super) const RECENT_BLOB_CACHE_MAX_BYTES: usize = 4 * 1024 * 1024;
42pub(super) const RECENT_BLOB_CACHE_MAX_TOTAL_BYTES: usize = 256 * 1024 * 1024;
49
50thread_local! {
51 static SNAPSHOT_WRITE_BATCH_DEPTHS: std::cell::RefCell<HashMap<PathBuf, usize>> =
52 std::cell::RefCell::new(HashMap::new());
53}
54
55#[derive(Clone, Copy, Debug, Eq, PartialEq)]
56pub enum LooseObjectWriteMode {
57 Durable,
58 BatchDirectorySync,
59}
60
61#[derive(Debug)]
83pub(super) struct RecentObjectCache<K, V> {
84 entries: HashMap<K, RecentObjectCacheEntry<V>>,
85 eviction_clock: VecDeque<K>,
86 capacity: usize,
87 byte_budget: Option<usize>,
89 sizer: fn(&V) -> usize,
92 cached_bytes: usize,
94}
95
96#[derive(Debug)]
97struct RecentObjectCacheEntry<V> {
98 value: V,
99 recently_accessed: AtomicBool,
100}
101
102impl<K, V> RecentObjectCache<K, V>
103where
104 K: Copy + Eq + Hash,
105{
106 pub(super) fn with_capacity(capacity: usize) -> Self {
110 Self {
111 entries: HashMap::new(),
112 eviction_clock: VecDeque::new(),
113 capacity,
114 byte_budget: None,
115 sizer: |_| 0,
116 cached_bytes: 0,
117 }
118 }
119
120 pub(super) fn with_byte_budget(
124 capacity: usize,
125 byte_budget: usize,
126 sizer: fn(&V) -> usize,
127 ) -> Self {
128 Self {
129 entries: HashMap::new(),
130 eviction_clock: VecDeque::new(),
131 capacity,
132 byte_budget: Some(byte_budget),
133 sizer,
134 cached_bytes: 0,
135 }
136 }
137
138 pub(super) fn get(&self, key: &K) -> Option<&V> {
141 let entry = self.entries.get(key)?;
142 entry.recently_accessed.store(true, Ordering::Relaxed);
143 Some(&entry.value)
144 }
145
146 pub(super) fn contains(&self, key: &K) -> bool {
151 self.entries.contains_key(key)
152 }
153
154 #[cfg(test)]
163 pub(super) fn remove(&mut self, key: &K) -> Option<V> {
164 let removed = self.entries.remove(key)?.value;
165 self.cached_bytes = self.cached_bytes.saturating_sub((self.sizer)(&removed));
166 Some(removed)
167 }
168
169 pub(super) fn insert(&mut self, key: K, value: V) {
170 if self.capacity == 0 {
171 return;
172 }
173 let new_bytes = self.byte_budget.map(|_| (self.sizer)(&value)).unwrap_or(0);
174 let entry = RecentObjectCacheEntry {
175 value,
176 recently_accessed: AtomicBool::new(false),
177 };
178 if let Some(old) = self.entries.insert(key, entry) {
179 self.cached_bytes = self.cached_bytes.saturating_sub(
180 self.byte_budget
181 .map(|_| (self.sizer)(&old.value))
182 .unwrap_or(0),
183 );
184 } else {
185 self.eviction_clock.push_back(key);
186 }
187 self.cached_bytes += new_bytes;
188 self.evict_to_fit(key);
189 }
190
191 fn evict_to_fit(&mut self, admitted_key: K) {
199 loop {
200 let over_count = self.entries.len() > self.capacity;
201 let over_bytes = self
202 .byte_budget
203 .is_some_and(|budget| self.cached_bytes > budget && self.entries.len() > 1);
204 if !over_count && !over_bytes {
205 break;
206 }
207 let Some(candidate) = self.eviction_clock.pop_front() else {
208 break;
209 };
210 let Some(entry) = self.entries.get(&candidate) else {
211 continue;
212 };
213 if candidate == admitted_key && self.entries.len() > 1 {
218 self.eviction_clock.push_back(candidate);
219 continue;
220 }
221 if entry.recently_accessed.swap(false, Ordering::Relaxed) {
222 self.eviction_clock.push_back(candidate);
223 continue;
224 }
225 if let Some(evicted) = self.entries.remove(&candidate) {
226 self.cached_bytes = self.cached_bytes.saturating_sub(
227 self.byte_budget
228 .map(|_| (self.sizer)(&evicted.value))
229 .unwrap_or(0),
230 );
231 }
232 }
233 }
234}
235
236pub struct FsStore {
257 pub(super) root: PathBuf,
258 pub(super) compression: CompressionConfig,
259 pub(super) snapshot_delta_search: bool,
260 pack_manager: RwLock<SnapshotPackManager>,
261 npk1_manager: RwLock<Npk1Manager>,
262 pub(super) recent_blobs: RwLock<RecentObjectCache<ContentHash, Blob>>,
263 pub(super) recent_trees: RwLock<RecentObjectCache<ContentHash, Tree>>,
264 pub(super) recent_states: RwLock<RecentObjectCache<StateId, State>>,
265 pub(super) external_source: Option<Arc<dyn super::super::ExternalObjectSource>>,
266 loose_object_write_mode: LooseObjectWriteMode,
267 pending_directory_syncs: Mutex<BTreeSet<PathBuf>>,
268 #[cfg(test)]
269 snapshot_batch_flushes: AtomicUsize,
270 pub(super) verified_loose_blobs: RwLock<RecentObjectCache<ContentHash, ()>>,
291}
292
293impl Clone for FsStore {
294 fn clone(&self) -> Self {
295 let mut cloned = Self::with_compression(&self.root, self.compression);
296 cloned.snapshot_delta_search = self.snapshot_delta_search;
297 cloned.loose_object_write_mode = self.loose_object_write_mode;
298 cloned.external_source = self.external_source.clone();
299 cloned
300 }
301}
302
303impl FsStore {
304 pub fn new(root: impl AsRef<Path>) -> Self {
308 let root = root.as_ref().to_path_buf();
309 let pack_manager = SnapshotPackManager::new(packs_dir(&root));
310 let npk1_manager = Npk1Manager::new(packs_dir(&root));
311 Self {
312 root,
313 compression: CompressionConfig::default(),
314 snapshot_delta_search: false,
315 pack_manager: RwLock::new(pack_manager),
316 npk1_manager: RwLock::new(npk1_manager),
317 recent_blobs: RwLock::new(RecentObjectCache::with_byte_budget(
318 RECENT_BLOB_CACHE_CAPACITY,
319 RECENT_BLOB_CACHE_MAX_TOTAL_BYTES,
320 |blob: &Blob| blob.content().len(),
321 )),
322 recent_trees: RwLock::new(RecentObjectCache::with_capacity(RECENT_TREE_CACHE_CAPACITY)),
323 recent_states: RwLock::new(RecentObjectCache::with_capacity(
324 RECENT_TREE_CACHE_CAPACITY,
325 )),
326 external_source: None,
327 loose_object_write_mode: LooseObjectWriteMode::Durable,
328 pending_directory_syncs: Mutex::new(BTreeSet::new()),
329 #[cfg(test)]
330 snapshot_batch_flushes: AtomicUsize::new(0),
331 verified_loose_blobs: RwLock::new(RecentObjectCache::with_capacity(
332 VERIFIED_LOOSE_BLOB_CACHE_CAPACITY,
333 )),
334 }
335 }
336
337 pub fn with_compression(root: impl AsRef<Path>, compression: CompressionConfig) -> Self {
339 let root = root.as_ref().to_path_buf();
340 let pack_manager = SnapshotPackManager::new(packs_dir(&root));
341 let npk1_manager = Npk1Manager::new(packs_dir(&root));
342 Self {
343 root,
344 compression,
345 snapshot_delta_search: false,
346 pack_manager: RwLock::new(pack_manager),
347 npk1_manager: RwLock::new(npk1_manager),
348 recent_blobs: RwLock::new(RecentObjectCache::with_byte_budget(
349 RECENT_BLOB_CACHE_CAPACITY,
350 RECENT_BLOB_CACHE_MAX_TOTAL_BYTES,
351 |blob: &Blob| blob.content().len(),
352 )),
353 recent_trees: RwLock::new(RecentObjectCache::with_capacity(RECENT_TREE_CACHE_CAPACITY)),
354 recent_states: RwLock::new(RecentObjectCache::with_capacity(
355 RECENT_TREE_CACHE_CAPACITY,
356 )),
357 external_source: None,
358 loose_object_write_mode: LooseObjectWriteMode::Durable,
359 pending_directory_syncs: Mutex::new(BTreeSet::new()),
360 #[cfg(test)]
361 snapshot_batch_flushes: AtomicUsize::new(0),
362 verified_loose_blobs: RwLock::new(RecentObjectCache::with_capacity(
363 VERIFIED_LOOSE_BLOB_CACHE_CAPACITY,
364 )),
365 }
366 }
367
368 pub fn init(&self) -> Result<()> {
370 crate::fs_atomic::create_dir_all_durable(&blobs_dir(&self.root))?;
373 crate::fs_atomic::create_dir_all_durable(&trees_dir(&self.root))?;
374 crate::fs_atomic::create_dir_all_durable(&tree_lineage_dir(&self.root))?;
375 crate::fs_atomic::create_dir_all_durable(&states_dir(&self.root))?;
376 crate::fs_atomic::create_dir_all_durable(&actions_dir(&self.root))?;
377 crate::fs_atomic::create_dir_all_durable(&packs_dir(&self.root))?;
378 Ok(())
379 }
380
381 pub fn root(&self) -> &Path {
383 &self.root
384 }
385
386 pub fn compression(&self) -> CompressionConfig {
388 self.compression
389 }
390
391 pub fn set_compression(&mut self, compression: CompressionConfig) {
393 self.compression = compression;
394 }
395
396 pub fn set_snapshot_delta_search(&mut self, enabled: bool) {
398 self.snapshot_delta_search = enabled;
399 }
400
401 pub fn loose_object_write_mode(&self) -> LooseObjectWriteMode {
402 self.loose_object_write_mode
403 }
404
405 pub fn set_loose_object_write_mode(&mut self, mode: LooseObjectWriteMode) {
406 self.loose_object_write_mode = mode;
407 }
408
409 pub fn set_external_source(&mut self, source: Arc<dyn super::super::ExternalObjectSource>) {
412 self.external_source = Some(source);
413 }
414
415 fn flush_pending_directory_syncs(&self) -> Result<usize> {
416 let pending_dirs = {
417 let mut guard = self.pending_directory_syncs.lock().map_err(|_| {
418 crate::store::HeddleError::Config(
419 "Failed to acquire pending directory sync lock".to_string(),
420 )
421 })?;
422 if guard.is_empty() {
423 return Ok(0);
424 }
425 let dirs = guard.iter().cloned().collect::<Vec<_>>();
426 guard.clear();
427 dirs
428 };
429
430 for (index, dir) in pending_dirs.iter().enumerate() {
431 if let Err(error) = sync_directory(dir) {
432 if let Ok(mut guard) = self.pending_directory_syncs.lock() {
433 guard.extend(pending_dirs[index..].iter().cloned());
434 }
435 return Err(error.into());
436 }
437 }
438
439 Ok(pending_dirs.len())
440 }
441
442 pub fn reload_packs(&self) -> Result<()> {
448 let packs = packs_dir(&self.root);
449 let _ = super::pack_install_journal::recover_pack_install_intents_with_ttl(
450 &packs,
451 Some(super::pack_install_journal::DEFAULT_PACK_INSTALL_INTENT_TTL_SECS),
452 )?;
453 let _ = super::fs_pack::prune_unpaired_pack_files(&packs)?;
455 let mut manager = self.pack_manager.write().map_err(|_| {
456 crate::store::HeddleError::Config("Failed to acquire pack manager lock".to_string())
457 })?;
458 manager.reload()?;
459 drop(manager);
460 let mut npk1 = self.npk1_manager.write().map_err(|_| {
461 crate::store::HeddleError::Config("Failed to acquire NPK1 manager lock".to_string())
462 })?;
463 npk1.reload()
464 }
465
466 pub(super) fn reload_packs_if_stale(&self) -> Result<bool> {
480 let generic_stale = {
482 let manager = self.pack_manager.read().map_err(|_| {
483 crate::store::HeddleError::Config("Failed to acquire pack manager lock".to_string())
484 })?;
485 manager.needs_reload()?
486 };
487 let npk1_stale = {
488 let manager = self.npk1_manager.read().map_err(|_| {
489 crate::store::HeddleError::Config("Failed to acquire NPK1 manager lock".to_string())
490 })?;
491 manager.needs_reload()?
492 };
493 if !generic_stale && !npk1_stale {
494 return Ok(false);
495 }
496 let mut manager = self.pack_manager.write().map_err(|_| {
500 crate::store::HeddleError::Config("Failed to acquire pack manager lock".to_string())
501 })?;
502 let generic_reloaded = manager.reload_if_stale()?;
503 drop(manager);
504 let mut npk1 = self.npk1_manager.write().map_err(|_| {
505 crate::store::HeddleError::Config("Failed to acquire NPK1 manager lock".to_string())
506 })?;
507 let npk1_reloaded = if npk1.needs_reload()? {
508 npk1.reload()?;
509 true
510 } else {
511 false
512 };
513 Ok(generic_reloaded || npk1_reloaded)
514 }
515
516 pub fn pack_manager(&self) -> &RwLock<SnapshotPackManager> {
518 &self.pack_manager
519 }
520
521 pub(super) fn npk1_manager(&self) -> &RwLock<Npk1Manager> {
522 &self.npk1_manager
523 }
524
525 pub fn clear_recent_object_caches(&self) {
526 if let Ok(mut blobs) = self.recent_blobs.write() {
527 *blobs = RecentObjectCache::with_byte_budget(
528 RECENT_BLOB_CACHE_CAPACITY,
529 RECENT_BLOB_CACHE_MAX_TOTAL_BYTES,
530 |blob: &Blob| blob.content().len(),
531 );
532 }
533 if let Ok(mut trees) = self.recent_trees.write() {
534 *trees = RecentObjectCache::with_capacity(RECENT_TREE_CACHE_CAPACITY);
535 }
536 if let Ok(mut states) = self.recent_states.write() {
537 *states = RecentObjectCache::with_capacity(RECENT_TREE_CACHE_CAPACITY);
538 }
539 }
540
541 #[cfg(test)]
549 pub(super) fn evict_recent_blob(&self, hash: &ContentHash) {
550 if let Ok(mut cache) = self.recent_blobs.write() {
551 cache.remove(hash);
552 }
553 }
554
555 pub fn pack_ids(&self) -> Result<Vec<PackObjectId>> {
556 let manager = self.pack_manager.read().map_err(|_| {
557 crate::store::HeddleError::Config("Failed to acquire pack manager lock".to_string())
558 })?;
559 let mut ids = manager.list_all_ids()?;
560 drop(manager);
561 let npk1 = self.npk1_manager.read().map_err(|_| {
562 crate::store::HeddleError::Config("Failed to acquire NPK1 manager lock".to_string())
563 })?;
564 ids.extend(npk1.list_ids()?.into_iter().map(PackObjectId::Hash));
565 ids.sort();
566 ids.dedup();
567 Ok(ids)
568 }
569
570 pub(super) fn write_loose_object_atomic(&self, path: &Path, data: &[u8]) -> Result<()> {
571 let batch_active = SNAPSHOT_WRITE_BATCH_DEPTHS
572 .with(|depths| depths.borrow().get(&self.root).copied().unwrap_or_default() > 0);
573 let configured_mode = if batch_active {
574 LooseObjectWriteMode::BatchDirectorySync
575 } else {
576 self.loose_object_write_mode
577 };
578
579 let mode = match configured_mode {
580 LooseObjectWriteMode::Durable => AtomicWriteMode::Durable,
581 LooseObjectWriteMode::BatchDirectorySync => AtomicWriteMode::BatchDirectorySync,
582 };
583 write_atomic(path, data, mode, Some(&self.pending_directory_syncs))
584 }
585
586 #[allow(dead_code)]
590 pub(super) fn write_pack_atomic(&self, path: &Path, data: &[u8]) -> Result<()> {
591 write_atomic(path, data, AtomicWriteMode::Durable, None)
592 }
593
594 pub(super) fn write_loose_object_cache(&self, path: &Path, data: &[u8]) -> Result<()> {
615 self.write_reconstructible_cache(path, data)
616 }
617
618 pub(super) fn write_reconstructible_cache(&self, path: &Path, data: &[u8]) -> Result<()> {
622 write_atomic(path, data, AtomicWriteMode::NoSync, None)
623 }
624
625 pub(super) fn begin_snapshot_write_batch_impl(&self) -> Result<()> {
626 SNAPSHOT_WRITE_BATCH_DEPTHS.with(|depths| {
627 *depths.borrow_mut().entry(self.root.clone()).or_default() += 1;
628 });
629 Ok(())
630 }
631
632 pub(super) fn flush_snapshot_write_batch_impl(&self) -> Result<()> {
633 let had_batch = SNAPSHOT_WRITE_BATCH_DEPTHS.with(|depths| {
634 let mut depths = depths.borrow_mut();
635 let Some(depth) = depths.get_mut(&self.root) else {
636 return false;
637 };
638 *depth -= 1;
639 if *depth == 0 {
640 depths.remove(&self.root);
641 }
642 true
643 });
644 if !had_batch {
645 return Ok(());
646 }
647
648 #[cfg(test)]
649 self.snapshot_batch_flushes.fetch_add(1, Ordering::Relaxed);
650
651 let _ = self.flush_pending_directory_syncs()?;
658 Ok(())
659 }
660
661 pub(super) fn abort_snapshot_write_batch_impl(&self) {
662 let should_flush = SNAPSHOT_WRITE_BATCH_DEPTHS.with(|depths| {
663 let mut depths = depths.borrow_mut();
664 let Some(depth) = depths.get_mut(&self.root) else {
665 return true;
669 };
670 *depth -= 1;
671 if *depth == 0 {
672 depths.remove(&self.root);
673 true
674 } else {
675 false
676 }
677 });
678 if should_flush {
683 let _ = self.flush_pending_directory_syncs();
684 }
685 }
686
687 #[cfg(test)]
688 pub(super) fn pending_directory_sync_count(&self) -> usize {
689 self.pending_directory_syncs
690 .lock()
691 .map(|pending| pending.len())
692 .unwrap_or(0)
693 }
694
695 #[cfg(test)]
696 pub(super) fn snapshot_batch_flush_count(&self) -> usize {
697 self.snapshot_batch_flushes.load(Ordering::Relaxed)
698 }
699}