1use std::{
5 borrow::Cow,
6 fs, io,
7 io::Write,
8 ops::Deref,
9 path::{Path, PathBuf},
10 sync::{
11 Arc,
12 atomic::{AtomicBool, AtomicUsize, Ordering},
13 },
14};
15
16use lru_mem::{HeapSize, MemSize};
17use rkyv::{
18 Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize,
19 api::high::{HighSerializer, HighValidator},
20 bytecheck::CheckBytes,
21 de::Pool,
22 rancor::{Error as RkyvError, Strategy},
23 ser::allocator::ArenaHandle,
24 util::AlignedVec,
25};
26use tempfile::NamedTempFile;
27use threadpool::ThreadPool;
28
29use crate::{
30 hash::ObjectHash,
31 internal::{
32 metadata::{EntryMeta, MetaAttached},
33 object::types::ObjectType,
34 pack::{entry::Entry, utils},
35 },
36};
37
38pub trait FileLoadStore: Sized {
42 fn f_load(path: &Path) -> Result<Self, io::Error>;
43 fn f_save(&self, path: &Path) -> Result<(), io::Error>;
44}
45
46fn write_bytes_atomically(path: &Path, data: &[u8]) -> Result<(), io::Error> {
47 if path.exists() {
48 return Ok(());
49 }
50 let dir = path
51 .parent()
52 .filter(|parent| !parent.as_os_str().is_empty())
53 .unwrap_or_else(|| Path::new("."));
54
55 let mut temp_file = NamedTempFile::new_in(dir)?;
56 temp_file.write_all(data)?;
57
58 match temp_file.persist_noclobber(path) {
59 Ok(_persisted) => Ok(()),
60 Err(err) if err.error.kind() == io::ErrorKind::AlreadyExists && path.exists() => Ok(()),
61 Err(err) => Err(err.error),
62 }
63}
64
65impl<T> FileLoadStore for T
66where
67 T: Archive,
68 T: for<'a> RkyvSerialize<HighSerializer<AlignedVec, ArenaHandle<'a>, RkyvError>>,
69 T::Archived: for<'a> CheckBytes<HighValidator<'a, RkyvError>>
70 + RkyvDeserialize<T, Strategy<Pool, RkyvError>>,
71{
72 fn f_load(path: &Path) -> Result<T, io::Error> {
74 let data = fs::read(path)?;
75 let obj = rkyv::from_bytes::<T, RkyvError>(&data).map_err(io::Error::other)?;
76 Ok(obj)
77 }
78
79 fn f_save(&self, path: &Path) -> Result<(), io::Error> {
80 if path.exists() {
81 return Ok(());
82 }
83
84 let data = rkyv::to_bytes::<RkyvError>(self).map_err(io::Error::other)?;
85 write_bytes_atomically(path, &data)
86 }
87}
88
89#[derive(
91 PartialEq,
92 Eq,
93 Clone,
94 Debug,
95 serde::Serialize,
96 serde::Deserialize,
97 rkyv::Archive,
98 rkyv::Serialize,
99 rkyv::Deserialize,
100)]
101pub(crate) enum CacheObjectInfo {
102 BaseObject(ObjectType, ObjectHash),
106 OffsetDelta(usize, usize),
109 OffsetZstdelta(usize, usize),
111 HashDelta(ObjectHash, usize),
114}
115
116impl CacheObjectInfo {
117 pub(crate) fn object_type(&self) -> ObjectType {
119 match self {
120 CacheObjectInfo::BaseObject(obj_type, _) => *obj_type,
121 CacheObjectInfo::OffsetDelta(_, _) => ObjectType::OffsetDelta,
122 CacheObjectInfo::OffsetZstdelta(_, _) => ObjectType::OffsetZstdelta,
123 CacheObjectInfo::HashDelta(_, _) => ObjectType::HashDelta,
124 }
125 }
126}
127
128#[derive(Debug)]
130pub struct CacheObject {
131 pub(crate) info: CacheObjectInfo,
132 pub offset: usize,
133 pub crc32: u32,
134 pub data_decompressed: Vec<u8>,
135 pub mem_recorder: Option<Arc<AtomicUsize>>, pub is_delta_in_pack: bool,
137 pub(crate) known_hash: Option<ObjectHash>,
138}
139
140#[derive(Debug, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)]
141struct CacheObjectOnDisk {
142 info: CacheObjectInfo,
143 offset: usize,
144 crc32: u32,
145 data_decompressed: Vec<u8>,
146 is_delta_in_pack: bool,
147}
148
149#[derive(Debug, rkyv::Archive, rkyv::Serialize)]
150struct CacheObjectOnDiskRef<'a> {
151 #[rkyv(with = rkyv::with::AsOwned)]
152 info: Cow<'a, CacheObjectInfo>,
153 offset: usize,
154 crc32: u32,
155 #[rkyv(with = rkyv::with::AsOwned)]
156 data_decompressed: Cow<'a, [u8]>,
157 is_delta_in_pack: bool,
158}
159
160impl<'a> From<&'a CacheObject> for CacheObjectOnDiskRef<'a> {
161 fn from(value: &'a CacheObject) -> Self {
162 Self {
163 info: Cow::Borrowed(&value.info),
164 offset: value.offset,
165 crc32: value.crc32,
166 data_decompressed: Cow::Borrowed(value.data_decompressed.as_slice()),
167 is_delta_in_pack: value.is_delta_in_pack,
168 }
169 }
170}
171
172impl From<CacheObjectOnDisk> for CacheObject {
173 fn from(value: CacheObjectOnDisk) -> Self {
174 Self {
175 info: value.info,
176 offset: value.offset,
177 crc32: value.crc32,
178 data_decompressed: value.data_decompressed,
179 mem_recorder: None,
180 is_delta_in_pack: value.is_delta_in_pack,
181 known_hash: None,
182 }
183 }
184}
185
186impl FileLoadStore for CacheObject {
187 fn f_load(path: &Path) -> Result<Self, io::Error> {
188 let obj = CacheObjectOnDisk::f_load(path)?;
189 Ok(obj.into())
190 }
191
192 fn f_save(&self, path: &Path) -> Result<(), io::Error> {
193 if path.exists() {
194 return Ok(());
195 }
196
197 let data = rkyv::to_bytes::<RkyvError>(&CacheObjectOnDiskRef::from(self))
198 .map_err(io::Error::other)?;
199 write_bytes_atomically(path, &data)
200 }
201}
202
203impl Clone for CacheObject {
204 fn clone(&self) -> Self {
205 let obj = CacheObject {
206 info: self.info.clone(),
207 offset: self.offset,
208 crc32: self.crc32,
209 data_decompressed: self.data_decompressed.clone(),
210 mem_recorder: self.mem_recorder.clone(),
211 is_delta_in_pack: self.is_delta_in_pack,
212 known_hash: self.known_hash,
213 };
214 obj.record_mem_size();
215 obj
216 }
217}
218
219impl HeapSize for CacheObject {
223 fn heap_size(&self) -> usize {
232 match &self.info {
233 CacheObjectInfo::BaseObject(_, _) => self.data_decompressed.heap_size(),
234 CacheObjectInfo::OffsetDelta(_, delta_final_size)
235 | CacheObjectInfo::OffsetZstdelta(_, delta_final_size)
236 | CacheObjectInfo::HashDelta(_, delta_final_size) => {
237 self.data_decompressed.heap_size() + delta_final_size
248 }
249 }
250 }
251}
252
253impl Drop for CacheObject {
254 fn drop(&mut self) {
257 if let Some(mem_recorder) = &self.mem_recorder {
259 mem_recorder.fetch_sub((*self).mem_size(), Ordering::Release);
260 }
261 }
262}
263
264pub trait MemSizeRecorder: MemSize {
271 fn record_mem_size(&self);
272 fn set_mem_recorder(&mut self, mem_size: Arc<AtomicUsize>);
273 }
275
276impl MemSizeRecorder for CacheObject {
277 fn record_mem_size(&self) {
280 if let Some(mem_recorder) = &self.mem_recorder {
281 mem_recorder.fetch_add(self.mem_size(), Ordering::Release);
282 }
283 }
284
285 fn set_mem_recorder(&mut self, mem_recorder: Arc<AtomicUsize>) {
286 self.mem_recorder = Some(mem_recorder);
287 }
288
289 }
293
294impl CacheObject {
295 pub fn new_for_undeltified(
297 obj_type: ObjectType,
298 data: Vec<u8>,
299 offset: usize,
300 crc32: u32,
301 ) -> Self {
302 let hash = utils::calculate_object_hash(obj_type, &data);
303 CacheObject {
304 info: CacheObjectInfo::BaseObject(obj_type, hash),
305 offset,
306 crc32,
307 data_decompressed: data,
308 mem_recorder: None,
309 is_delta_in_pack: false,
310 known_hash: None,
311 }
312 }
313
314 pub fn object_type(&self) -> ObjectType {
316 self.info.object_type()
317 }
318
319 pub fn base_object_hash(&self) -> Option<ObjectHash> {
323 match &self.info {
324 CacheObjectInfo::BaseObject(_, hash) => Some(*hash),
325 _ => None,
326 }
327 }
328
329 pub fn offset_delta(&self) -> Option<usize> {
337 match &self.info {
338 CacheObjectInfo::OffsetDelta(base_offset, _) => self.offset.checked_sub(*base_offset),
339 _ => None,
340 }
341 }
342
343 pub fn offset_delta_base_offset(&self) -> Option<usize> {
347 match &self.info {
348 CacheObjectInfo::OffsetDelta(base_offset, _) => Some(*base_offset),
349 _ => None,
350 }
351 }
352
353 pub fn hash_delta(&self) -> Option<ObjectHash> {
357 match &self.info {
358 CacheObjectInfo::HashDelta(hash, _) => Some(*hash),
359 _ => None,
360 }
361 }
362
363 pub fn to_entry(&self) -> Entry {
365 match self.info {
366 CacheObjectInfo::BaseObject(obj_type, hash) => Entry {
367 obj_type,
368 data: self.data_decompressed.clone(),
369 hash,
370 chain_len: 0,
371 },
372 _ => {
373 unreachable!("delta object should not persist!")
374 }
375 }
376 }
377
378 pub fn to_entry_metadata(&self) -> MetaAttached<Entry, EntryMeta> {
380 match self.info {
381 CacheObjectInfo::BaseObject(obj_type, hash) => {
382 let entry = Entry {
383 obj_type,
384 data: self.data_decompressed.clone(),
385 hash,
386 chain_len: 0,
387 };
388 let meta = EntryMeta {
389 pack_offset: Some(self.offset),
391 crc32: Some(self.crc32),
392 is_delta: Some(self.is_delta_in_pack),
393 ..Default::default()
394 };
395 MetaAttached { inner: entry, meta }
396 }
397
398 _ => {
399 unreachable!("delta object should not persist!")
400 }
401 }
402 }
403
404 pub fn into_entry_metadata(mut self) -> MetaAttached<Entry, EntryMeta> {
406 let recorded_size = self.mem_size();
407 if let Some(mem_recorder) = self.mem_recorder.take() {
408 mem_recorder.fetch_sub(recorded_size, Ordering::Release);
409 }
410
411 match self.info {
412 CacheObjectInfo::BaseObject(obj_type, hash) => {
413 let data = std::mem::take(&mut self.data_decompressed);
414 let entry = Entry {
415 obj_type,
416 data,
417 hash,
418 chain_len: 0,
419 };
420 let meta = EntryMeta {
421 pack_offset: Some(self.offset),
422 crc32: Some(self.crc32),
423 is_delta: Some(self.is_delta_in_pack),
424 ..Default::default()
425 };
426 MetaAttached { inner: entry, meta }
427 }
428
429 _ => {
430 unreachable!("delta object should not persist!")
431 }
432 }
433 }
434}
435
436pub trait ArcWrapperBounds: HeapSize + FileLoadStore + Send + Sync + 'static {}
438impl<T: HeapSize + FileLoadStore + Send + Sync + 'static> ArcWrapperBounds for T {}
441
442pub struct ArcWrapper<T: ArcWrapperBounds> {
446 pub data: Arc<T>,
447 complete_signal: Arc<AtomicBool>,
448 pool: Option<Arc<ThreadPool>>,
449 pub store_path: Option<PathBuf>, }
451impl<T: ArcWrapperBounds> ArcWrapper<T> {
452 pub fn new(data: Arc<T>, share_flag: Arc<AtomicBool>, pool: Option<Arc<ThreadPool>>) -> Self {
454 ArcWrapper {
455 data,
456 complete_signal: share_flag,
457 pool,
458 store_path: None,
459 }
460 }
461 pub fn set_store_path(&mut self, path: PathBuf) {
463 self.store_path = Some(path);
464 }
465}
466
467impl<T: ArcWrapperBounds> HeapSize for ArcWrapper<T> {
468 fn heap_size(&self) -> usize {
469 self.data.heap_size()
470 }
471}
472
473impl<T: ArcWrapperBounds> Clone for ArcWrapper<T> {
474 fn clone(&self) -> Self {
476 ArcWrapper {
477 data: self.data.clone(),
478 complete_signal: self.complete_signal.clone(),
479 pool: self.pool.clone(),
480 store_path: None,
481 }
482 }
483}
484
485impl<T: ArcWrapperBounds> Deref for ArcWrapper<T> {
486 type Target = Arc<T>;
487 fn deref(&self) -> &Self::Target {
488 &self.data
489 }
490}
491impl<T: ArcWrapperBounds> Drop for ArcWrapper<T> {
492 fn drop(&mut self) {
495 if !self.complete_signal.load(Ordering::Acquire)
496 && let Some(path) = &self.store_path
497 {
498 match &self.pool {
499 Some(pool) => {
500 let data_copy = self.data.clone();
501 let path_copy = path.clone();
502 let complete_signal = self.complete_signal.clone();
503 while pool.queued_count() > 2000 {
506 std::thread::yield_now();
507 }
508 pool.execute(move || {
509 if !complete_signal.load(Ordering::Acquire) {
510 let res = data_copy.f_save(&path_copy);
511 if let Err(e) = res {
512 println!("[f_save] {path_copy:?} error: {e:?}");
513 }
514 }
515 });
516 }
517 None => {
518 let res = self.data.f_save(path);
519 if let Err(e) = res {
520 println!("[f_save] {path:?} error: {e:?}");
521 }
522 }
523 }
524 }
525 }
526}
527#[cfg(test)]
528mod test {
529 use std::{fs, sync::Mutex};
530
531 use lru_mem::LruCache;
532 use tempfile::tempdir;
533
534 use super::*;
535 use crate::hash::{HashKind, set_hash_kind_for_test};
536
537 fn make_obj(size: usize) -> CacheObject {
539 CacheObject {
540 info: CacheObjectInfo::BaseObject(ObjectType::Blob, ObjectHash::default()),
541 offset: 0,
542 crc32: 0,
543 data_decompressed: vec![0; size],
544 mem_recorder: None,
545 is_delta_in_pack: false,
546 known_hash: None,
547 }
548 }
549
550 #[test]
551 fn offset_delta_accessor_returns_encoded_distance() {
552 let base_offset = 1024;
553 let delta_distance = 128;
554 let object_offset = base_offset + delta_distance;
555
556 let obj = CacheObject {
557 info: CacheObjectInfo::OffsetDelta(base_offset, 32),
558 offset: object_offset,
559 crc32: 0,
560 data_decompressed: Vec::new(),
561 mem_recorder: None,
562 is_delta_in_pack: true,
563 known_hash: None,
564 };
565
566 assert_eq!(obj.offset_delta(), Some(delta_distance));
567 assert_eq!(obj.offset_delta_base_offset(), Some(base_offset));
568 }
569
570 #[test]
572 fn test_heap_size_record() {
573 for (kind, size) in [(HashKind::Sha1, 1024usize), (HashKind::Sha256, 2048usize)] {
574 let _guard = set_hash_kind_for_test(kind);
575 let mut obj = make_obj(size);
576 let mem = Arc::new(AtomicUsize::default());
577 assert_eq!(mem.load(Ordering::Relaxed), 0);
578 obj.set_mem_recorder(mem.clone());
579 obj.record_mem_size();
580 assert_eq!(mem.load(Ordering::Relaxed), obj.mem_size());
581 drop(obj);
582 assert_eq!(mem.load(Ordering::Relaxed), 0);
583 }
584 }
585
586 #[test]
588 fn test_cache_object_with_same_size() {
589 for (kind, size) in [(HashKind::Sha1, 1024usize), (HashKind::Sha256, 2048usize)] {
590 let _guard = set_hash_kind_for_test(kind);
591 let a = make_obj(size);
592 assert_eq!(a.heap_size(), size);
593
594 let b = ArcWrapper::new(Arc::new(a.clone()), Arc::new(AtomicBool::new(false)), None);
595 assert_eq!(b.heap_size(), size);
596 }
597 }
598
599 #[test]
601 fn test_cache_object_with_lru() {
602 for (kind, cap, size_a, size_b) in [
603 (
604 HashKind::Sha1,
605 2048usize,
606 1024usize,
607 (1024.0 * 1.5) as usize,
608 ),
609 (
610 HashKind::Sha256,
611 4096usize,
612 2048usize,
613 (2048.0 * 1.5) as usize,
614 ),
615 ] {
616 let _guard = set_hash_kind_for_test(kind);
617 let mut cache = LruCache::new(cap);
618
619 let hash_a = ObjectHash::default();
620 let hash_b = ObjectHash::new(b"b"); let a = make_obj(size_a);
622 let b = make_obj(size_b);
623
624 {
625 let r = cache.insert(
626 hash_a.to_string(),
627 ArcWrapper::new(Arc::new(a.clone()), Arc::new(AtomicBool::new(true)), None),
628 );
629 assert!(r.is_ok());
630 }
631
632 {
633 let r = cache.try_insert(
634 hash_b.to_string(),
635 ArcWrapper::new(Arc::new(b.clone()), Arc::new(AtomicBool::new(true)), None),
636 );
637 assert!(r.is_err());
638 if let Err(lru_mem::TryInsertError::WouldEjectLru { .. }) = r {
639 } else {
641 panic!("Expected WouldEjectLru error");
642 }
643
644 let r = cache.insert(
645 hash_b.to_string(),
646 ArcWrapper::new(Arc::new(b.clone()), Arc::new(AtomicBool::new(true)), None),
647 );
648 assert!(r.is_ok());
649 }
650
651 {
652 let r = cache.get(&hash_a.to_string());
653 assert!(r.is_none());
654 }
655 }
656 }
657
658 #[derive(
660 serde::Serialize, serde::Deserialize, rkyv::Archive, rkyv::Serialize, rkyv::Deserialize,
661 )]
662 struct Test {
663 a: usize,
664 }
665 impl Drop for Test {
666 fn drop(&mut self) {
667 println!("drop Test");
668 }
669 }
670 impl HeapSize for Test {
671 fn heap_size(&self) -> usize {
672 self.a
673 }
674 }
675 #[test]
676 fn test_lru_drop() {
677 println!("insert a");
678 let cache = LruCache::new(2048);
679 let cache = Arc::new(Mutex::new(cache));
680 {
681 let mut c = cache.as_ref().lock().unwrap();
682 let _ = c.insert(
683 "a",
684 ArcWrapper::new(
685 Arc::new(Test { a: 1024 }),
686 Arc::new(AtomicBool::new(true)),
687 None,
688 ),
689 );
690 }
691 println!("insert b, a should be ejected");
692 {
693 let mut c = cache.as_ref().lock().unwrap();
694 let _ = c.insert(
695 "b",
696 ArcWrapper::new(
697 Arc::new(Test { a: 1200 }),
698 Arc::new(AtomicBool::new(true)),
699 None,
700 ),
701 );
702 }
703 let b = {
704 let mut c = cache.as_ref().lock().unwrap();
705 c.get("b").cloned()
706 };
707 println!("insert c, b should not be ejected");
708 {
709 let mut c = cache.as_ref().lock().unwrap();
710 let _ = c.insert(
711 "c",
712 ArcWrapper::new(
713 Arc::new(Test { a: 1200 }),
714 Arc::new(AtomicBool::new(true)),
715 None,
716 ),
717 );
718 }
719 println!("user b: {}", b.as_ref().unwrap().a);
720 println!("test over, enject all");
721 }
722
723 #[test]
724 fn test_cache_object_serialize() {
725 for (kind, size) in [(HashKind::Sha1, 1024usize), (HashKind::Sha256, 2048usize)] {
726 let _guard = set_hash_kind_for_test(kind);
727 let a = make_obj(size);
728 let s = rkyv::to_bytes::<RkyvError>(&CacheObjectOnDiskRef::from(&a)).unwrap();
729 let b: CacheObject = rkyv::from_bytes::<CacheObjectOnDisk, RkyvError>(&s)
730 .unwrap()
731 .into();
732 assert_eq!(a.info, b.info);
733 assert_eq!(a.data_decompressed, b.data_decompressed);
734 assert_eq!(a.offset, b.offset);
735 assert!(b.mem_recorder.is_none());
736 }
737 }
738
739 #[test]
740 fn test_write_bytes_atomically_creates_file_when_missing() {
741 let dir = tempdir().unwrap();
742 let path = dir.path().join("object");
743
744 write_bytes_atomically(&path, b"fresh").unwrap();
745
746 assert_eq!(fs::read(&path).unwrap(), b"fresh");
747 }
748
749 #[test]
750 fn test_write_bytes_atomically_returns_when_target_exists() {
751 let dir = tempdir().unwrap();
752 let path = dir.path().join("object");
753
754 fs::write(&path, b"existing").unwrap();
755 write_bytes_atomically(&path, b"new-data").unwrap();
756
757 assert_eq!(fs::read(&path).unwrap(), b"existing");
758 }
759
760 #[test]
761 fn test_cache_object_file_roundtrip() {
762 let dir = tempdir().unwrap();
763 let path = dir.path().join("object");
764 let a = make_obj(1024);
765
766 a.f_save(&path).unwrap();
767 let b = CacheObject::f_load(&path).unwrap();
768
769 assert_eq!(a.info, b.info);
770 assert_eq!(a.data_decompressed, b.data_decompressed);
771 assert_eq!(a.offset, b.offset);
772 assert!(b.mem_recorder.is_none());
773 }
774
775 #[test]
776 fn test_arc_wrapper_drop_store() {
777 let mut path = PathBuf::from(".cache_temp/test_arc_wrapper_drop_store");
778 fs::create_dir_all(&path).unwrap();
779 path.push("test_obj");
780 let mut a = ArcWrapper::new(Arc::new(1024), Arc::new(AtomicBool::new(false)), None);
781 a.set_store_path(path.clone());
782 drop(a);
783
784 assert!(path.exists());
785 path.pop();
786 fs::remove_dir_all(&path).unwrap();
787 let _ = fs::remove_dir(".cache_temp");
789 }
790
791 #[test]
792 fn test_arc_wrapper_with_lru() {
794 let mut cache = LruCache::new(1500);
795 let path = PathBuf::from(".cache_temp/test_arc_wrapper_with_lru");
796 let _ = fs::remove_dir_all(&path);
797 fs::create_dir_all(&path).unwrap();
798 let shared_flag = Arc::new(AtomicBool::new(false));
799
800 let a_path = path.join("a");
802 {
803 let mut a = ArcWrapper::new(Arc::new(Test { a: 1024 }), shared_flag.clone(), None);
804 a.set_store_path(a_path.clone());
805 let b = ArcWrapper::new(Arc::new(1024), shared_flag.clone(), None);
806 assert!(b.store_path.is_none());
807
808 println!("insert a with heap size: {:?}", a.heap_size());
809 let rt = cache.insert("a", a);
810 if let Err(e) = rt {
811 panic!("{}", format!("insert a failed: {:?}", e.to_string()));
812 }
813 println!("after insert a, cache used = {}", cache.current_size());
814 }
815 assert!(!a_path.exists());
816
817 let b_path = path.join("b");
818 {
820 let mut b = ArcWrapper::new(Arc::new(Test { a: 996 }), shared_flag.clone(), None);
821 b.set_store_path(b_path.clone());
822 let rt = cache.insert("b", b);
823 if let Err(e) = rt {
824 panic!("{}", format!("insert a failed: {:?}", e.to_string()));
825 }
826 println!("after insert b, cache used = {}", cache.current_size());
827 }
828 assert!(a_path.exists());
829 assert!(!b_path.exists());
830 shared_flag.store(true, Ordering::Release);
831 fs::remove_dir_all(path).unwrap();
832 let _ = fs::remove_dir(".cache_temp");
834 }
836}