1use std::{
5 fs, io,
6 path::{Path, PathBuf},
7 sync::{
8 Arc, Mutex, Once,
9 atomic::{AtomicBool, Ordering},
10 },
11 thread::sleep,
12};
13
14use dashmap::{DashMap, DashSet};
15use lru_mem::{LruCache, entry_size};
16use threadpool::ThreadPool;
17
18use crate::{
19 hash::ObjectHash,
20 internal::pack::cache_object::{ArcWrapper, CacheObject, FileLoadStore, MemSizeRecorder},
21 time_it,
22};
23
24const CACHE_LAYOUT_VERSION: &str = "rkyv-v1";
28
29pub trait _Cache {
33 fn new(mem_size: Option<usize>, tmp_path: PathBuf, thread_num: usize) -> Self
34 where
35 Self: Sized;
36 fn get_hash(&self, offset: usize) -> Option<ObjectHash>;
37 fn insert(&self, offset: usize, hash: ObjectHash, obj: CacheObject) -> Arc<CacheObject>;
38 fn get_by_offset(&self, offset: usize) -> Option<Arc<CacheObject>>;
39 fn get_by_hash(&self, h: ObjectHash) -> Option<Arc<CacheObject>>;
40 fn total_inserted(&self) -> usize;
41 fn memory_used(&self) -> usize;
42 fn clear(&self);
43}
44
45impl lru_mem::HeapSize for ObjectHash {
46 fn heap_size(&self) -> usize {
47 0
48 }
49}
50
51pub struct Caches {
55 map_offset: DashMap<usize, ObjectHash>, hash_set: DashSet<ObjectHash>, resident_hash_set: DashSet<ObjectHash>, lru_cache: Mutex<LruCache<ObjectHash, ArcWrapper<CacheObject>>>,
63 unbounded_cache: Option<DashMap<ObjectHash, Arc<CacheObject>>>,
64 unbounded_offset_cache: Option<DashMap<usize, Arc<CacheObject>>>,
65 mem_size: Option<usize>,
66 tmp_path: PathBuf,
67 path_prefixes: [Once; 256],
68 pool: Arc<ThreadPool>,
69 complete_signal: Arc<AtomicBool>,
70}
71
72impl Caches {
73 fn try_get(&self, hash: ObjectHash) -> Option<Arc<CacheObject>> {
75 let mut map = self.lru_cache.lock().unwrap();
76 map.get(&hash).map(|x| x.data.clone())
77 }
78
79 fn insert_lru_resident(
80 &self,
81 map: &mut LruCache<ObjectHash, ArcWrapper<CacheObject>>,
82 hash: ObjectHash,
83 obj: ArcWrapper<CacheObject>,
84 ) {
85 let size = entry_size(&hash, &obj);
86 if size <= map.max_size() {
87 while map.current_size() + size > map.max_size() {
88 if let Some((evicted_hash, _)) = map.remove_lru() {
89 self.resident_hash_set.remove(&evicted_hash);
90 } else {
91 break;
92 }
93 }
94 }
95
96 if map.insert(hash, obj).is_ok() {
97 self.resident_hash_set.insert(hash);
98 } else {
99 self.resident_hash_set.remove(&hash);
100 }
101 }
102
103 fn get_fallback(&self, hash: ObjectHash) -> io::Result<Arc<CacheObject>> {
106 let path = self.generate_temp_path(&self.tmp_path, hash);
107 let obj = {
109 loop {
110 match Self::read_from_temp(&path) {
111 Ok(x) => break x,
112 Err(e) if e.kind() == io::ErrorKind::NotFound => {
113 sleep(std::time::Duration::from_millis(10));
114 continue;
115 }
116 Err(e) => return Err(e), }
118 }
119 };
120
121 let mut map = self.lru_cache.lock().unwrap();
122 let obj = Arc::new(obj);
123 let mut x = ArcWrapper::new(
124 obj.clone(),
125 self.complete_signal.clone(),
126 Some(self.pool.clone()),
127 );
128 x.set_store_path(path);
129 self.insert_lru_resident(&mut map, hash, x);
130 Ok(obj)
131 }
132
133 fn generate_temp_path(&self, tmp_path: &Path, hash: ObjectHash) -> PathBuf {
135 let mut path =
137 PathBuf::with_capacity(self.tmp_path.capacity() + hash.to_string().len() + 5);
138 path.push(tmp_path);
139 path.push(CACHE_LAYOUT_VERSION);
140 let hash_str = hash._to_string();
141 path.push(&hash_str[..2]); self.path_prefixes[hash.as_ref()[0] as usize].call_once(|| {
143 if !path.exists() {
145 fs::create_dir_all(&path).unwrap();
146 }
147 });
148 path.push(hash_str);
149 path
150 }
151
152 fn read_from_temp(path: &Path) -> io::Result<CacheObject> {
154 let obj = CacheObject::f_load(path)?;
155 obj.record_mem_size();
158 Ok(obj)
159 }
160
161 pub fn queued_tasks(&self) -> usize {
163 self.pool.queued_count()
164 }
165
166 pub(crate) fn is_unbounded(&self) -> bool {
167 self.mem_size.is_none()
168 }
169
170 pub fn memory_used_index(&self) -> usize {
172 let hash_cache_size = if let Some(cache) = &self.unbounded_cache {
173 cache.capacity()
174 * (std::mem::size_of::<ObjectHash>() + std::mem::size_of::<Arc<CacheObject>>())
175 } else {
176 self.hash_set.capacity() * std::mem::size_of::<ObjectHash>()
177 };
178 let offset_cache_size = if let Some(cache) = &self.unbounded_offset_cache {
179 cache.capacity()
180 * (std::mem::size_of::<usize>() + std::mem::size_of::<Arc<CacheObject>>())
181 } else {
182 self.map_offset.capacity()
183 * (std::mem::size_of::<usize>() + std::mem::size_of::<ObjectHash>())
184 };
185 hash_cache_size + offset_cache_size
186 }
187
188 pub(crate) fn shutdown(&self) {
189 time_it!("Caches clear", {
190 self.complete_signal.store(true, Ordering::Release);
191 self.pool.join();
192 self.lru_cache
193 .lock()
194 .unwrap_or_else(|e| e.into_inner())
195 .clear();
196 if let Some(cache) = &self.unbounded_cache {
197 cache.clear();
198 cache.shrink_to_fit();
199 }
200 if let Some(cache) = &self.unbounded_offset_cache {
201 cache.clear();
202 cache.shrink_to_fit();
203 }
204 self.hash_set.clear();
205 self.hash_set.shrink_to_fit();
206 self.resident_hash_set.clear();
207 self.resident_hash_set.shrink_to_fit();
208 self.map_offset.clear();
209 self.map_offset.shrink_to_fit();
210 });
211 }
212
213 pub fn remove_tmp_dir(&self) -> io::Result<()> {
215 time_it!("Remove tmp dir", {
216 if self.tmp_path.exists() {
217 match fs::remove_dir_all(&self.tmp_path) {
218 Ok(()) => {}
219 Err(e) if e.kind() == io::ErrorKind::NotFound => {}
220 Err(e) => return Err(e),
221 }
222
223 if let Some(parent) = self.tmp_path.parent() {
224 let is_cache_temp = parent
225 .file_name()
226 .and_then(|n| n.to_str())
227 .map(|n| n == ".cache_temp")
228 .unwrap_or(false);
229 if is_cache_temp {
230 match fs::remove_dir(parent) {
231 Ok(()) => {}
232 Err(e)
233 if matches!(
234 e.kind(),
235 io::ErrorKind::DirectoryNotEmpty | io::ErrorKind::NotFound
236 ) => {}
237 Err(e) => return Err(e),
238 }
239 }
240 }
241 }
242
243 Ok(())
244 })
245 }
246
247 pub fn remove_unbounded(&self, offset: usize, hash: ObjectHash) {
248 if self.unbounded_cache.is_some() && self.lru_cache.lock().unwrap().remove(&hash).is_some()
249 {
250 self.resident_hash_set.remove(&hash);
251 }
252 if let Some(cache) = &self.unbounded_offset_cache {
253 cache.remove(&offset);
254 }
255 if let Some(cache) = &self.unbounded_cache {
256 cache.remove(&hash);
257 }
258 }
259}
260
261impl _Cache for Caches {
262 fn new(mem_size: Option<usize>, tmp_path: PathBuf, thread_num: usize) -> Self
265 where
266 Self: Sized,
267 {
268 if mem_size.is_some() {
270 fs::create_dir_all(&tmp_path).unwrap();
271 }
272
273 Caches {
274 map_offset: DashMap::new(),
275 hash_set: DashSet::new(),
276 resident_hash_set: DashSet::new(),
277 lru_cache: Mutex::new(LruCache::new(mem_size.unwrap_or(usize::MAX))),
278 unbounded_cache: mem_size.is_none().then(DashMap::new),
279 unbounded_offset_cache: mem_size.is_none().then(DashMap::new),
280 mem_size,
281 tmp_path,
282 path_prefixes: [const { Once::new() }; 256],
283 pool: Arc::new(ThreadPool::new(thread_num)),
284 complete_signal: Arc::new(AtomicBool::new(false)),
285 }
286 }
287
288 fn get_hash(&self, offset: usize) -> Option<ObjectHash> {
289 if let Some(cache) = &self.unbounded_offset_cache {
290 return cache.get(&offset).and_then(|obj| obj.base_object_hash());
291 }
292 self.map_offset.get(&offset).map(|x| *x)
293 }
294
295 fn insert(&self, offset: usize, hash: ObjectHash, obj: CacheObject) -> Arc<CacheObject> {
296 let obj_arc = Arc::new(obj);
297 if let Some(cache) = &self.unbounded_cache {
298 cache.insert(hash, obj_arc.clone());
299 if let Some(offset_cache) = &self.unbounded_offset_cache {
300 offset_cache.insert(offset, obj_arc.clone());
301 }
302 let mut map = self.lru_cache.lock().unwrap();
303 let a_obj = ArcWrapper::new(
304 obj_arc.clone(),
305 self.complete_signal.clone(),
306 Some(self.pool.clone()),
307 );
308 self.insert_lru_resident(&mut map, hash, a_obj);
309 } else {
310 {
320 let mut map = self.lru_cache.lock().unwrap();
321 let mut a_obj = ArcWrapper::new(
322 obj_arc.clone(),
323 self.complete_signal.clone(),
324 Some(self.pool.clone()),
325 );
326 if self.mem_size.is_some() {
327 a_obj.set_store_path(self.generate_temp_path(&self.tmp_path, hash));
328 }
329 self.insert_lru_resident(&mut map, hash, a_obj);
330 }
331 self.hash_set.insert(hash);
332 self.map_offset.insert(offset, hash);
335 }
336
337 obj_arc
338 }
339
340 fn get_by_offset(&self, offset: usize) -> Option<Arc<CacheObject>> {
342 if let Some(cache) = &self.unbounded_offset_cache {
351 let hash = cache.get(&offset).and_then(|obj| obj.base_object_hash());
352 return hash.and_then(|hash| self.try_get(hash));
353 }
354
355 let hash = self.map_offset.get(&offset).map(|x| *x);
356 hash.and_then(|hash| self.get_by_hash(hash))
357 }
358
359 fn get_by_hash(&self, hash: ObjectHash) -> Option<Arc<CacheObject>> {
361 if self.mem_size.is_none() {
362 if let Some(cache) = &self.unbounded_cache
363 && !cache.contains_key(&hash)
364 {
365 return None;
366 }
367 return self.try_get(hash);
368 }
369
370 if self.hash_set.contains(&hash) {
372 if !self.resident_hash_set.contains(&hash) {
373 return self.get_fallback(hash).ok();
374 }
375 match self.try_get(hash) {
376 Some(x) => Some(x),
377 None => {
378 if self.mem_size.is_none() {
379 panic!("should not be here when mem_size is not set")
380 }
381 self.get_fallback(hash).ok()
382 }
383 }
384 } else {
385 None
386 }
387 }
388
389 fn total_inserted(&self) -> usize {
390 if self.mem_size.is_some() {
391 self.hash_set.len()
392 } else if let Some(cache) = &self.unbounded_offset_cache {
393 cache.len()
394 } else {
395 self.map_offset.len()
396 }
397 }
398 fn memory_used(&self) -> usize {
399 self.lru_cache.lock().unwrap().current_size() + self.memory_used_index()
400 }
401 fn clear(&self) {
402 self.shutdown();
403
404 assert_eq!(self.pool.queued_count(), 0);
405 assert_eq!(self.pool.active_count(), 0);
406 assert_eq!(self.lru_cache.lock().unwrap().len(), 0);
407 }
408}
409
410#[cfg(test)]
411mod test {
412 use std::{env, sync::Arc, thread};
413
414 use super::*;
415 use crate::{
416 hash::{HashKind, ObjectHash, set_hash_kind_for_test},
417 internal::{object::types::ObjectType, pack::cache_object::CacheObjectInfo},
418 };
419
420 fn make_obj(size: usize, hash: ObjectHash) -> CacheObject {
422 CacheObject {
423 info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
424 data_decompressed: vec![0; size],
425 mem_recorder: None,
426 offset: 0,
427 crc32: 0,
428 is_delta_in_pack: false,
429 known_hash: None,
430 }
431 }
432
433 #[test]
435 fn test_cache_single_thread() {
436 for (kind, cap, size_ab, size_c, tmp_dir) in [
437 (
438 HashKind::Sha1,
439 2048usize,
440 800usize,
441 1700usize,
442 "tests/.cache_tmp",
443 ),
444 (
445 HashKind::Sha256,
446 4096usize,
447 1500usize,
448 3000usize,
449 "tests/.cache_tmp_sha256",
450 ),
451 ] {
452 let _guard = set_hash_kind_for_test(kind);
453 let source = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
454 let tmp_path = source.clone().join(tmp_dir);
455 if tmp_path.exists() {
456 fs::remove_dir_all(&tmp_path).unwrap();
457 }
458
459 let cache = Caches::new(Some(cap), tmp_path, 1);
460 let a_hash = ObjectHash::new(String::from("a").as_bytes());
461 let b_hash = ObjectHash::new(String::from("b").as_bytes());
462 let c_hash = ObjectHash::new(String::from("c").as_bytes());
463
464 let a = make_obj(size_ab, a_hash);
465 let b = make_obj(size_ab, b_hash);
466 let c = make_obj(size_c, c_hash);
467
468 cache.insert(a.offset, a_hash, a.clone());
470 assert!(cache.hash_set.contains(&a_hash));
471 assert!(cache.try_get(a_hash).is_some());
472
473 cache.insert(b.offset, b_hash, b.clone());
475 assert!(cache.hash_set.contains(&b_hash));
476 assert!(cache.try_get(b_hash).is_some());
477 assert!(cache.try_get(a_hash).is_some());
478
479 cache.insert(c.offset, c_hash, c.clone());
481 assert!(cache.try_get(a_hash).is_none());
482 assert!(cache.try_get(b_hash).is_none());
483 assert!(cache.try_get(c_hash).is_some());
484 assert!(cache.get_by_hash(c_hash).is_some());
485 }
486 }
487
488 #[test]
490 fn test_cache_multi_thread_mixed_hash_kinds() {
491 let base = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
492 let tmp_path = base.join("tests/.cache_tmp_mixed");
493 if tmp_path.exists() {
494 fs::remove_dir_all(&tmp_path).unwrap();
495 }
496
497 let cache = Arc::new(Caches::new(Some(4096), tmp_path, 2));
498
499 let cache_sha1 = Arc::clone(&cache);
500 let handle_sha1 = thread::spawn(move || {
501 let _g = set_hash_kind_for_test(HashKind::Sha1);
502 let hash = ObjectHash::new(b"sha1-entry");
503 let obj = CacheObject {
504 info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
505 data_decompressed: vec![0; 800],
506 mem_recorder: None,
507 offset: 1,
508 crc32: 0,
509 is_delta_in_pack: false,
510 known_hash: None,
511 };
512 cache_sha1.insert(obj.offset, hash, obj.clone());
513 assert!(cache_sha1.hash_set.contains(&hash));
514 assert!(cache_sha1.try_get(hash).is_some());
515 });
516
517 let cache_sha256 = Arc::clone(&cache);
518 let handle_sha256 = thread::spawn(move || {
519 let _g = set_hash_kind_for_test(HashKind::Sha256);
520 let hash = ObjectHash::new(b"sha256-entry");
521 let obj = CacheObject {
522 info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
523 data_decompressed: vec![0; 1500],
524 mem_recorder: None,
525 offset: 2,
526 crc32: 0,
527 is_delta_in_pack: false,
528 known_hash: None,
529 };
530 cache_sha256.insert(obj.offset, hash, obj.clone());
531 assert!(cache_sha256.hash_set.contains(&hash));
532 assert!(cache_sha256.try_get(hash).is_some());
533 });
534
535 handle_sha1.join().unwrap();
536 handle_sha256.join().unwrap();
537
538 assert_eq!(cache.total_inserted(), 2);
539 }
540
541 #[test]
542 fn test_remove_tmp_dir_does_not_panic_when_cleanup_fails() {
543 let dir = tempfile::tempdir().unwrap();
544 let tmp_path = dir.path().join("not-a-directory");
545 fs::write(&tmp_path, b"cache marker").unwrap();
546 let cache = Caches::new(None, tmp_path, 1);
547
548 let result =
549 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| cache.remove_tmp_dir()));
550
551 assert!(result.is_ok(), "cache cleanup should not panic");
552 assert!(
553 result.unwrap().is_err(),
554 "cleanup failure should be returned"
555 );
556 }
557
558 #[test]
559 fn test_unbounded_cache_skips_hash_set_index() {
560 let _guard = set_hash_kind_for_test(HashKind::Sha1);
561 let base = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
562 let tmp_path = base.join("tests/.cache_tmp_unbounded");
563 let cache = Caches::new(None, tmp_path, 1);
564 let hash = ObjectHash::new(b"unbounded-entry");
565 let obj = make_obj(64, hash);
566
567 cache.insert(obj.offset, hash, obj);
568
569 assert!(cache.hash_set.is_empty());
570 assert!(cache.resident_hash_set.contains(&hash));
571 assert_eq!(cache.total_inserted(), 1);
572 assert_eq!(cache.get_hash(0), Some(hash));
573 assert!(cache.try_get(hash).is_some());
574 assert!(cache.get_by_hash(hash).is_some());
575 assert!(
576 cache
577 .get_by_hash(ObjectHash::new(b"missing-entry"))
578 .is_none()
579 );
580 assert!(cache.get_by_offset(0).is_some());
581 assert!(cache.memory_used_index() > 0);
582 }
583
584 #[test]
585 fn test_bounded_cache_tracks_resident_entries() {
586 let _guard = set_hash_kind_for_test(HashKind::Sha1);
587 let base = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
588 let tmp_path = base.join("tests/.cache_tmp_resident");
589 if tmp_path.exists() {
590 fs::remove_dir_all(&tmp_path).unwrap();
591 }
592
593 let cache = Caches::new(Some(2048), tmp_path, 1);
594 let a_hash = ObjectHash::new(b"resident-a");
595 let b_hash = ObjectHash::new(b"resident-b");
596 let a = make_obj(1800, a_hash);
597 let b = make_obj(1800, b_hash);
598
599 cache.insert(a.offset, a_hash, a);
600 assert!(cache.hash_set.contains(&a_hash));
601 assert!(cache.resident_hash_set.contains(&a_hash));
602
603 cache.insert(b.offset + 1, b_hash, b);
604 assert!(cache.hash_set.contains(&a_hash));
605 assert!(cache.hash_set.contains(&b_hash));
606 assert!(!cache.resident_hash_set.contains(&a_hash));
607 assert!(cache.resident_hash_set.contains(&b_hash));
608 }
609
610 #[test]
626 fn test_cache_concurrent_insert_get_by_offset_no_deadlock() {
627 let _guard = set_hash_kind_for_test(HashKind::Sha1);
628 let tmp = std::env::temp_dir().join(format!(
629 "gi_cache_abba_{}_{:?}",
630 std::process::id(),
631 thread::current().id()
632 ));
633 if tmp.exists() {
634 let _ = fs::remove_dir_all(&tmp);
635 }
636
637 let cache = Arc::new(Caches::new(Some(1 << 20), tmp.clone(), 8));
643
644 const OFFSETS: usize = 8;
647 const ITERS: usize = 20_000;
648 const WRITERS: usize = 6;
649 const READERS: usize = 6;
650
651 let hashes: Arc<Vec<ObjectHash>> = Arc::new(
652 (0..OFFSETS)
653 .map(|i| ObjectHash::new(format!("obj-{i}").as_bytes()))
654 .collect(),
655 );
656
657 let (done_tx, done_rx) = std::sync::mpsc::channel();
658 let mut handles = Vec::new();
659
660 for _ in 0..WRITERS {
661 let cache = cache.clone();
662 let hashes = hashes.clone();
663 let done_tx = done_tx.clone();
664 handles.push(thread::spawn(move || {
665 let _g = set_hash_kind_for_test(HashKind::Sha1);
666 for k in 0..ITERS {
667 let o = k % OFFSETS;
668 let obj = CacheObject {
669 info: CacheObjectInfo::BaseObject(ObjectType::Blob, hashes[o]),
670 data_decompressed: vec![0u8; 64],
671 mem_recorder: None,
672 offset: o,
673 crc32: 0,
674 is_delta_in_pack: false,
675 known_hash: None,
676 };
677 cache.insert(o, hashes[o], obj);
678 }
679 let _ = done_tx.send(());
680 }));
681 }
682
683 for _ in 0..READERS {
684 let cache = cache.clone();
685 let done_tx = done_tx.clone();
686 handles.push(thread::spawn(move || {
687 let _g = set_hash_kind_for_test(HashKind::Sha1);
688 for k in 0..ITERS {
689 let o = k % OFFSETS;
690 let _ = cache.get_by_offset(o);
691 }
692 let _ = done_tx.send(());
693 }));
694 }
695 drop(done_tx);
696
697 let workers = WRITERS + READERS;
701 for finished in 0..workers {
702 if done_rx
703 .recv_timeout(std::time::Duration::from_secs(30))
704 .is_err()
705 {
706 panic!(
707 "cache deadlock: only {finished}/{workers} workers finished within 30s — \
708 the lru_cache <-> map_offset ABBA lock-order inversion has regressed"
709 );
710 }
711 }
712
713 for h in handles {
714 h.join().unwrap();
715 }
716 cache.clear();
719 let _ = fs::remove_dir_all(&tmp);
720 }
721}