From 7216204ee736f4dd7ca63bffc8af8114b43e888a Mon Sep 17 00:00:00 2001 From: Iris Shi <0.0@owo.li> Date: Fri, 13 Dec 2024 14:19:30 +0800 Subject: [PATCH 1/2] fix(memory): fix a crash when decoding without memory limit and related test --- mercury/src/internal/pack/cache.rs | 4 +++- mercury/src/internal/pack/decode.rs | 18 ++++++++++++++++-- mercury/src/internal/pack/mod.rs | 2 +- 3 files changed, 20 insertions(+), 4 deletions(-) diff --git a/mercury/src/internal/pack/cache.rs b/mercury/src/internal/pack/cache.rs index b53873d35..81376499d 100644 --- a/mercury/src/internal/pack/cache.rs +++ b/mercury/src/internal/pack/cache.rs @@ -153,7 +153,9 @@ impl _Cache for Caches { self.complete_signal.clone(), Some(self.pool.clone()), ); - a_obj.set_store_path(Caches::generate_temp_path(&self.tmp_path, hash)); + if self.mem_size.is_some() { + a_obj.set_store_path(Caches::generate_temp_path(&self.tmp_path, hash)); + } let _ = map.insert(hash.to_string(), a_obj); } //order maters as for reading in 'get_by_offset()' diff --git a/mercury/src/internal/pack/decode.rs b/mercury/src/internal/pack/decode.rs index 2766e5960..13a599e12 100644 --- a/mercury/src/internal/pack/decode.rs +++ b/mercury/src/internal/pack/decode.rs @@ -75,7 +75,7 @@ impl Pack { pool: Arc::new(ThreadPool::new(thread_num)), waitlist: Arc::new(Waitlist::new()), caches: Arc::new(Caches::new(cache_mem_size, temp_path, thread_num)), - mem_limit: mem_limit.unwrap_or(usize::MAX), + mem_limit: mem_limit, cache_objs_mem: Arc::new(AtomicUsize::default()), clean_tmp, } @@ -363,7 +363,8 @@ impl Pack { } // 3 parts: Waitlist + TheadPool + Caches // hardcode the limit of the tasks of threads_pool queue, to limit memory - while self.memory_used() > self.mem_limit || self.pool.queued_count() > 2000 { + while self.pool.queued_count() > 2000 + || self.mem_limit.map(|limit| self.memory_used() > limit).unwrap_or(false) { thread::yield_now(); } let r: Result = self.decode_pack_object(&mut reader, &mut offset); @@ -709,6 +710,19 @@ mod tests { p.decode(&mut buffered,|_,_|{}).unwrap(); } + #[test] + fn test_pack_decode_no_mem_limit() { + let mut source = PathBuf::from(env::current_dir().unwrap().parent().unwrap()); + source.push("tests/data/packs/pack-1d0e6c14760c956c173ede71cb28f33d921e232f.pack"); + + let tmp = PathBuf::from("/tmp/.cache_temp"); + + let f = fs::File::open(source).unwrap(); + let mut buffered = BufReader::new(f); + let mut p = Pack::new(None, None, Some(tmp), true); + p.decode(&mut buffered, |_,_|{}).unwrap(); + } + #[test] fn test_pack_decode_with_large_file_with_delta_without_ref() { init_logger(); diff --git a/mercury/src/internal/pack/mod.rs b/mercury/src/internal/pack/mod.rs index 3f942ca2a..360204ce2 100644 --- a/mercury/src/internal/pack/mod.rs +++ b/mercury/src/internal/pack/mod.rs @@ -29,7 +29,7 @@ pub struct Pack { pub pool: Arc, pub waitlist: Arc, pub caches: Arc, - pub mem_limit: usize, + pub mem_limit: Option, pub cache_objs_mem: Arc, // the memory size of CacheObjects in this Pack pub clean_tmp: bool, } From 1f1f2055eda2c3c56043c9ecf6bbfe696607ae41 Mon Sep 17 00:00:00 2001 From: Iris Shi <0.0@owo.li> Date: Fri, 13 Dec 2024 14:25:24 +0800 Subject: [PATCH 2/2] apply lint --- mercury/src/internal/pack/decode.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mercury/src/internal/pack/decode.rs b/mercury/src/internal/pack/decode.rs index 13a599e12..7755a03e2 100644 --- a/mercury/src/internal/pack/decode.rs +++ b/mercury/src/internal/pack/decode.rs @@ -75,7 +75,7 @@ impl Pack { pool: Arc::new(ThreadPool::new(thread_num)), waitlist: Arc::new(Waitlist::new()), caches: Arc::new(Caches::new(cache_mem_size, temp_path, thread_num)), - mem_limit: mem_limit, + mem_limit, cache_objs_mem: Arc::new(AtomicUsize::default()), clean_tmp, }