git_internal/internal/pack/
cache.rs

1//! Multi-tier cache for pack decoding that combines an in-memory LRU with spill-to-disk storage and
2//! bookkeeping for concurrent rebuild tasks.
3
4use 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
24/// Cache format version appended to the disk path so that caches written with an
25/// incompatible serialization format (for example the previous bincode layout)
26/// are ignored instead of causing deserialization errors.
27const CACHE_LAYOUT_VERSION: &str = "rkyv-v1";
28
29/// Trait defining the interface for a multi-tier cache system.
30/// This cache supports insertion and retrieval of objects by both offset and hash,
31/// as well as memory usage tracking and clearing functionality.
32pub 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
51/// Multi-tier cache implementation combining an in-memory LRU cache with spill-to-disk storage.
52/// It uses a DashMap for offset-to-hash mapping and a DashSet to track cached hashes.
53/// The cache supports concurrent rebuild tasks using a thread pool.
54pub struct Caches {
55    map_offset: DashMap<usize, ObjectHash>, // offset to hash
56    hash_set: DashSet<ObjectHash>,          // item in the cache
57    resident_hash_set: DashSet<ObjectHash>, // item currently held in the in-memory LRU
58    // dropping large lru cache will take a long time on Windows without multi-thread IO
59    // because "multi-thread IO" clone Arc<CacheObject>, so it won't be dropped in the main thread,
60    // and `CacheObjects` will be killed by OS after Process ends abnormally
61    // Solution: use `mimalloc`
62    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    /// only get object from memory, not from tmp file
74    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    /// !IMPORTANT: because of the process of pack, the file must be written / be writing before, so it won't be dead lock
104    /// fall back to temp to get item. **invoker should ensure the hash is in the cache, or it will block forever**
105    fn get_fallback(&self, hash: ObjectHash) -> io::Result<Arc<CacheObject>> {
106        let path = self.generate_temp_path(&self.tmp_path, hash);
107        // read from tmp file
108        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), // other error
117                }
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    /// generate the temp file path, hex string of the hash
134    fn generate_temp_path(&self, tmp_path: &Path, hash: ObjectHash) -> PathBuf {
135        // Reserve capacity for base path, 2-char subdir, hex hash string, and separators
136        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]); // use first 2 chars as the directory
142        self.path_prefixes[hash.as_ref()[0] as usize].call_once(|| {
143            // Check if the directory exists, if not, create it
144            if !path.exists() {
145                fs::create_dir_all(&path).unwrap();
146            }
147        });
148        path.push(hash_str);
149        path
150    }
151
152    /// read CacheObject from temp file
153    fn read_from_temp(path: &Path) -> io::Result<CacheObject> {
154        let obj = CacheObject::f_load(path)?;
155        // Deserializing will also create an object but without Construction outside and `::new()`
156        // So if you want to do sth. while Constructing, impl Deserialize trait yourself
157        obj.record_mem_size();
158        Ok(obj)
159    }
160
161    /// number of queued tasks in the thread pool
162    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    /// memory used by the index (exclude lru_cache which is contained in CacheObject::get_mem_size())
171    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    /// remove the tmp dir
214    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    /// @param size: the size of the memory lru cache. **None means no limit**
263    /// @param tmp_path: the path to store the cache object in the tmp file
264    fn new(mem_size: Option<usize>, tmp_path: PathBuf, thread_num: usize) -> Self
265    where
266        Self: Sized,
267    {
268        // `None` means no limit, so no need to create the tmp dir
269        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            // ? whether insert to cache directly or only write to tmp file
311            let mut map = self.lru_cache.lock().unwrap();
312            let mut a_obj = ArcWrapper::new(
313                obj_arc.clone(),
314                self.complete_signal.clone(),
315                Some(self.pool.clone()),
316            );
317            if self.mem_size.is_some() {
318                a_obj.set_store_path(self.generate_temp_path(&self.tmp_path, hash));
319            }
320            self.insert_lru_resident(&mut map, hash, a_obj);
321            self.hash_set.insert(hash);
322            // order matters as for reading in 'get_by_offset()'
323            self.map_offset.insert(offset, hash);
324        }
325
326        obj_arc
327    }
328
329    /// get object by offset, from memory or tmp file
330    fn get_by_offset(&self, offset: usize) -> Option<Arc<CacheObject>> {
331        if let Some(cache) = &self.unbounded_offset_cache {
332            return cache
333                .get(&offset)
334                .and_then(|obj| obj.base_object_hash())
335                .and_then(|hash| self.try_get(hash));
336        }
337
338        match self.map_offset.get(&offset) {
339            Some(x) => self.get_by_hash(*x),
340            None => None,
341        }
342    }
343
344    /// get object by hash, from memory or tmp file
345    fn get_by_hash(&self, hash: ObjectHash) -> Option<Arc<CacheObject>> {
346        if self.mem_size.is_none() {
347            if let Some(cache) = &self.unbounded_cache
348                && !cache.contains_key(&hash)
349            {
350                return None;
351            }
352            return self.try_get(hash);
353        }
354
355        // check if the hash is in the cache( lru or tmp file)
356        if self.hash_set.contains(&hash) {
357            if !self.resident_hash_set.contains(&hash) {
358                return self.get_fallback(hash).ok();
359            }
360            match self.try_get(hash) {
361                Some(x) => Some(x),
362                None => {
363                    if self.mem_size.is_none() {
364                        panic!("should not be here when mem_size is not set")
365                    }
366                    self.get_fallback(hash).ok()
367                }
368            }
369        } else {
370            None
371        }
372    }
373
374    fn total_inserted(&self) -> usize {
375        if self.mem_size.is_some() {
376            self.hash_set.len()
377        } else if let Some(cache) = &self.unbounded_offset_cache {
378            cache.len()
379        } else {
380            self.map_offset.len()
381        }
382    }
383    fn memory_used(&self) -> usize {
384        self.lru_cache.lock().unwrap().current_size() + self.memory_used_index()
385    }
386    fn clear(&self) {
387        self.shutdown();
388
389        assert_eq!(self.pool.queued_count(), 0);
390        assert_eq!(self.pool.active_count(), 0);
391        assert_eq!(self.lru_cache.lock().unwrap().len(), 0);
392    }
393}
394
395#[cfg(test)]
396mod test {
397    use std::{env, sync::Arc, thread};
398
399    use super::*;
400    use crate::{
401        hash::{HashKind, ObjectHash, set_hash_kind_for_test},
402        internal::{object::types::ObjectType, pack::cache_object::CacheObjectInfo},
403    };
404
405    /// Helper to build a base CacheObject with given size and hash.
406    fn make_obj(size: usize, hash: ObjectHash) -> CacheObject {
407        CacheObject {
408            info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
409            data_decompressed: vec![0; size],
410            mem_recorder: None,
411            offset: 0,
412            crc32: 0,
413            is_delta_in_pack: false,
414            known_hash: None,
415        }
416    }
417
418    /// test single-threaded cache behavior with different hash kinds and capacities
419    #[test]
420    fn test_cache_single_thread() {
421        for (kind, cap, size_ab, size_c, tmp_dir) in [
422            (
423                HashKind::Sha1,
424                2048usize,
425                800usize,
426                1700usize,
427                "tests/.cache_tmp",
428            ),
429            (
430                HashKind::Sha256,
431                4096usize,
432                1500usize,
433                3000usize,
434                "tests/.cache_tmp_sha256",
435            ),
436        ] {
437            let _guard = set_hash_kind_for_test(kind);
438            let source = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
439            let tmp_path = source.clone().join(tmp_dir);
440            if tmp_path.exists() {
441                fs::remove_dir_all(&tmp_path).unwrap();
442            }
443
444            let cache = Caches::new(Some(cap), tmp_path, 1);
445            let a_hash = ObjectHash::new(String::from("a").as_bytes());
446            let b_hash = ObjectHash::new(String::from("b").as_bytes());
447            let c_hash = ObjectHash::new(String::from("c").as_bytes());
448
449            let a = make_obj(size_ab, a_hash);
450            let b = make_obj(size_ab, b_hash);
451            let c = make_obj(size_c, c_hash);
452
453            // insert a
454            cache.insert(a.offset, a_hash, a.clone());
455            assert!(cache.hash_set.contains(&a_hash));
456            assert!(cache.try_get(a_hash).is_some());
457
458            // insert b, a should still be in cache
459            cache.insert(b.offset, b_hash, b.clone());
460            assert!(cache.hash_set.contains(&b_hash));
461            assert!(cache.try_get(b_hash).is_some());
462            assert!(cache.try_get(a_hash).is_some());
463
464            // insert c which will evict both a and b
465            cache.insert(c.offset, c_hash, c.clone());
466            assert!(cache.try_get(a_hash).is_none());
467            assert!(cache.try_get(b_hash).is_none());
468            assert!(cache.try_get(c_hash).is_some());
469            assert!(cache.get_by_hash(c_hash).is_some());
470        }
471    }
472
473    /// consider the multi-threaded scenario where different threads use different hash kinds
474    #[test]
475    fn test_cache_multi_thread_mixed_hash_kinds() {
476        let base = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
477        let tmp_path = base.join("tests/.cache_tmp_mixed");
478        if tmp_path.exists() {
479            fs::remove_dir_all(&tmp_path).unwrap();
480        }
481
482        let cache = Arc::new(Caches::new(Some(4096), tmp_path, 2));
483
484        let cache_sha1 = Arc::clone(&cache);
485        let handle_sha1 = thread::spawn(move || {
486            let _g = set_hash_kind_for_test(HashKind::Sha1);
487            let hash = ObjectHash::new(b"sha1-entry");
488            let obj = CacheObject {
489                info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
490                data_decompressed: vec![0; 800],
491                mem_recorder: None,
492                offset: 1,
493                crc32: 0,
494                is_delta_in_pack: false,
495                known_hash: None,
496            };
497            cache_sha1.insert(obj.offset, hash, obj.clone());
498            assert!(cache_sha1.hash_set.contains(&hash));
499            assert!(cache_sha1.try_get(hash).is_some());
500        });
501
502        let cache_sha256 = Arc::clone(&cache);
503        let handle_sha256 = thread::spawn(move || {
504            let _g = set_hash_kind_for_test(HashKind::Sha256);
505            let hash = ObjectHash::new(b"sha256-entry");
506            let obj = CacheObject {
507                info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
508                data_decompressed: vec![0; 1500],
509                mem_recorder: None,
510                offset: 2,
511                crc32: 0,
512                is_delta_in_pack: false,
513                known_hash: None,
514            };
515            cache_sha256.insert(obj.offset, hash, obj.clone());
516            assert!(cache_sha256.hash_set.contains(&hash));
517            assert!(cache_sha256.try_get(hash).is_some());
518        });
519
520        handle_sha1.join().unwrap();
521        handle_sha256.join().unwrap();
522
523        assert_eq!(cache.total_inserted(), 2);
524    }
525
526    #[test]
527    fn test_remove_tmp_dir_does_not_panic_when_cleanup_fails() {
528        let dir = tempfile::tempdir().unwrap();
529        let tmp_path = dir.path().join("not-a-directory");
530        fs::write(&tmp_path, b"cache marker").unwrap();
531        let cache = Caches::new(None, tmp_path, 1);
532
533        let result =
534            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| cache.remove_tmp_dir()));
535
536        assert!(result.is_ok(), "cache cleanup should not panic");
537        assert!(
538            result.unwrap().is_err(),
539            "cleanup failure should be returned"
540        );
541    }
542
543    #[test]
544    fn test_unbounded_cache_skips_hash_set_index() {
545        let _guard = set_hash_kind_for_test(HashKind::Sha1);
546        let base = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
547        let tmp_path = base.join("tests/.cache_tmp_unbounded");
548        let cache = Caches::new(None, tmp_path, 1);
549        let hash = ObjectHash::new(b"unbounded-entry");
550        let obj = make_obj(64, hash);
551
552        cache.insert(obj.offset, hash, obj);
553
554        assert!(cache.hash_set.is_empty());
555        assert!(cache.resident_hash_set.contains(&hash));
556        assert_eq!(cache.total_inserted(), 1);
557        assert_eq!(cache.get_hash(0), Some(hash));
558        assert!(cache.try_get(hash).is_some());
559        assert!(cache.get_by_hash(hash).is_some());
560        assert!(
561            cache
562                .get_by_hash(ObjectHash::new(b"missing-entry"))
563                .is_none()
564        );
565        assert!(cache.get_by_offset(0).is_some());
566        assert!(cache.memory_used_index() > 0);
567    }
568
569    #[test]
570    fn test_bounded_cache_tracks_resident_entries() {
571        let _guard = set_hash_kind_for_test(HashKind::Sha1);
572        let base = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
573        let tmp_path = base.join("tests/.cache_tmp_resident");
574        if tmp_path.exists() {
575            fs::remove_dir_all(&tmp_path).unwrap();
576        }
577
578        let cache = Caches::new(Some(2048), tmp_path, 1);
579        let a_hash = ObjectHash::new(b"resident-a");
580        let b_hash = ObjectHash::new(b"resident-b");
581        let a = make_obj(1800, a_hash);
582        let b = make_obj(1800, b_hash);
583
584        cache.insert(a.offset, a_hash, a);
585        assert!(cache.hash_set.contains(&a_hash));
586        assert!(cache.resident_hash_set.contains(&a_hash));
587
588        cache.insert(b.offset + 1, b_hash, b);
589        assert!(cache.hash_set.contains(&a_hash));
590        assert!(cache.hash_set.contains(&b_hash));
591        assert!(!cache.resident_hash_set.contains(&a_hash));
592        assert!(cache.resident_hash_set.contains(&b_hash));
593    }
594}
git_internal/internal/pack/cache.rs

git_internal/internal/pack/
cache.rs
