1use std::{
62 fs::{File, OpenOptions},
63 io::{self, BufWriter, Cursor, Read, Seek, SeekFrom, Write},
64 path::PathBuf,
65};
66
67use heddle_format::compression::CompressionConfig;
68
69use super::{ObjectType, PackObjectId, PackStats, pack_container_spec, write_container_header};
70
71#[cfg(feature = "zstd")]
77const CSIZE_PLACEHOLDER_LEN: usize = 10;
78use crate::{
79 object::ContentHash,
80 store::{Result, StoreError},
81};
82
83const BUCKETS_PER_VARIANT: usize = 256;
89const TOTAL_BUCKETS: usize = BUCKETS_PER_VARIANT * 2;
91const MAX_OPEN_BUCKET_WRITERS: usize = 32;
95
96const HASH_VARIANT: usize = 0;
100const CHANGEID_VARIANT: usize = 1;
101
102pub trait SyncData {
108 fn sync_data_for_durability(&mut self) -> io::Result<()>;
109}
110
111impl SyncData for File {
112 fn sync_data_for_durability(&mut self) -> io::Result<()> {
113 self.sync_all()
114 }
115}
116
117impl SyncData for Cursor<Vec<u8>> {
118 fn sync_data_for_durability(&mut self) -> io::Result<()> {
119 Ok(())
120 }
121}
122
123pub struct StreamingPackBuilder<W: Write + Read + Seek> {
126 pack_writer: Option<BufWriter<W>>,
132 header_offset: u64,
136 pack_position: u64,
139 object_count: u64,
140 declared_object_count: Option<u64>,
141 total_uncompressed: u64,
142 total_compressed: u64,
143 #[cfg_attr(not(feature = "zstd"), allow(dead_code))]
148 compression: CompressionConfig,
149 bucket_dir: PathBuf,
152 bucket_writers: Vec<Option<BucketWriter>>,
156 open_bucket_writers: usize,
157 bucket_access_tick: u64,
158 bucket_paths: Vec<PathBuf>,
159 index_path: PathBuf,
163 finalized: bool,
166}
167
168struct BucketWriter {
169 writer: BufWriter<File>,
170 last_used: u64,
171}
172
173#[cfg(feature = "zstd")]
174struct CountingWriter<'a, W: Write> {
175 inner: &'a mut W,
176 written: u64,
177}
178
179#[cfg(feature = "zstd")]
180impl<'a, W: Write> CountingWriter<'a, W> {
181 fn new(inner: &'a mut W) -> Self {
182 Self { inner, written: 0 }
183 }
184}
185
186#[cfg(feature = "zstd")]
187impl<W: Write> Write for CountingWriter<'_, W> {
188 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
189 let written = self.inner.write(buf)?;
190 self.written = self.written.saturating_add(written as u64);
191 Ok(written)
192 }
193
194 fn flush(&mut self) -> std::io::Result<()> {
195 self.inner.flush()
196 }
197}
198
199impl<W: Write + Read + Seek + SyncData> StreamingPackBuilder<W> {
200 pub fn new(
216 pack_writer: W,
217 index_path: PathBuf,
218 compression: CompressionConfig,
219 bucket_dir: PathBuf,
220 ) -> Result<Self> {
221 Self::new_inner(pack_writer, index_path, compression, bucket_dir, None)
222 }
223
224 pub fn new_with_object_count(
231 pack_writer: W,
232 index_path: PathBuf,
233 compression: CompressionConfig,
234 bucket_dir: PathBuf,
235 object_count: u64,
236 ) -> Result<Self> {
237 Self::new_inner(
238 pack_writer,
239 index_path,
240 compression,
241 bucket_dir,
242 Some(object_count),
243 )
244 }
245
246 fn new_inner(
247 mut pack_writer: W,
248 index_path: PathBuf,
249 compression: CompressionConfig,
250 bucket_dir: PathBuf,
251 declared_object_count: Option<u64>,
252 ) -> Result<Self> {
253 crate::fs_atomic::create_dir_all_durable(&bucket_dir).map_err(StoreError::from)?;
254 let header_offset = pack_writer.stream_position().map_err(StoreError::from)?;
255
256 let mut header_bytes = Vec::with_capacity(16);
261 write_container_header(
262 &mut header_bytes,
263 pack_container_spec(),
264 declared_object_count.unwrap_or(0),
265 );
266 pack_writer
267 .write_all(&header_bytes)
268 .map_err(StoreError::from)?;
269
270 let bucket_paths: Vec<PathBuf> = (0..TOTAL_BUCKETS)
271 .map(|i| {
272 let variant = if i < BUCKETS_PER_VARIANT { 'h' } else { 'c' };
273 let prefix = i % BUCKETS_PER_VARIANT;
274 bucket_dir.join(format!("bucket-{variant}-{prefix:02x}"))
275 })
276 .collect();
277 for path in &bucket_paths {
278 let _ = std::fs::remove_file(path);
279 }
280
281 Ok(Self {
282 pack_writer: Some(BufWriter::new(pack_writer)),
283 header_offset,
284 pack_position: header_offset + header_bytes.len() as u64,
285 object_count: 0,
286 declared_object_count,
287 total_uncompressed: 0,
288 total_compressed: 0,
289 compression,
290 bucket_dir,
291 bucket_writers: (0..TOTAL_BUCKETS).map(|_| None).collect(),
292 open_bucket_writers: 0,
293 bucket_access_tick: 0,
294 bucket_paths,
295 index_path,
296 finalized: false,
297 })
298 }
299
300 pub fn flush_pack(&mut self) -> Result<()> {
305 if let Some(writer) = self.pack_writer.as_mut() {
306 writer.flush().map_err(StoreError::from)?;
307 }
308 Ok(())
309 }
310
311 pub fn add(&mut self, hash: ContentHash, obj_type: ObjectType, data: Vec<u8>) -> Result<()> {
313 self.add_id(PackObjectId::Hash(hash), obj_type, data)
314 }
315
316 pub fn add_id(&mut self, id: PackObjectId, obj_type: ObjectType, data: Vec<u8>) -> Result<()> {
339 let pw = self
344 .pack_writer
345 .as_mut()
346 .expect("add_id called after finalize");
347 let entry_start = self.pack_position;
348 let offset = entry_start
349 .checked_sub(self.header_offset)
350 .expect("header_offset should never be past current position");
351
352 self.total_uncompressed += data.len() as u64;
353
354 let mut header_buf = Vec::with_capacity(40);
357 id.encode_tagged(&mut header_buf);
358 super::varint::encode_type_and_size(obj_type, data.len() as u64, &mut header_buf);
359 pw.write_all(&header_buf).map_err(StoreError::from)?;
360 self.pack_position = self
361 .pack_position
362 .checked_add(header_buf.len() as u64)
363 .ok_or_else(|| {
364 StoreError::InvalidObject("streaming pack position overflow".to_string())
365 })?;
366 #[cfg(feature = "zstd")]
370 let csize_pos = entry_start + header_buf.len() as u64;
371
372 let want_compress: bool;
385 #[cfg(feature = "zstd")]
386 {
387 want_compress = self.compression.enabled && data.len() >= self.compression.min_size;
388 }
389 #[cfg(not(feature = "zstd"))]
390 {
391 want_compress = false;
392 }
393 if !want_compress {
394 let mut csize_buf = Vec::with_capacity(10);
397 super::varint::encode_varint(data.len() as u64, &mut csize_buf);
398 pw.write_all(&csize_buf).map_err(StoreError::from)?;
399 self.pack_position = self
400 .pack_position
401 .checked_add(csize_buf.len() as u64)
402 .ok_or_else(|| {
403 StoreError::InvalidObject("streaming pack position overflow".to_string())
404 })?;
405 pw.write_all(&data).map_err(StoreError::from)?;
406 self.pack_position = self
407 .pack_position
408 .checked_add(data.len() as u64)
409 .ok_or_else(|| {
410 StoreError::InvalidObject("streaming pack position overflow".to_string())
411 })?;
412 self.total_compressed += data.len() as u64;
413 } else {
414 #[cfg(feature = "zstd")]
415 {
416 pw.write_all(&[0u8; CSIZE_PLACEHOLDER_LEN])
419 .map_err(StoreError::from)?;
420 self.pack_position = self
421 .pack_position
422 .checked_add(CSIZE_PLACEHOLDER_LEN as u64)
423 .ok_or_else(|| {
424 StoreError::InvalidObject("streaming pack position overflow".to_string())
425 })?;
426 let body_start = self.pack_position;
427 let compressed_size;
428 {
429 let mut counting = CountingWriter::new(&mut *pw);
430 let mut enc =
431 zstd::stream::write::Encoder::new(&mut counting, self.compression.level)
432 .map_err(StoreError::from)?;
433 enc.set_pledged_src_size(Some(data.len() as u64))
438 .map_err(StoreError::from)?;
439 enc.write_all(&data).map_err(StoreError::from)?;
440 enc.finish().map_err(StoreError::from)?;
441 compressed_size = counting.written;
442 }
443 self.pack_position =
444 self.pack_position
445 .checked_add(compressed_size)
446 .ok_or_else(|| {
447 StoreError::InvalidObject(
448 "streaming pack position overflow".to_string(),
449 )
450 })?;
451 let body_end = body_start.checked_add(compressed_size).ok_or_else(|| {
452 StoreError::InvalidObject("streaming pack position overflow".to_string())
453 })?;
454 self.total_compressed += compressed_size;
455
456 let mut csize_bytes = [0u8; CSIZE_PLACEHOLDER_LEN];
461 encode_varint_padded_to_10(compressed_size, &mut csize_bytes);
462 pw.flush().map_err(StoreError::from)?;
463 let inner = pw.get_mut();
464 inner
465 .seek(SeekFrom::Start(csize_pos))
466 .map_err(StoreError::from)?;
467 inner.write_all(&csize_bytes).map_err(StoreError::from)?;
468 inner
469 .seek(SeekFrom::Start(body_end))
470 .map_err(StoreError::from)?;
471 }
472 #[cfg(not(feature = "zstd"))]
473 {
474 unreachable!("compression branch reached without `zstd` feature");
477 }
478 }
479
480 let bucket_idx = bucket_index_for(&id);
484 let bucket = self.get_or_open_bucket(bucket_idx)?;
485 let mut idx_entry = Vec::with_capacity(33 + 8);
486 id.encode_tagged(&mut idx_entry);
487 idx_entry.extend_from_slice(&offset.to_be_bytes());
488 bucket.write_all(&idx_entry).map_err(StoreError::from)?;
489
490 self.object_count += 1;
491 Ok(())
492 }
493
494 fn get_or_open_bucket(&mut self, idx: usize) -> Result<&mut BufWriter<File>> {
495 self.bucket_access_tick = self.bucket_access_tick.wrapping_add(1);
496 let last_used = self.bucket_access_tick;
497 if self.bucket_writers[idx].is_none() {
498 if self.open_bucket_writers >= MAX_OPEN_BUCKET_WRITERS {
499 self.evict_lru_bucket()?;
500 }
501 let path = &self.bucket_paths[idx];
502 let f = OpenOptions::new()
503 .create(true)
504 .append(true)
505 .open(path)
506 .map_err(StoreError::from)?;
507 self.bucket_writers[idx] = Some(BucketWriter {
508 writer: BufWriter::new(f),
509 last_used,
510 });
511 self.open_bucket_writers += 1;
512 } else if let Some(bucket) = self.bucket_writers[idx].as_mut() {
513 bucket.last_used = last_used;
514 }
515 Ok(&mut self.bucket_writers[idx]
516 .as_mut()
517 .expect("just inserted above")
518 .writer)
519 }
520
521 fn evict_lru_bucket(&mut self) -> Result<()> {
522 let Some((idx, _)) = self
523 .bucket_writers
524 .iter()
525 .enumerate()
526 .filter_map(|(idx, bucket)| bucket.as_ref().map(|bucket| (idx, bucket.last_used)))
527 .min_by_key(|(_, last_used)| *last_used)
528 else {
529 return Ok(());
530 };
531
532 if let Some(mut bucket) = self.bucket_writers[idx].take() {
533 bucket.writer.flush().map_err(StoreError::from)?;
534 self.open_bucket_writers -= 1;
535 }
536 Ok(())
537 }
538
539 pub fn finalize(mut self) -> Result<(W, PackStats)> {
550 for bucket in self.bucket_writers.iter_mut().flatten() {
553 bucket.writer.flush().map_err(StoreError::from)?;
554 }
555 for slot in self.bucket_writers.iter_mut() {
558 *slot = None;
559 }
560 self.open_bucket_writers = 0;
561
562 let bw = self
566 .pack_writer
567 .take()
568 .expect("finalize called twice — pack_writer already consumed");
569 let mut writer = bw
570 .into_inner()
571 .map_err(|e| StoreError::from(std::io::Error::other(e.to_string())))?;
572 if let Some(expected) = self.declared_object_count {
573 if expected != self.object_count {
574 return Err(StoreError::InvalidObject(format!(
575 "streaming pack declared {expected} object(s) but added {}",
576 self.object_count
577 )));
578 }
579 } else {
580 writer
581 .seek(SeekFrom::Start(self.header_offset))
582 .map_err(StoreError::from)?;
583 let mut header_bytes = Vec::with_capacity(16);
584 write_container_header(&mut header_bytes, pack_container_spec(), self.object_count);
585 writer.write_all(&header_bytes).map_err(StoreError::from)?;
586 }
587
588 writer
593 .seek(SeekFrom::Start(self.header_offset))
594 .map_err(StoreError::from)?;
595 let mut hasher = blake3::Hasher::new();
596 let mut buf = vec![0u8; 64 * 1024];
597 loop {
598 let n = writer.read(&mut buf).map_err(StoreError::from)?;
599 if n == 0 {
600 break;
601 }
602 hasher.update(&buf[..n]);
603 }
604 let checksum = hasher.finalize();
605
606 writer.seek(SeekFrom::End(0)).map_err(StoreError::from)?;
608 writer
609 .write_all(checksum.as_bytes())
610 .map_err(StoreError::from)?;
611 writer.flush().map_err(StoreError::from)?;
612 writer
614 .sync_data_for_durability()
615 .map_err(StoreError::from)?;
616
617 let idx_file = File::create(&self.index_path).map_err(StoreError::from)?;
630 let mut idx_writer = BufWriter::new(idx_file);
631 write_index_header(&mut idx_writer, self.object_count)?;
632 let mut entries_written: u64 = 0;
633 for path in self.bucket_paths.iter() {
634 if !path.exists() {
635 continue;
636 }
637 let bucket_bytes = std::fs::read(path).map_err(StoreError::from)?;
638 let mut entries = decode_bucket_file(&bucket_bytes)?;
639 entries.sort_by_key(|(id, _)| *id);
644 for (id, offset) in entries {
645 write_index_entry(&mut idx_writer, id, offset)?;
646 entries_written += 1;
647 }
648 }
649 idx_writer.flush().map_err(StoreError::from)?;
650 let idx_file = idx_writer
652 .into_inner()
653 .map_err(|e| StoreError::from(std::io::Error::other(e.to_string())))?;
654 idx_file.sync_all().map_err(StoreError::from)?;
655 if let Some(parent) = self.index_path.parent() {
656 crate::fs_atomic::sync_directory(parent).map_err(StoreError::from)?;
657 }
658 debug_assert_eq!(
659 entries_written, self.object_count,
660 "streaming index entry count drifted from add() count"
661 );
662
663 for path in self.bucket_paths.iter() {
668 let _ = std::fs::remove_file(path);
669 }
670 let _ = std::fs::remove_dir(&self.bucket_dir);
671 self.finalized = true;
672
673 let stats = PackStats {
674 object_count: self.object_count,
675 total_uncompressed: self.total_uncompressed,
676 total_compressed: self.total_compressed,
677 delta_count: 0,
678 compression_ratio: if self.total_uncompressed == 0 {
679 0.0
680 } else {
681 self.total_compressed as f64 / self.total_uncompressed as f64
682 },
683 };
684
685 Ok((writer, stats))
686 }
687}
688
689fn write_index_header<W: Write>(out: &mut W, count: u64) -> Result<()> {
694 super::pack_index::index_header().write_to(out, count)
695}
696
697fn write_index_entry<W: Write>(out: &mut W, id: PackObjectId, offset: u64) -> Result<()> {
701 let mut buf = Vec::with_capacity(33 + 8);
702 id.encode_tagged(&mut buf);
703 buf.extend_from_slice(&offset.to_be_bytes());
704 out.write_all(&buf).map_err(StoreError::from)
705}
706
707#[cfg(feature = "zstd")]
720fn encode_varint_padded_to_10(value: u64, out: &mut [u8; 10]) {
721 let mut v = value;
722 for slot in out.iter_mut().take(9) {
723 *slot = 0x80 | ((v & 0x7F) as u8);
724 v >>= 7;
725 }
726 out[9] = (v & 0x7F) as u8;
727}
728
729impl<W: Write + Read + Seek> Drop for StreamingPackBuilder<W> {
730 fn drop(&mut self) {
731 if self.finalized {
732 return;
733 }
734 for path in self.bucket_paths.iter() {
737 let _ = std::fs::remove_file(path);
738 }
739 let _ = std::fs::remove_dir(&self.bucket_dir);
740 }
741}
742
743fn bucket_index_for(id: &PackObjectId) -> usize {
747 match id {
748 PackObjectId::Hash(h) => HASH_VARIANT * BUCKETS_PER_VARIANT + h.as_bytes()[0] as usize,
749 PackObjectId::ChangeId(c) => {
750 CHANGEID_VARIANT * BUCKETS_PER_VARIANT + c.as_bytes()[0] as usize
751 }
752 }
753}
754
755fn decode_bucket_file(bytes: &[u8]) -> Result<Vec<(PackObjectId, u64)>> {
760 let mut out = Vec::new();
761 let mut pos = 0;
762 while pos < bytes.len() {
763 let (id, id_len) = PackObjectId::decode_tagged(&bytes[pos..])?;
764 pos += id_len;
765 if pos + 8 > bytes.len() {
766 return Err(StoreError::InvalidObject(
767 "streaming bucket entry truncated at offset".to_string(),
768 ));
769 }
770 let offset = u64::from_be_bytes(bytes[pos..pos + 8].try_into().map_err(|_| {
771 StoreError::InvalidObject("streaming bucket bad offset slice".to_string())
772 })?);
773 pos += 8;
774 out.push((id, offset));
775 }
776 Ok(out)
777}
778
779#[cfg(test)]
782mod tests {
783 use std::io::Cursor;
784
785 use super::*;
786 use crate::{
787 object::ChangeId,
788 store::pack::{PackReader, PackStats},
789 };
790
791 fn deterministic_hash(seed: u8) -> ContentHash {
792 let mut bytes = [0u8; 32];
796 bytes[0] = seed;
797 for (i, b) in bytes.iter_mut().enumerate().skip(1) {
798 *b = seed.wrapping_mul(31).wrapping_add(i as u8);
799 }
800 ContentHash::from_bytes(bytes)
801 }
802
803 fn deterministic_change_id(seed: u8) -> ChangeId {
804 let mut bytes = [0u8; 16];
805 bytes[0] = seed;
806 for (i, b) in bytes.iter_mut().enumerate().skip(1) {
807 *b = seed.wrapping_add(i as u8 * 7);
808 }
809 ChangeId::from_bytes(bytes)
810 }
811
812 fn fresh_builder(
817 tmp: &tempfile::TempDir,
818 ) -> (StreamingPackBuilder<Cursor<Vec<u8>>>, PathBuf, PathBuf) {
819 let bucket_dir = tmp.path().join("buckets");
820 let index_path = tmp.path().join("test.idx");
821 let cursor = Cursor::new(Vec::<u8>::new());
822 let b = StreamingPackBuilder::new(
823 cursor,
824 index_path.clone(),
825 CompressionConfig::default(),
826 bucket_dir.clone(),
827 )
828 .unwrap();
829 (b, bucket_dir, index_path)
830 }
831
832 fn finalize_cursor(
837 b: StreamingPackBuilder<Cursor<Vec<u8>>>,
838 index_path: &std::path::Path,
839 ) -> (Vec<u8>, Vec<u8>, PackStats) {
840 let (cursor, stats) = b.finalize().unwrap();
841 let index_bytes = std::fs::read(index_path).unwrap();
842 (cursor.into_inner(), index_bytes, stats)
843 }
844
845 #[test]
846 fn empty_pack_finalizes_to_valid_zero_count_pack() {
847 let tmp = tempfile::TempDir::new().unwrap();
848 let (b, bucket_dir, idx_path) = fresh_builder(&tmp);
849 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
850
851 assert_eq!(stats.object_count, 0);
852 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
854 assert!(reader.list_ids().is_empty());
855 assert!(
857 !bucket_dir.exists(),
858 "bucket dir should be cleaned on successful finalize"
859 );
860 }
861
862 #[test]
863 fn single_blob_with_hash_id_round_trips() {
864 let tmp = tempfile::TempDir::new().unwrap();
865 let (mut b, _, idx_path) = fresh_builder(&tmp);
866 let hash = deterministic_hash(0x42);
867 let payload = b"hello, streaming pack".to_vec();
868 b.add(hash, ObjectType::Blob, payload.clone()).unwrap();
869 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
870
871 assert_eq!(stats.object_count, 1);
872 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
873 let id = PackObjectId::Hash(hash);
874 assert!(reader.has_object(&id));
875 let (got_type, got_data) = reader.get_object(&id).unwrap().unwrap();
876 assert_eq!(got_type, ObjectType::Blob);
877 assert_eq!(got_data, payload);
878 }
879
880 #[test]
881 fn single_state_with_change_id_round_trips() {
882 let tmp = tempfile::TempDir::new().unwrap();
883 let (mut b, _, idx_path) = fresh_builder(&tmp);
884 let cid = deterministic_change_id(0xa5);
885 let payload = b"serialized-state-bytes".to_vec();
886 b.add_id(
887 PackObjectId::ChangeId(cid),
888 ObjectType::State,
889 payload.clone(),
890 )
891 .unwrap();
892 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
893
894 assert_eq!(stats.object_count, 1);
895 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
896 let id = PackObjectId::ChangeId(cid);
897 let (ty, data) = reader.get_object(&id).unwrap().unwrap();
898 assert_eq!(ty, ObjectType::State);
899 assert_eq!(data, payload);
900 }
901
902 #[test]
903 fn mixed_hash_and_changeid_ids_all_retrievable() {
904 let tmp = tempfile::TempDir::new().unwrap();
905 let (mut b, _, idx_path) = fresh_builder(&tmp);
906 let blob_hash = deterministic_hash(0x10);
907 let tree_hash = deterministic_hash(0x20);
908 let state_cid = deterministic_change_id(0x80);
909
910 b.add(blob_hash, ObjectType::Blob, b"blob-bytes".to_vec())
911 .unwrap();
912 b.add(tree_hash, ObjectType::Tree, b"serialized-tree".to_vec())
913 .unwrap();
914 b.add_id(
915 PackObjectId::ChangeId(state_cid),
916 ObjectType::State,
917 b"serialized-state".to_vec(),
918 )
919 .unwrap();
920
921 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
922 assert_eq!(stats.object_count, 3);
923 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
924 assert_eq!(
925 reader
926 .get_object(&PackObjectId::Hash(blob_hash))
927 .unwrap()
928 .unwrap()
929 .1,
930 b"blob-bytes".to_vec()
931 );
932 assert_eq!(
933 reader
934 .get_object(&PackObjectId::Hash(tree_hash))
935 .unwrap()
936 .unwrap()
937 .1,
938 b"serialized-tree".to_vec()
939 );
940 assert_eq!(
941 reader
942 .get_object(&PackObjectId::ChangeId(state_cid))
943 .unwrap()
944 .unwrap()
945 .1,
946 b"serialized-state".to_vec()
947 );
948 }
949
950 #[test]
951 fn ten_thousand_objects_round_trip_correctly() {
952 let tmp = tempfile::TempDir::new().unwrap();
956 let (mut b, _, idx_path) = fresh_builder(&tmp);
957 let mut hashes = Vec::with_capacity(10_000);
958 for i in 0..10_000u32 {
959 let h = blake3::hash(&i.to_le_bytes());
962 let hash = ContentHash::from_bytes(*h.as_bytes());
963 hashes.push(hash);
964 b.add(hash, ObjectType::Blob, format!("payload-{i}").into_bytes())
965 .unwrap();
966 }
967 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
968 assert_eq!(stats.object_count, 10_000);
969
970 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
971 assert_eq!(reader.list_ids().len(), 10_000);
972 for i in [0, 1, 99, 1234, 5_000, 9_999] {
974 let id = PackObjectId::Hash(hashes[i]);
975 let (_ty, data) = reader.get_object(&id).unwrap().unwrap();
976 assert_eq!(data, format!("payload-{i}").into_bytes());
977 }
978 }
979
980 #[test]
981 fn bucket_writers_are_lru_capped_below_fd_limit() {
982 let tmp = tempfile::TempDir::new().unwrap();
983 let (mut b, _bucket_dir, idx_path) = fresh_builder(&tmp);
984 let mut ids = Vec::new();
985
986 for i in 0..BUCKETS_PER_VARIANT {
987 let hash = deterministic_hash(i as u8);
988 ids.push(PackObjectId::Hash(hash));
989 b.add(hash, ObjectType::Blob, format!("hash-{i}").into_bytes())
990 .unwrap();
991 assert!(
992 b.open_bucket_writers <= MAX_OPEN_BUCKET_WRITERS,
993 "open bucket writers should stay capped"
994 );
995 }
996
997 for i in 0..BUCKETS_PER_VARIANT {
998 let cid = deterministic_change_id(i as u8);
999 ids.push(PackObjectId::ChangeId(cid));
1000 b.add_id(
1001 PackObjectId::ChangeId(cid),
1002 ObjectType::State,
1003 format!("state-{i}").into_bytes(),
1004 )
1005 .unwrap();
1006 assert!(
1007 b.open_bucket_writers <= MAX_OPEN_BUCKET_WRITERS,
1008 "open bucket writers should stay capped"
1009 );
1010 }
1011
1012 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
1013 assert_eq!(stats.object_count, TOTAL_BUCKETS as u64);
1014 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1015 for id in ids {
1016 assert!(reader.has_object(&id), "missing id {id:?}");
1017 }
1018 }
1019
1020 #[test]
1021 fn index_id_sort_order_matches_packbuilder_output() {
1022 use crate::store::pack::PackBuilder;
1030 let payloads: Vec<(PackObjectId, ObjectType, Vec<u8>)> = (0..200u32)
1031 .map(|i| {
1032 let h = blake3::hash(&i.to_le_bytes());
1033 (
1034 PackObjectId::Hash(ContentHash::from_bytes(*h.as_bytes())),
1035 if i % 3 == 0 {
1036 ObjectType::Tree
1037 } else {
1038 ObjectType::Blob
1039 },
1040 format!("body-{i}").into_bytes(),
1041 )
1042 })
1043 .collect();
1044
1045 let compression = CompressionConfig {
1048 max_delta_size: 0,
1049 ..CompressionConfig::default()
1050 };
1051 let mut classic = PackBuilder::new(compression);
1052 for (id, ty, data) in payloads.iter() {
1053 classic.add_id(*id, *ty, data.clone());
1054 }
1055 let (classic_pack, classic_index, _) = classic.build().unwrap();
1056 let classic_reader = PackReader::from_bytes(classic_pack, classic_index).unwrap();
1057
1058 let tmp = tempfile::TempDir::new().unwrap();
1059 let bucket_dir = tmp.path().join("buckets");
1060 let idx_path = tmp.path().join("test.idx");
1061 let cursor = Cursor::new(Vec::<u8>::new());
1062 let mut streaming =
1063 StreamingPackBuilder::new(cursor, idx_path.clone(), compression, bucket_dir).unwrap();
1064 for (id, ty, data) in payloads.iter() {
1065 streaming.add_id(*id, *ty, data.clone()).unwrap();
1066 }
1067 let (streaming_pack, streaming_index, _) = finalize_cursor(streaming, &idx_path);
1068 let streaming_reader = PackReader::from_bytes(streaming_pack, streaming_index).unwrap();
1069
1070 assert_eq!(
1073 streaming_reader.list_ids(),
1074 classic_reader.list_ids(),
1075 "streaming and classic indices should report the same id sequence"
1076 );
1077 for (id, _ty, want) in payloads.iter().take(10).chain(payloads.iter().skip(190)) {
1080 let (_, got) = streaming_reader.get_object(id).unwrap().unwrap();
1081 assert_eq!(&got, want);
1082 let (_, classic_got) = classic_reader.get_object(id).unwrap().unwrap();
1083 assert_eq!(got, classic_got);
1084 }
1085 }
1086
1087 #[test]
1088 fn corrupted_pack_fails_checksum_verification() {
1089 let tmp = tempfile::TempDir::new().unwrap();
1090 let (mut b, _, idx_path) = fresh_builder(&tmp);
1091 b.add(
1092 deterministic_hash(0x01),
1093 ObjectType::Blob,
1094 b"some bytes".to_vec(),
1095 )
1096 .unwrap();
1097 let (mut pack_data, index_data, _) = finalize_cursor(b, &idx_path);
1098 let body_byte = 18; pack_data[body_byte] ^= 0xff;
1101 let result = PackReader::from_bytes(pack_data, index_data);
1102 assert!(
1103 result.is_err(),
1104 "PackReader should reject pack with mutated body"
1105 );
1106 }
1107
1108 #[test]
1109 fn pack_count_in_header_matches_index_entry_count() {
1110 let tmp = tempfile::TempDir::new().unwrap();
1111 let (mut b, _, idx_path) = fresh_builder(&tmp);
1112 for i in 0..7u8 {
1113 b.add(
1114 deterministic_hash(i),
1115 ObjectType::Blob,
1116 format!("p{i}").into_bytes(),
1117 )
1118 .unwrap();
1119 }
1120 let (pack_data, index_data, _) = finalize_cursor(b, &idx_path);
1121 let count = u64::from_be_bytes(pack_data[8..16].try_into().unwrap());
1123 assert_eq!(count, 7);
1124 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1125 assert_eq!(reader.list_ids().len(), 7);
1126 }
1127
1128 #[test]
1129 fn declared_pack_count_is_written_before_finalize() {
1130 let tmp = tempfile::TempDir::new().unwrap();
1131 let bucket_dir = tmp.path().join("buckets");
1132 let idx_path = tmp.path().join("test.idx");
1133 let cursor = Cursor::new(Vec::<u8>::new());
1134 let mut b = StreamingPackBuilder::new_with_object_count(
1135 cursor,
1136 idx_path.clone(),
1137 CompressionConfig::default(),
1138 bucket_dir,
1139 2,
1140 )
1141 .unwrap();
1142
1143 b.flush_pack().unwrap();
1144 let initial = b.pack_writer.as_ref().unwrap().get_ref().get_ref().clone();
1145 assert_eq!(u64::from_be_bytes(initial[8..16].try_into().unwrap()), 2);
1146
1147 let hash = deterministic_hash(0x40);
1148 b.add(hash, ObjectType::Blob, b"known-count-entry".to_vec())
1149 .unwrap();
1150 b.flush_pack().unwrap();
1151 let after_add = b.pack_writer.as_ref().unwrap().get_ref().get_ref().clone();
1152 assert_eq!(u64::from_be_bytes(after_add[8..16].try_into().unwrap()), 2);
1153
1154 let second_hash = deterministic_hash(0x41);
1155 b.add(second_hash, ObjectType::Blob, b"second-entry".to_vec())
1156 .unwrap();
1157 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
1158
1159 assert_eq!(stats.object_count, 2);
1160 assert_eq!(u64::from_be_bytes(pack_data[8..16].try_into().unwrap()), 2);
1161 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1162 assert!(reader.has_object(&PackObjectId::Hash(hash)));
1163 assert!(reader.has_object(&PackObjectId::Hash(second_hash)));
1164 }
1165
1166 #[test]
1167 fn declared_pack_count_mismatch_fails_finalize() {
1168 let tmp = tempfile::TempDir::new().unwrap();
1169 let bucket_dir = tmp.path().join("buckets");
1170 let idx_path = tmp.path().join("test.idx");
1171 let cursor = Cursor::new(Vec::<u8>::new());
1172 let mut b = StreamingPackBuilder::new_with_object_count(
1173 cursor,
1174 idx_path,
1175 CompressionConfig::default(),
1176 bucket_dir,
1177 2,
1178 )
1179 .unwrap();
1180
1181 b.add(
1182 deterministic_hash(0x50),
1183 ObjectType::Blob,
1184 b"only-entry".to_vec(),
1185 )
1186 .unwrap();
1187 let error = b.finalize().unwrap_err();
1188
1189 assert!(
1190 error
1191 .to_string()
1192 .contains("streaming pack declared 2 object(s) but added 1")
1193 );
1194 }
1195
1196 #[test]
1197 fn bucket_files_are_cleaned_on_successful_finalize() {
1198 let tmp = tempfile::TempDir::new().unwrap();
1199 let bucket_dir = tmp.path().join("buckets");
1200 let idx_path = tmp.path().join("test.idx");
1201 let cursor = Cursor::new(Vec::<u8>::new());
1202 let mut b = StreamingPackBuilder::new(
1203 cursor,
1204 idx_path.clone(),
1205 CompressionConfig::default(),
1206 bucket_dir.clone(),
1207 )
1208 .unwrap();
1209 for i in 0..50u8 {
1210 b.add(deterministic_hash(i), ObjectType::Blob, vec![i; 32])
1211 .unwrap();
1212 }
1213 assert!(bucket_dir.exists());
1215 let bucket_count = std::fs::read_dir(&bucket_dir).unwrap().count();
1216 assert!(bucket_count > 0, "bucket dir should hold some files");
1217 let _ = finalize_cursor(b, &idx_path);
1218 assert!(
1219 !bucket_dir.exists(),
1220 "bucket dir should be removed on finalize"
1221 );
1222 }
1223
1224 #[test]
1225 fn bucket_files_are_cleaned_on_drop_without_finalize() {
1226 let tmp = tempfile::TempDir::new().unwrap();
1227 let bucket_dir = tmp.path().join("buckets");
1228 let idx_path = tmp.path().join("test.idx");
1229 {
1230 let cursor = Cursor::new(Vec::<u8>::new());
1231 let mut b = StreamingPackBuilder::new(
1232 cursor,
1233 idx_path.clone(),
1234 CompressionConfig::default(),
1235 bucket_dir.clone(),
1236 )
1237 .unwrap();
1238 for i in 0..10u8 {
1239 b.add(deterministic_hash(i), ObjectType::Blob, vec![0; 32])
1240 .unwrap();
1241 }
1242 assert!(bucket_dir.exists());
1243 }
1245 assert!(
1246 !idx_path.exists(),
1247 "no index file should have been created without finalize"
1248 );
1249 assert!(
1250 !bucket_dir.exists(),
1251 "bucket dir should be removed on Drop when finalize never ran"
1252 );
1253 }
1254
1255 #[test]
1256 fn large_blob_streams_to_disk_without_double_buffering() {
1257 let tmp = tempfile::TempDir::new().unwrap();
1262 let bucket_dir = tmp.path().join("buckets");
1263 let pack_path = tmp.path().join("pack.dat");
1264 let idx_path = tmp.path().join("pack.idx");
1265 let file = std::fs::OpenOptions::new()
1266 .read(true)
1267 .write(true)
1268 .create(true)
1269 .truncate(true)
1270 .open(&pack_path)
1271 .unwrap();
1272 let mut b = StreamingPackBuilder::new(
1273 file,
1274 idx_path.clone(),
1275 CompressionConfig::default(),
1276 bucket_dir,
1277 )
1278 .unwrap();
1279 let payload: Vec<u8> = (0..4 * 1024 * 1024u32).map(|i| (i & 0xff) as u8).collect();
1280 let hash = deterministic_hash(0xff);
1281 b.add(hash, ObjectType::Blob, payload.clone()).unwrap();
1282 let (_, stats) = b.finalize().unwrap();
1283 let index_data = std::fs::read(&idx_path).unwrap();
1284 assert_eq!(stats.object_count, 1);
1285 let pack_bytes = std::fs::read(&pack_path).unwrap();
1286 let reader = PackReader::from_bytes(pack_bytes, index_data).unwrap();
1289 let (_ty, got) = reader
1290 .get_object(&PackObjectId::Hash(hash))
1291 .unwrap()
1292 .unwrap();
1293 assert_eq!(got, payload);
1294 }
1295
1296 #[test]
1297 fn bucket_distribution_for_random_hashes_is_roughly_uniform() {
1298 let tmp = tempfile::TempDir::new().unwrap();
1305 let bucket_dir = tmp.path().join("buckets");
1306 let idx_path = tmp.path().join("test.idx");
1307 let cursor = Cursor::new(Vec::<u8>::new());
1308 let mut b = StreamingPackBuilder::new(
1309 cursor,
1310 idx_path.clone(),
1311 CompressionConfig::default(),
1312 bucket_dir.clone(),
1313 )
1314 .unwrap();
1315 for i in 0..1024u32 {
1316 let h = blake3::hash(&i.to_le_bytes());
1317 let hash = ContentHash::from_bytes(*h.as_bytes());
1318 b.add(hash, ObjectType::Blob, b"x".to_vec()).unwrap();
1319 }
1320 b.pack_writer.as_mut().unwrap().flush().unwrap();
1322 let mut max_entries = 0usize;
1323 let entry_size = 33 + 8; for path in b.bucket_paths.iter() {
1325 if path.exists() {
1326 let size = std::fs::metadata(path).unwrap().len() as usize;
1327 let entries = size / entry_size;
1328 if entries > max_entries {
1329 max_entries = entries;
1330 }
1331 }
1332 }
1333 assert!(
1336 max_entries <= 16,
1337 "max bucket has {max_entries} entries; uniform expected ~4"
1338 );
1339 let _ = finalize_cursor(b, &idx_path);
1340 }
1341
1342 #[test]
1343 fn finalize_returns_correct_stats() {
1344 let tmp = tempfile::TempDir::new().unwrap();
1345 let (mut b, _, idx_path) = fresh_builder(&tmp);
1346 let payload = vec![0xabu8; 1024];
1347 for i in 0..5u8 {
1348 b.add(deterministic_hash(i), ObjectType::Blob, payload.clone())
1349 .unwrap();
1350 }
1351 let (_, _, stats) = finalize_cursor(b, &idx_path);
1352 assert_eq!(stats.object_count, 5);
1353 assert_eq!(stats.total_uncompressed, 5 * 1024);
1354 assert!(stats.total_compressed > 0);
1355 assert!(stats.compression_ratio > 0.0);
1356 assert_eq!(stats.delta_count, 0, "streaming builder never deltas");
1357 }
1358
1359 #[cfg(feature = "zstd")]
1360 #[test]
1361 fn streaming_compression_roundtrips_through_zstd_frame() {
1362 let tmp = tempfile::TempDir::new().unwrap();
1370 let (mut b, _, idx_path) = fresh_builder(&tmp);
1371 let payload = vec![0u8; 64 * 1024];
1374 let hash = deterministic_hash(0x77);
1375 b.add(hash, ObjectType::Blob, payload.clone()).unwrap();
1376 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
1377 assert!(
1378 stats.total_compressed < stats.total_uncompressed,
1379 "expected compression ratio < 1.0, got {}/{}",
1380 stats.total_compressed,
1381 stats.total_uncompressed
1382 );
1383 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1384 let (_ty, got) = reader
1385 .get_object(&PackObjectId::Hash(hash))
1386 .unwrap()
1387 .unwrap();
1388 assert_eq!(got, payload);
1389 }
1390
1391 #[cfg(feature = "zstd")]
1392 #[test]
1393 fn padded_varint_decodes_to_original_value_for_canonical_decoder() {
1394 let cases: &[u64] = &[0, 1, 127, 128, 4096, 1_000_000, 1_000_000_000_000, u64::MAX];
1400 for &value in cases {
1401 let mut buf = [0u8; 10];
1402 super::encode_varint_padded_to_10(value, &mut buf);
1403 let (decoded, consumed) = super::super::varint::decode_varint(&buf)
1404 .expect("padded varint should always decode");
1405 assert_eq!(decoded, value, "varint roundtrip failed for {value}");
1406 assert_eq!(
1407 consumed, 10,
1408 "padded encoding should consume all 10 bytes for {value}"
1409 );
1410 }
1411 }
1412
1413 #[cfg(feature = "zstd")]
1414 #[test]
1415 fn streaming_path_does_not_buffer_compressed_payload_in_memory() {
1416 let tmp = tempfile::TempDir::new().unwrap();
1429 let bucket_dir = tmp.path().join("buckets");
1430 let pack_path = tmp.path().join("pack.dat");
1431 let idx_path = tmp.path().join("pack.idx");
1432 let file = std::fs::OpenOptions::new()
1433 .read(true)
1434 .write(true)
1435 .create(true)
1436 .truncate(true)
1437 .open(&pack_path)
1438 .unwrap();
1439 let mut b = StreamingPackBuilder::new(
1440 file,
1441 idx_path.clone(),
1442 CompressionConfig::default(),
1443 bucket_dir,
1444 )
1445 .unwrap();
1446 let payload = vec![0xa5u8; 8 * 1024 * 1024];
1447 let hash = deterministic_hash(0x66);
1448 b.add(hash, ObjectType::Blob, payload.clone()).unwrap();
1449 let mid_size = std::fs::metadata(&pack_path).unwrap().len();
1453 assert!(
1454 mid_size > 16 + 40,
1455 "pack file should hold real entry data after add; size={mid_size}"
1456 );
1457 let (_, _) = b.finalize().unwrap();
1458 let pack_bytes = std::fs::read(&pack_path).unwrap();
1459 let index_bytes = std::fs::read(&idx_path).unwrap();
1460 let reader = PackReader::from_bytes(pack_bytes, index_bytes).unwrap();
1461 let (_ty, got) = reader
1462 .get_object(&PackObjectId::Hash(hash))
1463 .unwrap()
1464 .unwrap();
1465 assert_eq!(got, payload);
1466 }
1467
1468 #[test]
1469 fn list_ids_returns_all_added_ids_sorted() {
1470 let tmp = tempfile::TempDir::new().unwrap();
1471 let (mut b, _, idx_path) = fresh_builder(&tmp);
1472 let mut added: Vec<PackObjectId> = Vec::new();
1473 for seed in [0x05u8, 0xa0, 0x12, 0x9f, 0x33] {
1475 let id = PackObjectId::Hash(deterministic_hash(seed));
1476 b.add_id(id, ObjectType::Blob, vec![seed; 4]).unwrap();
1477 added.push(id);
1478 }
1479 for seed in [0x80u8, 0x10, 0xff] {
1480 let id = PackObjectId::ChangeId(deterministic_change_id(seed));
1481 b.add_id(id, ObjectType::State, vec![seed; 4]).unwrap();
1482 added.push(id);
1483 }
1484 let (pack_data, index_data, _) = finalize_cursor(b, &idx_path);
1485 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1486 let mut got = reader.list_ids();
1487 let mut sorted = got.clone();
1490 sorted.sort();
1491 assert_eq!(got, sorted, "list_ids must come back sorted");
1492 added.sort();
1494 got.sort();
1495 assert_eq!(got, added);
1496 }
1497}