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;
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    // dropping large lru cache will take a long time on Windows without multi-thread IO
58    // because "multi-thread IO" clone Arc<CacheObject>, so it won't be dropped in the main thread,
59    // and `CacheObjects` will be killed by OS after Process ends abnormally
60    // Solution: use `mimalloc`
61    lru_cache: Mutex<LruCache<ObjectHash, ArcWrapper<CacheObject>>>,
62    mem_size: Option<usize>,
63    tmp_path: PathBuf,
64    path_prefixes: [Once; 256],
65    pool: Arc<ThreadPool>,
66    complete_signal: Arc<AtomicBool>,
67}
68
69impl Caches {
70    /// only get object from memory, not from tmp file
71    fn try_get(&self, hash: ObjectHash) -> Option<Arc<CacheObject>> {
72        let mut map = self.lru_cache.lock().unwrap();
73        map.get(&hash).map(|x| x.data.clone())
74    }
75
76    /// !IMPORTANT: because of the process of pack, the file must be written / be writing before, so it won't be dead lock
77    /// fall back to temp to get item. **invoker should ensure the hash is in the cache, or it will block forever**
78    fn get_fallback(&self, hash: ObjectHash) -> io::Result<Arc<CacheObject>> {
79        let path = self.generate_temp_path(&self.tmp_path, hash);
80        // read from tmp file
81        let obj = {
82            loop {
83                match Self::read_from_temp(&path) {
84                    Ok(x) => break x,
85                    Err(e) if e.kind() == io::ErrorKind::NotFound => {
86                        sleep(std::time::Duration::from_millis(10));
87                        continue;
88                    }
89                    Err(e) => return Err(e), // other error
90                }
91            }
92        };
93
94        let mut map = self.lru_cache.lock().unwrap();
95        let obj = Arc::new(obj);
96        let mut x = ArcWrapper::new(
97            obj.clone(),
98            self.complete_signal.clone(),
99            Some(self.pool.clone()),
100        );
101        x.set_store_path(path);
102        let _ = map.insert(hash, x); // handle the error
103        Ok(obj)
104    }
105
106    /// generate the temp file path, hex string of the hash
107    fn generate_temp_path(&self, tmp_path: &Path, hash: ObjectHash) -> PathBuf {
108        // Reserve capacity for base path, 2-char subdir, hex hash string, and separators
109        let mut path =
110            PathBuf::with_capacity(self.tmp_path.capacity() + hash.to_string().len() + 5);
111        path.push(tmp_path);
112        path.push(CACHE_LAYOUT_VERSION);
113        let hash_str = hash._to_string();
114        path.push(&hash_str[..2]); // use first 2 chars as the directory
115        self.path_prefixes[hash.as_ref()[0] as usize].call_once(|| {
116            // Check if the directory exists, if not, create it
117            if !path.exists() {
118                fs::create_dir_all(&path).unwrap();
119            }
120        });
121        path.push(hash_str);
122        path
123    }
124
125    /// read CacheObject from temp file
126    fn read_from_temp(path: &Path) -> io::Result<CacheObject> {
127        let obj = CacheObject::f_load(path)?;
128        // Deserializing will also create an object but without Construction outside and `::new()`
129        // So if you want to do sth. while Constructing, impl Deserialize trait yourself
130        obj.record_mem_size();
131        Ok(obj)
132    }
133
134    /// number of queued tasks in the thread pool
135    pub fn queued_tasks(&self) -> usize {
136        self.pool.queued_count()
137    }
138
139    /// memory used by the index (exclude lru_cache which is contained in CacheObject::get_mem_size())
140    pub fn memory_used_index(&self) -> usize {
141        self.map_offset.capacity()
142            * (std::mem::size_of::<usize>() + std::mem::size_of::<ObjectHash>())
143            + self.hash_set.capacity() * (std::mem::size_of::<ObjectHash>())
144    }
145
146    pub(crate) fn shutdown(&self) {
147        time_it!("Caches clear", {
148            self.complete_signal.store(true, Ordering::Release);
149            self.pool.join();
150            self.lru_cache
151                .lock()
152                .unwrap_or_else(|e| e.into_inner())
153                .clear();
154            self.hash_set.clear();
155            self.hash_set.shrink_to_fit();
156            self.map_offset.clear();
157            self.map_offset.shrink_to_fit();
158        });
159    }
160
161    /// remove the tmp dir
162    pub fn remove_tmp_dir(&self) -> io::Result<()> {
163        time_it!("Remove tmp dir", {
164            if self.tmp_path.exists() {
165                match fs::remove_dir_all(&self.tmp_path) {
166                    Ok(()) => {}
167                    Err(e) if e.kind() == io::ErrorKind::NotFound => {}
168                    Err(e) => return Err(e),
169                }
170
171                if let Some(parent) = self.tmp_path.parent() {
172                    let is_cache_temp = parent
173                        .file_name()
174                        .and_then(|n| n.to_str())
175                        .map(|n| n == ".cache_temp")
176                        .unwrap_or(false);
177                    if is_cache_temp {
178                        match fs::remove_dir(parent) {
179                            Ok(()) => {}
180                            Err(e)
181                                if matches!(
182                                    e.kind(),
183                                    io::ErrorKind::DirectoryNotEmpty | io::ErrorKind::NotFound
184                                ) => {}
185                            Err(e) => return Err(e),
186                        }
187                    }
188                }
189            }
190
191            Ok(())
192        })
193    }
194}
195
196impl _Cache for Caches {
197    /// @param size: the size of the memory lru cache. **None means no limit**
198    /// @param tmp_path: the path to store the cache object in the tmp file
199    fn new(mem_size: Option<usize>, tmp_path: PathBuf, thread_num: usize) -> Self
200    where
201        Self: Sized,
202    {
203        // `None` means no limit, so no need to create the tmp dir
204        if mem_size.is_some() {
205            fs::create_dir_all(&tmp_path).unwrap();
206        }
207
208        Caches {
209            map_offset: DashMap::new(),
210            hash_set: DashSet::new(),
211            lru_cache: Mutex::new(LruCache::new(mem_size.unwrap_or(usize::MAX))),
212            mem_size,
213            tmp_path,
214            path_prefixes: [const { Once::new() }; 256],
215            pool: Arc::new(ThreadPool::new(thread_num)),
216            complete_signal: Arc::new(AtomicBool::new(false)),
217        }
218    }
219
220    fn get_hash(&self, offset: usize) -> Option<ObjectHash> {
221        self.map_offset.get(&offset).map(|x| *x)
222    }
223
224    fn insert(&self, offset: usize, hash: ObjectHash, obj: CacheObject) -> Arc<CacheObject> {
225        let obj_arc = Arc::new(obj);
226        {
227            // ? whether insert to cache directly or only write to tmp file
228            let mut map = self.lru_cache.lock().unwrap();
229            let mut a_obj = ArcWrapper::new(
230                obj_arc.clone(),
231                self.complete_signal.clone(),
232                Some(self.pool.clone()),
233            );
234            if self.mem_size.is_some() {
235                a_obj.set_store_path(self.generate_temp_path(&self.tmp_path, hash));
236            }
237            let _ = map.insert(hash, a_obj);
238        }
239        //order maters as for reading in 'get_by_offset()'
240        self.hash_set.insert(hash);
241        self.map_offset.insert(offset, hash);
242
243        obj_arc
244    }
245
246    /// get object by offset, from memory or tmp file
247    fn get_by_offset(&self, offset: usize) -> Option<Arc<CacheObject>> {
248        match self.map_offset.get(&offset) {
249            Some(x) => self.get_by_hash(*x),
250            None => None,
251        }
252    }
253
254    /// get object by hash, from memory or tmp file
255    fn get_by_hash(&self, hash: ObjectHash) -> Option<Arc<CacheObject>> {
256        // check if the hash is in the cache( lru or tmp file)
257        if self.hash_set.contains(&hash) {
258            match self.try_get(hash) {
259                Some(x) => Some(x),
260                None => {
261                    if self.mem_size.is_none() {
262                        panic!("should not be here when mem_size is not set")
263                    }
264                    self.get_fallback(hash).ok()
265                }
266            }
267        } else {
268            None
269        }
270    }
271
272    fn total_inserted(&self) -> usize {
273        self.hash_set.len()
274    }
275    fn memory_used(&self) -> usize {
276        self.lru_cache.lock().unwrap().current_size() + self.memory_used_index()
277    }
278    fn clear(&self) {
279        self.shutdown();
280
281        assert_eq!(self.pool.queued_count(), 0);
282        assert_eq!(self.pool.active_count(), 0);
283        assert_eq!(self.lru_cache.lock().unwrap().len(), 0);
284    }
285}
286
287#[cfg(test)]
288mod test {
289    use std::{env, sync::Arc, thread};
290
291    use super::*;
292    use crate::{
293        hash::{HashKind, ObjectHash, set_hash_kind_for_test},
294        internal::{object::types::ObjectType, pack::cache_object::CacheObjectInfo},
295    };
296
297    /// Helper to build a base CacheObject with given size and hash.
298    fn make_obj(size: usize, hash: ObjectHash) -> CacheObject {
299        CacheObject {
300            info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
301            data_decompressed: vec![0; size],
302            mem_recorder: None,
303            offset: 0,
304            crc32: 0,
305            is_delta_in_pack: false,
306        }
307    }
308
309    /// test single-threaded cache behavior with different hash kinds and capacities
310    #[test]
311    fn test_cache_single_thread() {
312        for (kind, cap, size_ab, size_c, tmp_dir) in [
313            (
314                HashKind::Sha1,
315                2048usize,
316                800usize,
317                1700usize,
318                "tests/.cache_tmp",
319            ),
320            (
321                HashKind::Sha256,
322                4096usize,
323                1500usize,
324                3000usize,
325                "tests/.cache_tmp_sha256",
326            ),
327        ] {
328            let _guard = set_hash_kind_for_test(kind);
329            let source = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
330            let tmp_path = source.clone().join(tmp_dir);
331            if tmp_path.exists() {
332                fs::remove_dir_all(&tmp_path).unwrap();
333            }
334
335            let cache = Caches::new(Some(cap), tmp_path, 1);
336            let a_hash = ObjectHash::new(String::from("a").as_bytes());
337            let b_hash = ObjectHash::new(String::from("b").as_bytes());
338            let c_hash = ObjectHash::new(String::from("c").as_bytes());
339
340            let a = make_obj(size_ab, a_hash);
341            let b = make_obj(size_ab, b_hash);
342            let c = make_obj(size_c, c_hash);
343
344            // insert a
345            cache.insert(a.offset, a_hash, a.clone());
346            assert!(cache.hash_set.contains(&a_hash));
347            assert!(cache.try_get(a_hash).is_some());
348
349            // insert b, a should still be in cache
350            cache.insert(b.offset, b_hash, b.clone());
351            assert!(cache.hash_set.contains(&b_hash));
352            assert!(cache.try_get(b_hash).is_some());
353            assert!(cache.try_get(a_hash).is_some());
354
355            // insert c which will evict both a and b
356            cache.insert(c.offset, c_hash, c.clone());
357            assert!(cache.try_get(a_hash).is_none());
358            assert!(cache.try_get(b_hash).is_none());
359            assert!(cache.try_get(c_hash).is_some());
360            assert!(cache.get_by_hash(c_hash).is_some());
361        }
362    }
363
364    /// consider the multi-threaded scenario where different threads use different hash kinds
365    #[test]
366    fn test_cache_multi_thread_mixed_hash_kinds() {
367        let base = PathBuf::from(env::current_dir().unwrap().parent().unwrap());
368        let tmp_path = base.join("tests/.cache_tmp_mixed");
369        if tmp_path.exists() {
370            fs::remove_dir_all(&tmp_path).unwrap();
371        }
372
373        let cache = Arc::new(Caches::new(Some(4096), tmp_path, 2));
374
375        let cache_sha1 = Arc::clone(&cache);
376        let handle_sha1 = thread::spawn(move || {
377            let _g = set_hash_kind_for_test(HashKind::Sha1);
378            let hash = ObjectHash::new(b"sha1-entry");
379            let obj = CacheObject {
380                info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
381                data_decompressed: vec![0; 800],
382                mem_recorder: None,
383                offset: 1,
384                crc32: 0,
385                is_delta_in_pack: false,
386            };
387            cache_sha1.insert(obj.offset, hash, obj.clone());
388            assert!(cache_sha1.hash_set.contains(&hash));
389            assert!(cache_sha1.try_get(hash).is_some());
390        });
391
392        let cache_sha256 = Arc::clone(&cache);
393        let handle_sha256 = thread::spawn(move || {
394            let _g = set_hash_kind_for_test(HashKind::Sha256);
395            let hash = ObjectHash::new(b"sha256-entry");
396            let obj = CacheObject {
397                info: CacheObjectInfo::BaseObject(ObjectType::Blob, hash),
398                data_decompressed: vec![0; 1500],
399                mem_recorder: None,
400                offset: 2,
401                crc32: 0,
402                is_delta_in_pack: false,
403            };
404            cache_sha256.insert(obj.offset, hash, obj.clone());
405            assert!(cache_sha256.hash_set.contains(&hash));
406            assert!(cache_sha256.try_get(hash).is_some());
407        });
408
409        handle_sha1.join().unwrap();
410        handle_sha256.join().unwrap();
411
412        assert_eq!(cache.total_inserted(), 2);
413    }
414
415    #[test]
416    fn test_remove_tmp_dir_does_not_panic_when_cleanup_fails() {
417        let dir = tempfile::tempdir().unwrap();
418        let tmp_path = dir.path().join("not-a-directory");
419        fs::write(&tmp_path, b"cache marker").unwrap();
420        let cache = Caches::new(None, tmp_path, 1);
421
422        let result =
423            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| cache.remove_tmp_dir()));
424
425        assert!(result.is_ok(), "cache cleanup should not panic");
426        assert!(
427            result.unwrap().is_err(),
428            "cleanup failure should be returned"
429        );
430    }
431}
git_internal/internal/pack/cache.rs

git_internal/internal/pack/
cache.rs
