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 durable: bool,
167}
168
169struct BucketWriter {
170 writer: BufWriter<File>,
171 last_used: u64,
172}
173
174#[cfg(feature = "zstd")]
175struct CountingWriter<'a, W: Write> {
176 inner: &'a mut W,
177 written: u64,
178}
179
180#[cfg(feature = "zstd")]
181impl<'a, W: Write> CountingWriter<'a, W> {
182 fn new(inner: &'a mut W) -> Self {
183 Self { inner, written: 0 }
184 }
185}
186
187#[cfg(feature = "zstd")]
188impl<W: Write> Write for CountingWriter<'_, W> {
189 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
190 let written = self.inner.write(buf)?;
191 self.written = self.written.saturating_add(written as u64);
192 Ok(written)
193 }
194
195 fn flush(&mut self) -> std::io::Result<()> {
196 self.inner.flush()
197 }
198}
199
200impl<W: Write + Read + Seek + SyncData> StreamingPackBuilder<W> {
201 pub fn new(
217 pack_writer: W,
218 index_path: PathBuf,
219 compression: CompressionConfig,
220 bucket_dir: PathBuf,
221 ) -> Result<Self> {
222 Self::new_inner(pack_writer, index_path, compression, bucket_dir, None, true)
223 }
224
225 pub fn new_with_object_count(
232 pack_writer: W,
233 index_path: PathBuf,
234 compression: CompressionConfig,
235 bucket_dir: PathBuf,
236 object_count: u64,
237 ) -> Result<Self> {
238 Self::new_inner(
239 pack_writer,
240 index_path,
241 compression,
242 bucket_dir,
243 Some(object_count),
244 true,
245 )
246 }
247
248 pub fn new_with_object_count_ephemeral(
255 pack_writer: W,
256 index_path: PathBuf,
257 compression: CompressionConfig,
258 bucket_dir: PathBuf,
259 object_count: u64,
260 ) -> Result<Self> {
261 Self::new_inner(
262 pack_writer,
263 index_path,
264 compression,
265 bucket_dir,
266 Some(object_count),
267 false,
268 )
269 }
270
271 fn new_inner(
272 mut pack_writer: W,
273 index_path: PathBuf,
274 compression: CompressionConfig,
275 bucket_dir: PathBuf,
276 declared_object_count: Option<u64>,
277 durable: bool,
278 ) -> Result<Self> {
279 if durable {
280 heddle_fs_prims::fs_atomic::create_dir_all_durable(&bucket_dir)
281 .map_err(StoreError::from)?;
282 } else {
283 std::fs::create_dir_all(&bucket_dir).map_err(StoreError::from)?;
284 }
285 let header_offset = pack_writer.stream_position().map_err(StoreError::from)?;
286
287 let mut header_bytes = Vec::with_capacity(16);
292 write_container_header(
293 &mut header_bytes,
294 pack_container_spec(),
295 declared_object_count.unwrap_or(0),
296 );
297 pack_writer
298 .write_all(&header_bytes)
299 .map_err(StoreError::from)?;
300
301 let bucket_paths: Vec<PathBuf> = (0..TOTAL_BUCKETS)
302 .map(|i| {
303 let variant = if i < BUCKETS_PER_VARIANT { 'h' } else { 'c' };
304 let prefix = i % BUCKETS_PER_VARIANT;
305 bucket_dir.join(format!("bucket-{variant}-{prefix:02x}"))
306 })
307 .collect();
308 for path in &bucket_paths {
309 let _ = std::fs::remove_file(path);
310 }
311
312 Ok(Self {
313 pack_writer: Some(BufWriter::new(pack_writer)),
314 header_offset,
315 pack_position: header_offset + header_bytes.len() as u64,
316 object_count: 0,
317 declared_object_count,
318 total_uncompressed: 0,
319 total_compressed: 0,
320 compression,
321 bucket_dir,
322 bucket_writers: (0..TOTAL_BUCKETS).map(|_| None).collect(),
323 open_bucket_writers: 0,
324 bucket_access_tick: 0,
325 bucket_paths,
326 index_path,
327 finalized: false,
328 durable,
329 })
330 }
331
332 pub fn flush_pack(&mut self) -> Result<()> {
337 if let Some(writer) = self.pack_writer.as_mut() {
338 writer.flush().map_err(StoreError::from)?;
339 }
340 Ok(())
341 }
342
343 pub fn add(&mut self, hash: ContentHash, obj_type: ObjectType, data: Vec<u8>) -> Result<()> {
345 self.add_id(PackObjectId::Hash(hash), obj_type, data)
346 }
347
348 pub fn add_id(&mut self, id: PackObjectId, obj_type: ObjectType, data: Vec<u8>) -> Result<()> {
371 let pw = self
376 .pack_writer
377 .as_mut()
378 .expect("add_id called after finalize");
379 let entry_start = self.pack_position;
380 let offset = entry_start
381 .checked_sub(self.header_offset)
382 .expect("header_offset should never be past current position");
383
384 self.total_uncompressed += data.len() as u64;
385
386 let mut header_buf = Vec::with_capacity(40);
389 id.encode_tagged(&mut header_buf);
390 super::varint::encode_type_and_size(obj_type, data.len() as u64, &mut header_buf);
391 pw.write_all(&header_buf).map_err(StoreError::from)?;
392 self.pack_position = self
393 .pack_position
394 .checked_add(header_buf.len() as u64)
395 .ok_or_else(|| {
396 StoreError::InvalidObject("streaming pack position overflow".to_string())
397 })?;
398 #[cfg(feature = "zstd")]
402 let csize_pos = entry_start + header_buf.len() as u64;
403
404 let want_compress: bool;
417 #[cfg(feature = "zstd")]
418 {
419 want_compress = self.compression.enabled && data.len() >= self.compression.min_size;
420 }
421 #[cfg(not(feature = "zstd"))]
422 {
423 want_compress = false;
424 }
425 if !want_compress {
426 let mut csize_buf = Vec::with_capacity(10);
429 super::varint::encode_varint(data.len() as u64, &mut csize_buf);
430 pw.write_all(&csize_buf).map_err(StoreError::from)?;
431 self.pack_position = self
432 .pack_position
433 .checked_add(csize_buf.len() as u64)
434 .ok_or_else(|| {
435 StoreError::InvalidObject("streaming pack position overflow".to_string())
436 })?;
437 pw.write_all(&data).map_err(StoreError::from)?;
438 self.pack_position = self
439 .pack_position
440 .checked_add(data.len() as u64)
441 .ok_or_else(|| {
442 StoreError::InvalidObject("streaming pack position overflow".to_string())
443 })?;
444 self.total_compressed += data.len() as u64;
445 } else {
446 #[cfg(feature = "zstd")]
447 {
448 pw.write_all(&[0u8; CSIZE_PLACEHOLDER_LEN])
451 .map_err(StoreError::from)?;
452 self.pack_position = self
453 .pack_position
454 .checked_add(CSIZE_PLACEHOLDER_LEN as u64)
455 .ok_or_else(|| {
456 StoreError::InvalidObject("streaming pack position overflow".to_string())
457 })?;
458 let body_start = self.pack_position;
459 let compressed_size;
460 {
461 let mut counting = CountingWriter::new(&mut *pw);
462 let mut enc =
463 zstd::stream::write::Encoder::new(&mut counting, self.compression.level)
464 .map_err(StoreError::from)?;
465 enc.set_pledged_src_size(Some(data.len() as u64))
470 .map_err(StoreError::from)?;
471 enc.write_all(&data).map_err(StoreError::from)?;
472 enc.finish().map_err(StoreError::from)?;
473 compressed_size = counting.written;
474 }
475 self.pack_position =
476 self.pack_position
477 .checked_add(compressed_size)
478 .ok_or_else(|| {
479 StoreError::InvalidObject(
480 "streaming pack position overflow".to_string(),
481 )
482 })?;
483 let body_end = body_start.checked_add(compressed_size).ok_or_else(|| {
484 StoreError::InvalidObject("streaming pack position overflow".to_string())
485 })?;
486 self.total_compressed += compressed_size;
487
488 let mut csize_bytes = [0u8; CSIZE_PLACEHOLDER_LEN];
493 encode_varint_padded_to_10(compressed_size, &mut csize_bytes);
494 pw.flush().map_err(StoreError::from)?;
495 let inner = pw.get_mut();
496 inner
497 .seek(SeekFrom::Start(csize_pos))
498 .map_err(StoreError::from)?;
499 inner.write_all(&csize_bytes).map_err(StoreError::from)?;
500 inner
501 .seek(SeekFrom::Start(body_end))
502 .map_err(StoreError::from)?;
503 }
504 #[cfg(not(feature = "zstd"))]
505 {
506 unreachable!("compression branch reached without `zstd` feature");
509 }
510 }
511
512 let bucket_idx = bucket_index_for(&id);
516 let bucket = self.get_or_open_bucket(bucket_idx)?;
517 let mut idx_entry = Vec::with_capacity(33 + 8);
518 id.encode_tagged(&mut idx_entry);
519 idx_entry.extend_from_slice(&offset.to_be_bytes());
520 bucket.write_all(&idx_entry).map_err(StoreError::from)?;
521
522 self.object_count += 1;
523 Ok(())
524 }
525
526 fn get_or_open_bucket(&mut self, idx: usize) -> Result<&mut BufWriter<File>> {
527 self.bucket_access_tick = self.bucket_access_tick.wrapping_add(1);
528 let last_used = self.bucket_access_tick;
529 if self.bucket_writers[idx].is_none() {
530 if self.open_bucket_writers >= MAX_OPEN_BUCKET_WRITERS {
531 self.evict_lru_bucket()?;
532 }
533 let path = &self.bucket_paths[idx];
534 let f = OpenOptions::new()
535 .create(true)
536 .append(true)
537 .open(path)
538 .map_err(StoreError::from)?;
539 self.bucket_writers[idx] = Some(BucketWriter {
540 writer: BufWriter::new(f),
541 last_used,
542 });
543 self.open_bucket_writers += 1;
544 } else if let Some(bucket) = self.bucket_writers[idx].as_mut() {
545 bucket.last_used = last_used;
546 }
547 Ok(&mut self.bucket_writers[idx]
548 .as_mut()
549 .expect("just inserted above")
550 .writer)
551 }
552
553 fn evict_lru_bucket(&mut self) -> Result<()> {
554 let Some((idx, _)) = self
555 .bucket_writers
556 .iter()
557 .enumerate()
558 .filter_map(|(idx, bucket)| bucket.as_ref().map(|bucket| (idx, bucket.last_used)))
559 .min_by_key(|(_, last_used)| *last_used)
560 else {
561 return Ok(());
562 };
563
564 if let Some(mut bucket) = self.bucket_writers[idx].take() {
565 bucket.writer.flush().map_err(StoreError::from)?;
566 self.open_bucket_writers -= 1;
567 }
568 Ok(())
569 }
570
571 pub fn finalize(mut self) -> Result<(W, PackStats)> {
582 for bucket in self.bucket_writers.iter_mut().flatten() {
585 bucket.writer.flush().map_err(StoreError::from)?;
586 }
587 for slot in self.bucket_writers.iter_mut() {
590 *slot = None;
591 }
592 self.open_bucket_writers = 0;
593
594 let bw = self
598 .pack_writer
599 .take()
600 .expect("finalize called twice — pack_writer already consumed");
601 let mut writer = bw
602 .into_inner()
603 .map_err(|e| StoreError::from(std::io::Error::other(e.to_string())))?;
604 if let Some(expected) = self.declared_object_count {
605 if expected != self.object_count {
606 return Err(StoreError::InvalidObject(format!(
607 "streaming pack declared {expected} object(s) but added {}",
608 self.object_count
609 )));
610 }
611 } else {
612 writer
613 .seek(SeekFrom::Start(self.header_offset))
614 .map_err(StoreError::from)?;
615 let mut header_bytes = Vec::with_capacity(16);
616 write_container_header(&mut header_bytes, pack_container_spec(), self.object_count);
617 writer.write_all(&header_bytes).map_err(StoreError::from)?;
618 }
619
620 writer
625 .seek(SeekFrom::Start(self.header_offset))
626 .map_err(StoreError::from)?;
627 let mut hasher = blake3::Hasher::new();
628 let mut buf = vec![0u8; 64 * 1024];
629 loop {
630 let n = writer.read(&mut buf).map_err(StoreError::from)?;
631 if n == 0 {
632 break;
633 }
634 hasher.update(&buf[..n]);
635 }
636 let checksum = hasher.finalize();
637
638 writer.seek(SeekFrom::End(0)).map_err(StoreError::from)?;
640 writer
641 .write_all(checksum.as_bytes())
642 .map_err(StoreError::from)?;
643 writer.flush().map_err(StoreError::from)?;
644 if self.durable {
646 writer
647 .sync_data_for_durability()
648 .map_err(StoreError::from)?;
649 }
650
651 let idx_file = File::create(&self.index_path).map_err(StoreError::from)?;
664 let mut idx_writer = BufWriter::new(idx_file);
665 write_index_header(&mut idx_writer, self.object_count)?;
666 let mut entries_written: u64 = 0;
667 for path in self.bucket_paths.iter() {
668 if !path.exists() {
669 continue;
670 }
671 let bucket_bytes = std::fs::read(path).map_err(StoreError::from)?;
672 let mut entries = decode_bucket_file(&bucket_bytes)?;
673 entries.sort_by_key(|(id, _)| *id);
678 for (id, offset) in entries {
679 write_index_entry(&mut idx_writer, id, offset)?;
680 entries_written += 1;
681 }
682 }
683 idx_writer.flush().map_err(StoreError::from)?;
684 let idx_file = idx_writer
686 .into_inner()
687 .map_err(|e| StoreError::from(std::io::Error::other(e.to_string())))?;
688 if self.durable {
689 idx_file.sync_all().map_err(StoreError::from)?;
690 if let Some(parent) = self.index_path.parent() {
691 heddle_fs_prims::fs_atomic::sync_directory(parent).map_err(StoreError::from)?;
692 }
693 }
694 debug_assert_eq!(
695 entries_written, self.object_count,
696 "streaming index entry count drifted from add() count"
697 );
698
699 for path in self.bucket_paths.iter() {
704 let _ = std::fs::remove_file(path);
705 }
706 let _ = std::fs::remove_dir(&self.bucket_dir);
707 self.finalized = true;
708
709 let stats = PackStats {
710 object_count: self.object_count,
711 total_uncompressed: self.total_uncompressed,
712 total_compressed: self.total_compressed,
713 delta_count: 0,
714 compression_ratio: if self.total_uncompressed == 0 {
715 0.0
716 } else {
717 self.total_compressed as f64 / self.total_uncompressed as f64
718 },
719 };
720
721 Ok((writer, stats))
722 }
723}
724
725fn write_index_header<W: Write>(out: &mut W, count: u64) -> Result<()> {
730 super::pack_index::index_header().write_to(out, count)
731}
732
733fn write_index_entry<W: Write>(out: &mut W, id: PackObjectId, offset: u64) -> Result<()> {
737 let mut buf = Vec::with_capacity(33 + 8);
738 id.encode_tagged(&mut buf);
739 buf.extend_from_slice(&offset.to_be_bytes());
740 out.write_all(&buf).map_err(StoreError::from)
741}
742
743#[cfg(feature = "zstd")]
756fn encode_varint_padded_to_10(value: u64, out: &mut [u8; 10]) {
757 let mut v = value;
758 for slot in out.iter_mut().take(9) {
759 *slot = 0x80 | ((v & 0x7F) as u8);
760 v >>= 7;
761 }
762 out[9] = (v & 0x7F) as u8;
763}
764
765impl<W: Write + Read + Seek> Drop for StreamingPackBuilder<W> {
766 fn drop(&mut self) {
767 if self.finalized {
768 return;
769 }
770 for path in self.bucket_paths.iter() {
773 let _ = std::fs::remove_file(path);
774 }
775 let _ = std::fs::remove_dir(&self.bucket_dir);
776 }
777}
778
779fn bucket_index_for(id: &PackObjectId) -> usize {
783 match id {
784 PackObjectId::Hash(h) => HASH_VARIANT * BUCKETS_PER_VARIANT + h.as_bytes()[0] as usize,
785 PackObjectId::StateId(c) => {
786 CHANGEID_VARIANT * BUCKETS_PER_VARIANT + c.as_bytes()[0] as usize
787 }
788 }
789}
790
791fn decode_bucket_file(bytes: &[u8]) -> Result<Vec<(PackObjectId, u64)>> {
796 let mut out = Vec::new();
797 let mut pos = 0;
798 while pos < bytes.len() {
799 let (id, id_len) = PackObjectId::decode_tagged(&bytes[pos..])?;
800 pos += id_len;
801 if pos + 8 > bytes.len() {
802 return Err(StoreError::InvalidObject(
803 "streaming bucket entry truncated at offset".to_string(),
804 ));
805 }
806 let offset = u64::from_be_bytes(bytes[pos..pos + 8].try_into().map_err(|_| {
807 StoreError::InvalidObject("streaming bucket bad offset slice".to_string())
808 })?);
809 pos += 8;
810 out.push((id, offset));
811 }
812 Ok(out)
813}
814
815#[cfg(test)]
818mod tests {
819 use std::io::Cursor;
820
821 use super::*;
822 use crate::{
823 object::StateId,
824 store::pack::{PackReader, PackStats},
825 };
826
827 fn deterministic_hash(seed: u8) -> ContentHash {
828 let mut bytes = [0u8; 32];
832 bytes[0] = seed;
833 for (i, b) in bytes.iter_mut().enumerate().skip(1) {
834 *b = seed.wrapping_mul(31).wrapping_add(i as u8);
835 }
836 ContentHash::from_bytes(bytes)
837 }
838
839 fn deterministic_state_id(seed: u8) -> StateId {
840 let mut bytes = [0u8; 32];
841 bytes[0] = seed;
842 for (i, b) in bytes.iter_mut().enumerate().skip(1) {
843 *b = seed.wrapping_add(i as u8 * 7);
844 }
845 StateId::from_bytes(bytes)
846 }
847
848 fn fresh_builder(
853 tmp: &tempfile::TempDir,
854 ) -> (StreamingPackBuilder<Cursor<Vec<u8>>>, PathBuf, PathBuf) {
855 let bucket_dir = tmp.path().join("buckets");
856 let index_path = tmp.path().join("test.idx");
857 let cursor = Cursor::new(Vec::<u8>::new());
858 let b = StreamingPackBuilder::new(
859 cursor,
860 index_path.clone(),
861 CompressionConfig::default(),
862 bucket_dir.clone(),
863 )
864 .unwrap();
865 (b, bucket_dir, index_path)
866 }
867
868 fn finalize_cursor(
873 b: StreamingPackBuilder<Cursor<Vec<u8>>>,
874 index_path: &std::path::Path,
875 ) -> (Vec<u8>, Vec<u8>, PackStats) {
876 let (cursor, stats) = b.finalize().unwrap();
877 let index_bytes = std::fs::read(index_path).unwrap();
878 (cursor.into_inner(), index_bytes, stats)
879 }
880
881 #[test]
882 fn empty_pack_finalizes_to_valid_zero_count_pack() {
883 let tmp = tempfile::TempDir::new().unwrap();
884 let (b, bucket_dir, idx_path) = fresh_builder(&tmp);
885 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
886
887 assert_eq!(stats.object_count, 0);
888 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
890 assert!(reader.list_ids().is_empty());
891 assert!(
893 !bucket_dir.exists(),
894 "bucket dir should be cleaned on successful finalize"
895 );
896 }
897
898 #[test]
899 fn single_blob_with_hash_id_round_trips() {
900 let tmp = tempfile::TempDir::new().unwrap();
901 let (mut b, _, idx_path) = fresh_builder(&tmp);
902 let hash = deterministic_hash(0x42);
903 let payload = b"hello, streaming pack".to_vec();
904 b.add(hash, ObjectType::Blob, payload.clone()).unwrap();
905 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
906
907 assert_eq!(stats.object_count, 1);
908 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
909 let id = PackObjectId::Hash(hash);
910 assert!(reader.has_object(&id));
911 let (got_type, got_data) = reader.get_object(&id).unwrap().unwrap();
912 assert_eq!(got_type, ObjectType::Blob);
913 assert_eq!(got_data, payload);
914 }
915
916 #[test]
917 fn single_state_with_change_id_round_trips() {
918 let tmp = tempfile::TempDir::new().unwrap();
919 let (mut b, _, idx_path) = fresh_builder(&tmp);
920 let cid = deterministic_state_id(0xa5);
921 let payload = b"serialized-state-bytes".to_vec();
922 b.add_id(
923 PackObjectId::StateId(cid),
924 ObjectType::State,
925 payload.clone(),
926 )
927 .unwrap();
928 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
929
930 assert_eq!(stats.object_count, 1);
931 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
932 let id = PackObjectId::StateId(cid);
933 let (ty, data) = reader.get_object(&id).unwrap().unwrap();
934 assert_eq!(ty, ObjectType::State);
935 assert_eq!(data, payload);
936 }
937
938 #[test]
939 fn mixed_hash_and_changeid_ids_all_retrievable() {
940 let tmp = tempfile::TempDir::new().unwrap();
941 let (mut b, _, idx_path) = fresh_builder(&tmp);
942 let blob_hash = deterministic_hash(0x10);
943 let tree_hash = deterministic_hash(0x20);
944 let state_cid = deterministic_state_id(0x80);
945
946 b.add(blob_hash, ObjectType::Blob, b"blob-bytes".to_vec())
947 .unwrap();
948 b.add(tree_hash, ObjectType::Tree, b"serialized-tree".to_vec())
949 .unwrap();
950 b.add_id(
951 PackObjectId::StateId(state_cid),
952 ObjectType::State,
953 b"serialized-state".to_vec(),
954 )
955 .unwrap();
956
957 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
958 assert_eq!(stats.object_count, 3);
959 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
960 assert_eq!(
961 reader
962 .get_object(&PackObjectId::Hash(blob_hash))
963 .unwrap()
964 .unwrap()
965 .1,
966 b"blob-bytes".to_vec()
967 );
968 assert_eq!(
969 reader
970 .get_object(&PackObjectId::Hash(tree_hash))
971 .unwrap()
972 .unwrap()
973 .1,
974 b"serialized-tree".to_vec()
975 );
976 assert_eq!(
977 reader
978 .get_object(&PackObjectId::StateId(state_cid))
979 .unwrap()
980 .unwrap()
981 .1,
982 b"serialized-state".to_vec()
983 );
984 }
985
986 #[test]
987 fn ten_thousand_objects_round_trip_correctly() {
988 let tmp = tempfile::TempDir::new().unwrap();
992 let (mut b, _, idx_path) = fresh_builder(&tmp);
993 let mut hashes = Vec::with_capacity(10_000);
994 for i in 0..10_000u32 {
995 let h = blake3::hash(&i.to_le_bytes());
998 let hash = ContentHash::from_bytes(*h.as_bytes());
999 hashes.push(hash);
1000 b.add(hash, ObjectType::Blob, format!("payload-{i}").into_bytes())
1001 .unwrap();
1002 }
1003 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
1004 assert_eq!(stats.object_count, 10_000);
1005
1006 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1007 assert_eq!(reader.list_ids().len(), 10_000);
1008 for i in [0, 1, 99, 1234, 5_000, 9_999] {
1010 let id = PackObjectId::Hash(hashes[i]);
1011 let (_ty, data) = reader.get_object(&id).unwrap().unwrap();
1012 assert_eq!(data, format!("payload-{i}").into_bytes());
1013 }
1014 }
1015
1016 #[test]
1017 fn bucket_writers_are_lru_capped_below_fd_limit() {
1018 let tmp = tempfile::TempDir::new().unwrap();
1019 let (mut b, _bucket_dir, idx_path) = fresh_builder(&tmp);
1020 let mut ids = Vec::new();
1021
1022 for i in 0..BUCKETS_PER_VARIANT {
1023 let hash = deterministic_hash(i as u8);
1024 ids.push(PackObjectId::Hash(hash));
1025 b.add(hash, ObjectType::Blob, format!("hash-{i}").into_bytes())
1026 .unwrap();
1027 assert!(
1028 b.open_bucket_writers <= MAX_OPEN_BUCKET_WRITERS,
1029 "open bucket writers should stay capped"
1030 );
1031 }
1032
1033 for i in 0..BUCKETS_PER_VARIANT {
1034 let cid = deterministic_state_id(i as u8);
1035 ids.push(PackObjectId::StateId(cid));
1036 b.add_id(
1037 PackObjectId::StateId(cid),
1038 ObjectType::State,
1039 format!("state-{i}").into_bytes(),
1040 )
1041 .unwrap();
1042 assert!(
1043 b.open_bucket_writers <= MAX_OPEN_BUCKET_WRITERS,
1044 "open bucket writers should stay capped"
1045 );
1046 }
1047
1048 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
1049 assert_eq!(stats.object_count, TOTAL_BUCKETS as u64);
1050 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1051 for id in ids {
1052 assert!(reader.has_object(&id), "missing id {id:?}");
1053 }
1054 }
1055
1056 #[test]
1057 fn index_id_sort_order_matches_packbuilder_output() {
1058 use crate::store::pack::PackBuilder;
1066 let payloads: Vec<(PackObjectId, ObjectType, Vec<u8>)> = (0..200u32)
1067 .map(|i| {
1068 let h = blake3::hash(&i.to_le_bytes());
1069 (
1070 PackObjectId::Hash(ContentHash::from_bytes(*h.as_bytes())),
1071 if i % 3 == 0 {
1072 ObjectType::Tree
1073 } else {
1074 ObjectType::Blob
1075 },
1076 format!("body-{i}").into_bytes(),
1077 )
1078 })
1079 .collect();
1080
1081 let compression = CompressionConfig {
1084 max_delta_size: 0,
1085 ..CompressionConfig::default()
1086 };
1087 let mut classic = PackBuilder::new(compression);
1088 for (id, ty, data) in payloads.iter() {
1089 classic.add_id(*id, *ty, data.clone());
1090 }
1091 let (classic_pack, classic_index, _) = classic.build().unwrap();
1092 let classic_reader = PackReader::from_bytes(classic_pack, classic_index).unwrap();
1093
1094 let tmp = tempfile::TempDir::new().unwrap();
1095 let bucket_dir = tmp.path().join("buckets");
1096 let idx_path = tmp.path().join("test.idx");
1097 let cursor = Cursor::new(Vec::<u8>::new());
1098 let mut streaming =
1099 StreamingPackBuilder::new(cursor, idx_path.clone(), compression, bucket_dir).unwrap();
1100 for (id, ty, data) in payloads.iter() {
1101 streaming.add_id(*id, *ty, data.clone()).unwrap();
1102 }
1103 let (streaming_pack, streaming_index, _) = finalize_cursor(streaming, &idx_path);
1104 let streaming_reader = PackReader::from_bytes(streaming_pack, streaming_index).unwrap();
1105
1106 assert_eq!(
1109 streaming_reader.list_ids(),
1110 classic_reader.list_ids(),
1111 "streaming and classic indices should report the same id sequence"
1112 );
1113 for (id, _ty, want) in payloads.iter().take(10).chain(payloads.iter().skip(190)) {
1116 let (_, got) = streaming_reader.get_object(id).unwrap().unwrap();
1117 assert_eq!(&got, want);
1118 let (_, classic_got) = classic_reader.get_object(id).unwrap().unwrap();
1119 assert_eq!(got, classic_got);
1120 }
1121 }
1122
1123 #[test]
1124 fn corrupted_pack_fails_checksum_verification() {
1125 let tmp = tempfile::TempDir::new().unwrap();
1126 let (mut b, _, idx_path) = fresh_builder(&tmp);
1127 b.add(
1128 deterministic_hash(0x01),
1129 ObjectType::Blob,
1130 b"some bytes".to_vec(),
1131 )
1132 .unwrap();
1133 let (mut pack_data, index_data, _) = finalize_cursor(b, &idx_path);
1134 let body_byte = 18; pack_data[body_byte] ^= 0xff;
1137 let result = PackReader::from_bytes(pack_data, index_data);
1138 assert!(
1139 result.is_err(),
1140 "PackReader should reject pack with mutated body"
1141 );
1142 }
1143
1144 #[test]
1145 fn pack_count_in_header_matches_index_entry_count() {
1146 let tmp = tempfile::TempDir::new().unwrap();
1147 let (mut b, _, idx_path) = fresh_builder(&tmp);
1148 for i in 0..7u8 {
1149 b.add(
1150 deterministic_hash(i),
1151 ObjectType::Blob,
1152 format!("p{i}").into_bytes(),
1153 )
1154 .unwrap();
1155 }
1156 let (pack_data, index_data, _) = finalize_cursor(b, &idx_path);
1157 let count = u64::from_be_bytes(pack_data[8..16].try_into().unwrap());
1159 assert_eq!(count, 7);
1160 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1161 assert_eq!(reader.list_ids().len(), 7);
1162 }
1163
1164 #[test]
1165 fn declared_pack_count_is_written_before_finalize() {
1166 let tmp = tempfile::TempDir::new().unwrap();
1167 let bucket_dir = tmp.path().join("buckets");
1168 let idx_path = tmp.path().join("test.idx");
1169 let cursor = Cursor::new(Vec::<u8>::new());
1170 let mut b = StreamingPackBuilder::new_with_object_count(
1171 cursor,
1172 idx_path.clone(),
1173 CompressionConfig::default(),
1174 bucket_dir,
1175 2,
1176 )
1177 .unwrap();
1178
1179 b.flush_pack().unwrap();
1180 let initial = b.pack_writer.as_ref().unwrap().get_ref().get_ref().clone();
1181 assert_eq!(u64::from_be_bytes(initial[8..16].try_into().unwrap()), 2);
1182
1183 let hash = deterministic_hash(0x40);
1184 b.add(hash, ObjectType::Blob, b"known-count-entry".to_vec())
1185 .unwrap();
1186 b.flush_pack().unwrap();
1187 let after_add = b.pack_writer.as_ref().unwrap().get_ref().get_ref().clone();
1188 assert_eq!(u64::from_be_bytes(after_add[8..16].try_into().unwrap()), 2);
1189
1190 let second_hash = deterministic_hash(0x41);
1191 b.add(second_hash, ObjectType::Blob, b"second-entry".to_vec())
1192 .unwrap();
1193 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
1194
1195 assert_eq!(stats.object_count, 2);
1196 assert_eq!(u64::from_be_bytes(pack_data[8..16].try_into().unwrap()), 2);
1197 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1198 assert!(reader.has_object(&PackObjectId::Hash(hash)));
1199 assert!(reader.has_object(&PackObjectId::Hash(second_hash)));
1200 }
1201
1202 #[test]
1203 fn declared_pack_count_mismatch_fails_finalize() {
1204 let tmp = tempfile::TempDir::new().unwrap();
1205 let bucket_dir = tmp.path().join("buckets");
1206 let idx_path = tmp.path().join("test.idx");
1207 let cursor = Cursor::new(Vec::<u8>::new());
1208 let mut b = StreamingPackBuilder::new_with_object_count(
1209 cursor,
1210 idx_path,
1211 CompressionConfig::default(),
1212 bucket_dir,
1213 2,
1214 )
1215 .unwrap();
1216
1217 b.add(
1218 deterministic_hash(0x50),
1219 ObjectType::Blob,
1220 b"only-entry".to_vec(),
1221 )
1222 .unwrap();
1223 let error = b.finalize().unwrap_err();
1224
1225 assert!(
1226 error
1227 .to_string()
1228 .contains("streaming pack declared 2 object(s) but added 1")
1229 );
1230 }
1231
1232 #[test]
1233 fn bucket_files_are_cleaned_on_successful_finalize() {
1234 let tmp = tempfile::TempDir::new().unwrap();
1235 let bucket_dir = tmp.path().join("buckets");
1236 let idx_path = tmp.path().join("test.idx");
1237 let cursor = Cursor::new(Vec::<u8>::new());
1238 let mut b = StreamingPackBuilder::new(
1239 cursor,
1240 idx_path.clone(),
1241 CompressionConfig::default(),
1242 bucket_dir.clone(),
1243 )
1244 .unwrap();
1245 for i in 0..50u8 {
1246 b.add(deterministic_hash(i), ObjectType::Blob, vec![i; 32])
1247 .unwrap();
1248 }
1249 assert!(bucket_dir.exists());
1251 let bucket_count = std::fs::read_dir(&bucket_dir).unwrap().count();
1252 assert!(bucket_count > 0, "bucket dir should hold some files");
1253 let _ = finalize_cursor(b, &idx_path);
1254 assert!(
1255 !bucket_dir.exists(),
1256 "bucket dir should be removed on finalize"
1257 );
1258 }
1259
1260 #[test]
1261 fn bucket_files_are_cleaned_on_drop_without_finalize() {
1262 let tmp = tempfile::TempDir::new().unwrap();
1263 let bucket_dir = tmp.path().join("buckets");
1264 let idx_path = tmp.path().join("test.idx");
1265 {
1266 let cursor = Cursor::new(Vec::<u8>::new());
1267 let mut b = StreamingPackBuilder::new(
1268 cursor,
1269 idx_path.clone(),
1270 CompressionConfig::default(),
1271 bucket_dir.clone(),
1272 )
1273 .unwrap();
1274 for i in 0..10u8 {
1275 b.add(deterministic_hash(i), ObjectType::Blob, vec![0; 32])
1276 .unwrap();
1277 }
1278 assert!(bucket_dir.exists());
1279 }
1281 assert!(
1282 !idx_path.exists(),
1283 "no index file should have been created without finalize"
1284 );
1285 assert!(
1286 !bucket_dir.exists(),
1287 "bucket dir should be removed on Drop when finalize never ran"
1288 );
1289 }
1290
1291 #[test]
1292 fn large_blob_streams_to_disk_without_double_buffering() {
1293 let tmp = tempfile::TempDir::new().unwrap();
1298 let bucket_dir = tmp.path().join("buckets");
1299 let pack_path = tmp.path().join("pack.dat");
1300 let idx_path = tmp.path().join("pack.idx");
1301 let file = std::fs::OpenOptions::new()
1302 .read(true)
1303 .write(true)
1304 .create(true)
1305 .truncate(true)
1306 .open(&pack_path)
1307 .unwrap();
1308 let mut b = StreamingPackBuilder::new(
1309 file,
1310 idx_path.clone(),
1311 CompressionConfig::default(),
1312 bucket_dir,
1313 )
1314 .unwrap();
1315 let payload: Vec<u8> = (0..4 * 1024 * 1024u32).map(|i| (i & 0xff) as u8).collect();
1316 let hash = deterministic_hash(0xff);
1317 b.add(hash, ObjectType::Blob, payload.clone()).unwrap();
1318 let (_, stats) = b.finalize().unwrap();
1319 let index_data = std::fs::read(&idx_path).unwrap();
1320 assert_eq!(stats.object_count, 1);
1321 let pack_bytes = std::fs::read(&pack_path).unwrap();
1322 let reader = PackReader::from_bytes(pack_bytes, index_data).unwrap();
1325 let (_ty, got) = reader
1326 .get_object(&PackObjectId::Hash(hash))
1327 .unwrap()
1328 .unwrap();
1329 assert_eq!(got, payload);
1330 }
1331
1332 #[test]
1333 fn bucket_distribution_for_random_hashes_is_roughly_uniform() {
1334 let tmp = tempfile::TempDir::new().unwrap();
1341 let bucket_dir = tmp.path().join("buckets");
1342 let idx_path = tmp.path().join("test.idx");
1343 let cursor = Cursor::new(Vec::<u8>::new());
1344 let mut b = StreamingPackBuilder::new(
1345 cursor,
1346 idx_path.clone(),
1347 CompressionConfig::default(),
1348 bucket_dir.clone(),
1349 )
1350 .unwrap();
1351 for i in 0..1024u32 {
1352 let h = blake3::hash(&i.to_le_bytes());
1353 let hash = ContentHash::from_bytes(*h.as_bytes());
1354 b.add(hash, ObjectType::Blob, b"x".to_vec()).unwrap();
1355 }
1356 b.pack_writer.as_mut().unwrap().flush().unwrap();
1358 let mut max_entries = 0usize;
1359 let entry_size = 33 + 8; for path in b.bucket_paths.iter() {
1361 if path.exists() {
1362 let size = std::fs::metadata(path).unwrap().len() as usize;
1363 let entries = size / entry_size;
1364 if entries > max_entries {
1365 max_entries = entries;
1366 }
1367 }
1368 }
1369 assert!(
1372 max_entries <= 16,
1373 "max bucket has {max_entries} entries; uniform expected ~4"
1374 );
1375 let _ = finalize_cursor(b, &idx_path);
1376 }
1377
1378 #[test]
1379 fn finalize_returns_correct_stats() {
1380 let tmp = tempfile::TempDir::new().unwrap();
1381 let (mut b, _, idx_path) = fresh_builder(&tmp);
1382 let payload = vec![0xabu8; 1024];
1383 for i in 0..5u8 {
1384 b.add(deterministic_hash(i), ObjectType::Blob, payload.clone())
1385 .unwrap();
1386 }
1387 let (_, _, stats) = finalize_cursor(b, &idx_path);
1388 assert_eq!(stats.object_count, 5);
1389 assert_eq!(stats.total_uncompressed, 5 * 1024);
1390 assert!(stats.total_compressed > 0);
1391 assert!(stats.compression_ratio > 0.0);
1392 assert_eq!(stats.delta_count, 0, "streaming builder never deltas");
1393 }
1394
1395 #[cfg(feature = "zstd")]
1396 #[test]
1397 fn streaming_compression_roundtrips_through_zstd_frame() {
1398 let tmp = tempfile::TempDir::new().unwrap();
1406 let (mut b, _, idx_path) = fresh_builder(&tmp);
1407 let payload = vec![0u8; 64 * 1024];
1410 let hash = deterministic_hash(0x77);
1411 b.add(hash, ObjectType::Blob, payload.clone()).unwrap();
1412 let (pack_data, index_data, stats) = finalize_cursor(b, &idx_path);
1413 assert!(
1414 stats.total_compressed < stats.total_uncompressed,
1415 "expected compression ratio < 1.0, got {}/{}",
1416 stats.total_compressed,
1417 stats.total_uncompressed
1418 );
1419 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1420 let (_ty, got) = reader
1421 .get_object(&PackObjectId::Hash(hash))
1422 .unwrap()
1423 .unwrap();
1424 assert_eq!(got, payload);
1425 }
1426
1427 #[cfg(feature = "zstd")]
1428 #[test]
1429 fn padded_varint_decodes_to_original_value_for_canonical_decoder() {
1430 let cases: &[u64] = &[0, 1, 127, 128, 4096, 1_000_000, 1_000_000_000_000, u64::MAX];
1436 for &value in cases {
1437 let mut buf = [0u8; 10];
1438 super::encode_varint_padded_to_10(value, &mut buf);
1439 let (decoded, consumed) = super::super::varint::decode_varint(&buf)
1440 .expect("padded varint should always decode");
1441 assert_eq!(decoded, value, "varint roundtrip failed for {value}");
1442 assert_eq!(
1443 consumed, 10,
1444 "padded encoding should consume all 10 bytes for {value}"
1445 );
1446 }
1447 }
1448
1449 #[cfg(feature = "zstd")]
1450 #[test]
1451 fn streaming_path_does_not_buffer_compressed_payload_in_memory() {
1452 let tmp = tempfile::TempDir::new().unwrap();
1465 let bucket_dir = tmp.path().join("buckets");
1466 let pack_path = tmp.path().join("pack.dat");
1467 let idx_path = tmp.path().join("pack.idx");
1468 let file = std::fs::OpenOptions::new()
1469 .read(true)
1470 .write(true)
1471 .create(true)
1472 .truncate(true)
1473 .open(&pack_path)
1474 .unwrap();
1475 let mut b = StreamingPackBuilder::new(
1476 file,
1477 idx_path.clone(),
1478 CompressionConfig::default(),
1479 bucket_dir,
1480 )
1481 .unwrap();
1482 let payload = vec![0xa5u8; 8 * 1024 * 1024];
1483 let hash = deterministic_hash(0x66);
1484 b.add(hash, ObjectType::Blob, payload.clone()).unwrap();
1485 let mid_size = std::fs::metadata(&pack_path).unwrap().len();
1489 assert!(
1490 mid_size > 16 + 40,
1491 "pack file should hold real entry data after add; size={mid_size}"
1492 );
1493 let (_, _) = b.finalize().unwrap();
1494 let pack_bytes = std::fs::read(&pack_path).unwrap();
1495 let index_bytes = std::fs::read(&idx_path).unwrap();
1496 let reader = PackReader::from_bytes(pack_bytes, index_bytes).unwrap();
1497 let (_ty, got) = reader
1498 .get_object(&PackObjectId::Hash(hash))
1499 .unwrap()
1500 .unwrap();
1501 assert_eq!(got, payload);
1502 }
1503
1504 #[test]
1505 fn list_ids_returns_all_added_ids_sorted() {
1506 let tmp = tempfile::TempDir::new().unwrap();
1507 let (mut b, _, idx_path) = fresh_builder(&tmp);
1508 let mut added: Vec<PackObjectId> = Vec::new();
1509 for seed in [0x05u8, 0xa0, 0x12, 0x9f, 0x33] {
1511 let id = PackObjectId::Hash(deterministic_hash(seed));
1512 b.add_id(id, ObjectType::Blob, vec![seed; 4]).unwrap();
1513 added.push(id);
1514 }
1515 for seed in [0x80u8, 0x10, 0xff] {
1516 let id = PackObjectId::StateId(deterministic_state_id(seed));
1517 b.add_id(id, ObjectType::State, vec![seed; 4]).unwrap();
1518 added.push(id);
1519 }
1520 let (pack_data, index_data, _) = finalize_cursor(b, &idx_path);
1521 let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
1522 let mut got = reader.list_ids();
1523 let mut sorted = got.clone();
1526 sorted.sort();
1527 assert_eq!(got, sorted, "list_ids must come back sorted");
1528 added.sort();
1530 got.sort();
1531 assert_eq!(got, added);
1532 }
1533}