Skip to main content

heddle_pack/store/pack/
streaming_builder.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Streaming pack builder for bounded-memory imports.
3//!
4//! `PackBuilder` accumulates every `(id, type, data)` tuple in memory
5//! before producing a pack. That's fine for sync-protocol packets and
6//! small batches, but the import path can produce millions of objects
7//! and would OOM on large repos.
8//!
9//! `StreamingPackBuilder` removes the in-memory buffering by:
10//!
11//! 1. **Streaming pack data to disk** as objects are added. Compression
12//!    runs per-object (the existing zstd path is non-streaming, so the
13//!    one compressed payload is held briefly in a `Vec<u8>` before
14//!    being written), but the writer never holds more than one
15//!    object's worth of data plus its `BufWriter` capacity.
16//!
17//! 2. **External sorting the index** via 512 hash-prefix bucket files
18//!    on disk (256 for `Hash` ids, 256 for `StateId` ids). Each
19//!    `add()` appends one fixed-shape `(id, offset)` record to the
20//!    bucket whose first byte matches the id's first inner byte. At
21//!    finalize, each bucket is small enough to sort in memory; the
22//!    concatenation of `Hash` buckets followed by `StateId` buckets
23//!    in byte order produces the exact same global sort `PackBuilder`
24//!    would have via `entries.sort_by_key(|e| e.id)`.
25//!
26//! ## Memory bound
27//!
28//! - Pack data on disk: streamed; only one compressed object held in
29//!   memory at a time.
30//! - Index entries in bucket buffers: at most 32 bucket files are held
31//!   open at once, each behind a default-capacity `BufWriter` (~8 KB),
32//!   so peak buffering is ~256 KB.
33//! - Sort scratch at finalize: O(largest bucket). For uniformly-
34//!   distributed BLAKE3 hashes / ULID change-ids and N total objects,
35//!   the largest bucket is ~N/256 entries ≈ 40 bytes each. Even at
36//!   100 M objects that's ~16 MB peak.
37//!
38//! Net peak memory: ~20 MB regardless of repo size, modulo the size
39//! of the largest single object (which is unavoidable while the zstd
40//! API is non-streaming).
41//!
42//! ## Trade-offs vs `PackBuilder`
43//!
44//! - **No delta encoding.** Streaming and sliding-window deltas are
45//!   incompatible — delta search needs random access to recently-
46//!   written objects. The import path runs with deltas disabled
47//!   anyway (the cost-benefit is bad on real Heddle history), so this
48//!   is a non-issue for the call site that motivated this builder.
49//! - **No path-grouped reordering.** Entries land in the order added.
50//! - **Output is a pack file at a path** rather than `(Vec<u8>, Vec<u8>)`.
51//!   Callers pair this with `objects::store::ObjectStore::install_pack_from_path`
52//!   which moves/installs the pack without copying it through RAM.
53//! - **Re-reads the pack at finalize** to compute the BLAKE3 trailer
54//!   checksum (the pack format hashes header+body, and the count goes
55//!   in the header — we patch it on finalize via seek-back, then
56//!   re-stream the body to the hasher). 2× sequential disk I/O on the
57//!   pack data is the cost of sticking with the current format. A
58//!   future format change could put the count in the footer to avoid
59//!   the second pass.
60
61use 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/// How many bytes to reserve for the compressed-size varint in the
72/// streaming path. 10 is enough to encode any `u64` (max 9 7-bit
73/// continuation bytes plus 1 terminator). After streaming we patch
74/// the placeholder with a non-canonical varint that pads to exactly
75/// this length. Only the zstd-enabled compress path uses it.
76#[cfg(feature = "zstd")]
77const CSIZE_PLACEHOLDER_LEN: usize = 10;
78use crate::{
79    object::ContentHash,
80    store::{Result, StoreError},
81};
82
83/// Number of buckets per id variant. 256 = one bucket per first byte
84/// of the inner id. We want the bucket boundaries to align with the
85/// `PackObjectId`'s `Ord` derivation (variant tag major, inner bytes
86/// minor) so the concatenated bucket output matches what
87/// `PackIndex::sort()` would have produced.
88const BUCKETS_PER_VARIANT: usize = 256;
89/// 256 for Hash ids + 256 for StateId ids.
90const TOTAL_BUCKETS: usize = BUCKETS_PER_VARIANT * 2;
91/// Cap concurrently-open index-bucket files. macOS GUI-launched
92/// processes commonly inherit a 256-fd soft limit; imports also need
93/// room for Git pack/index files, sqlite maps, the output pack, etc.
94const MAX_OPEN_BUCKET_WRITERS: usize = 32;
95
96/// Variant indices into the `bucket_*` arrays. `Hash` ids fill the
97/// lower half (matches the variant order in `PackObjectId` which makes
98/// `Hash(_) < StateId(_)`).
99const HASH_VARIANT: usize = 0;
100const CHANGEID_VARIANT: usize = 1;
101
102/// Fsync staged pack bytes after finalize flush (Wave 5 L7).
103///
104/// Production writers are [`File`]; in-memory [`Cursor`] tests no-op.
105/// Publish still re-fsyncs at `publish_file_durable` install; this closes
106/// the pre-publish window if a caller inspects staged files after finalize.
107pub 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
123/// Streaming pack builder. Held generic over the pack writer (`File`
124/// in production, `Cursor<Vec<u8>>` in tests).
125pub struct StreamingPackBuilder<W: Write + Read + Seek> {
126    /// Writer for the pack's `[header][body]` content. The trailer
127    /// checksum is appended to the same writer at `finalize`.
128    /// Wrapped in `Option` so `finalize` can `.take()` it out without
129    /// running afoul of the `Drop` impl's restriction on moving fields.
130    /// `None` after `finalize` succeeds.
131    pack_writer: Option<BufWriter<W>>,
132    /// Position in the pack writer where the header was written, so
133    /// we can seek back at finalize and patch the real `object_count`
134    /// into bytes 8..16.
135    header_offset: u64,
136    /// Logical append position. Avoids flushing the buffered writer before
137    /// every object just to ask the file for its current offset.
138    pack_position: u64,
139    object_count: u64,
140    declared_object_count: Option<u64>,
141    total_uncompressed: u64,
142    total_compressed: u64,
143    /// Compression knobs. Only consulted when the `zstd` feature is on
144    /// (`enabled` and `min_size` decide whether each entry compresses;
145    /// `level` parameterizes the encoder). Without `zstd` every entry
146    /// takes the raw branch and this field is just along for the ride.
147    #[cfg_attr(not(feature = "zstd"), allow(dead_code))]
148    compression: CompressionConfig,
149    /// Directory holding the 512 bucket files. Owned by the builder
150    /// so we can clean up on `Drop` if `finalize` is never called.
151    bucket_dir: PathBuf,
152    /// Buckets `[variant][prefix_byte]` → optional buffered file.
153    /// Lazily opened on first write and capped with LRU eviction so a
154    /// large import cannot exhaust the process fd limit.
155    bucket_writers: Vec<Option<BucketWriter>>,
156    open_bucket_writers: usize,
157    bucket_access_tick: u64,
158    bucket_paths: Vec<PathBuf>,
159    /// File path where the pack index is materialized at `finalize`.
160    /// Bytes are written incrementally as buckets are sorted, so the
161    /// index never sits in memory in its entirety.
162    index_path: PathBuf,
163    /// Set true on `finalize` so `Drop` knows the bucket dir was
164    /// already cleaned and shouldn't be removed again.
165    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    /// Open a streaming builder against `pack_writer`, using
202    /// `bucket_dir` for transient index buckets and writing the
203    /// finalized index to `index_path`. The bucket dir is created if
204    /// it doesn't exist; on a successful `finalize` it's removed
205    /// (along with any bucket files left in it).
206    ///
207    /// `index_path` is *not* created by `new` — opening happens at
208    /// finalize so a misconfigured caller doesn't leave an empty index
209    /// file behind on early failure. It's still recorded here so
210    /// `finalize` can write to a known location and the caller can
211    /// install the file by path.
212    ///
213    /// The `pack_writer` must support `Read` because finalize re-streams
214    /// the body to compute the trailer checksum — see the module-level
215    /// note on the format.
216    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    /// Open a streaming builder whose final object count is known up front.
226    ///
227    /// This writes the real count into the pack header immediately, so callers
228    /// may safely stream already-flushed pack bytes before `finalize()` appends
229    /// the trailer checksum. `finalize()` still verifies that exactly this many
230    /// objects were added before producing the index.
231    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    /// Open a known-count builder for a transient transfer spool.
249    ///
250    /// The source object store remains authoritative, so these files only need
251    /// to be flushed for same-process readers; crash durability would add an
252    /// fsync to the pack, index, and directory for every unary push without
253    /// making the transfer any safer.
254    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        // Write a placeholder header with `count = 0` unless the caller knows
288        // the count up front. The known-count path lets network senders tail
289        // flushed pack bytes before finalize without later mutating bytes they
290        // already sent.
291        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    /// Flush pack bytes written so far to the underlying writer.
333    ///
334    /// Used by hosted sync's interleaved build/send path: after each complete
335    /// entry is added, the sender flushes and drains full chunks from the file.
336    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    /// Add an object with a content-hash id.
344    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    /// Add an object with an explicit id. Mirrors [`super::PackBuilder::add_id`].
349    ///
350    /// # Memory shape
351    ///
352    /// Per-entry, the only allocations are:
353    ///
354    /// - `data: Vec<u8>` (the input, owned by the caller — comes from
355    ///   gix' `find_object` and isn't ours to stream further).
356    /// - A ~40-byte stack scratch for the entry header.
357    /// - zstd's internal compression context (~128 KB constant).
358    /// - One 50-byte index-bucket entry buffered into the bucket's
359    ///   `BufWriter`.
360    ///
361    /// The compressed payload is **never materialized** as a `Vec<u8>` —
362    /// it streams directly through `zstd::stream::write::Encoder` into
363    /// the pack writer. The pack format requires a `compressed_size`
364    /// varint *before* the compressed bytes, which we don't know yet
365    /// when we write the header; we reserve a 10-byte placeholder and
366    /// seek-back to patch it after the encoder finishes. Heddle's
367    /// varint decoder accepts non-canonical encodings (it walks
368    /// continuation bits without enforcing minimum-byte form), so the
369    /// padded write decodes back to the same value any reader expects.
370    pub fn add_id(&mut self, id: PackObjectId, obj_type: ObjectType, data: Vec<u8>) -> Result<()> {
371        // Compute the entry's offset relative to the header from our logical
372        // append cursor. Asking the underlying file for its position would
373        // flush the BufWriter on every object, which defeats the streaming
374        // sender's chunk-sized drain cadence.
375        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        // Phase 1: write the entry header up to (but not including) the
387        // compressed-size varint. Always small, fits in `entry_header_buf`.
388        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        // Only consumed by the zstd-enabled streaming branch below, but
399        // we compute it here while we already have `header_buf`'s length
400        // in scope.
401        #[cfg(feature = "zstd")]
402        let csize_pos = entry_start + header_buf.len() as u64;
403
404        // Phase 2: stream the compressed payload. We branch here on
405        // whether to compress at all — for tiny objects (`< min_size`)
406        // the bulk path traditionally wrote raw bytes to skip zstd
407        // overhead, and the reader's existing `compressed_size ==
408        // uncompressed_size` heuristic in `pack_reader.rs:128` reads
409        // raw entries back unchanged. We preserve that policy.
410        // `want_compress` gates the zstd path. Even with the feature
411        // enabled we fall through to raw for tiny entries (where
412        // zstd's frame overhead dominates) or when the caller
413        // explicitly disabled compression in `CompressionConfig`.
414        // Without the `zstd` Cargo feature, every entry takes the raw
415        // branch — same fallback shape as `compress_pack_payload`.
416        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            // Raw entry: known compressed_size = data.len(). One canonical
427            // varint + the data itself. No seek-back needed.
428            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                // Streaming entry: reserve 10 bytes for compressed_size,
449                // stream-compress the payload, then seek back to patch.
450                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                    // Pass the source size so the zstd frame's optional
466                    // Frame Content Size field is set — lets decoders
467                    // preallocate output buffers and validates that we
468                    // wrote exactly what we promised at finish().
469                    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                // Seek back over the placeholder, write a 10-byte
489                // non-canonical varint encoding the actual compressed_size,
490                // then seek forward to where we left off so subsequent
491                // adds append correctly.
492                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: `want_compress` is forced to `false`
507                // when the `zstd` feature is off.
508                unreachable!("compression branch reached without `zstd` feature");
509            }
510        }
511
512        // Append the index entry (id || u64 BE offset) to the bucket
513        // matching the id's first inner byte. The bucket file is opened
514        // lazily so a sparse pack only creates files it actually uses.
515        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    /// Close the pack: patch the header count, append the BLAKE3
572    /// trailer, build the sorted index from bucket files, and clean up
573    /// the bucket directory. Returns `(pack_writer, index_bytes,
574    /// stats)` so the caller can install the pack into its store.
575    ///
576    /// On any failure the bucket dir is left in place; rerunning the
577    /// import will overwrite stale bucket files (they're keyed by
578    /// fixed name, not content) so this isn't a correctness issue —
579    /// just a small amount of disk churn until the next clean
580    /// finalize.
581    pub fn finalize(mut self) -> Result<(W, PackStats)> {
582        // 1. Flush every bucket so reads in the next phase see all
583        //    queued entries. `flatten()` skips the never-opened slots.
584        for bucket in self.bucket_writers.iter_mut().flatten() {
585            bucket.writer.flush().map_err(StoreError::from)?;
586        }
587        // Drop the writers so the OS file handles close before we
588        // re-open the same paths for reading.
589        for slot in self.bucket_writers.iter_mut() {
590            *slot = None;
591        }
592        self.open_bucket_writers = 0;
593
594        // 2. Patch the pack header with the real object count unless the
595        //    caller declared it up front, then re-stream the [header][body]
596        //    bytes to compute the trailer checksum.
597        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        // 3. Hash the on-disk content from header_offset to current
621        //    position (which is just past the body). One sequential
622        //    pass; the BufWriter we drained is gone so this read is
623        //    on the raw writer.
624        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        // 4. Append the trailer checksum.
639        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        // L7: durable staged pack before return (File fsync; Cursor no-op).
645        if self.durable {
646            writer
647                .sync_data_for_durability()
648                .map_err(StoreError::from)?;
649        }
650
651        // 5. Stream the final sorted index directly to disk. We open
652        //    a `BufWriter` against `index_path`, write the index
653        //    container header (magic + version + count — count is
654        //    already known from the per-add bookkeeping), then walk
655        //    the 512 buckets in `(variant, prefix)` order, sorting
656        //    each in memory and writing entries to the file as they
657        //    come off the sort. The intermediate `PackIndex` Vec —
658        //    O(K) in the previous implementation — is gone; the
659        //    largest in-memory state is one bucket's worth of entries.
660        //    Bucket distribution is uniform via BLAKE3 so each bucket
661        //    is ~K/256 entries × ~50 bytes; even at 100M objects that's
662        //    a ~16 MB sort scratch.
663        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            // Local sort by `PackObjectId` matches the global sort
674            // because all entries in a bucket share the same variant
675            // tag *and* the same first inner byte; only the remaining
676            // bytes differ between them.
677            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        // L7: durable staged index file + parent dirent for rename/read.
685        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        // 6. Clean up the bucket dir so the heddle store doesn't carry
700        //    transient artifacts. Deletion failures are non-fatal —
701        //    the dir is uniquely named per import so leftovers are at
702        //    worst stale, not corrupting.
703        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
725/// Write the index container header to `out`. Mirrors
726/// [`PackIndex::to_bytes`]'s prefix exactly (4-byte magic, 4-byte
727/// big-endian version, 8-byte big-endian count) so a reader written
728/// against the existing format works without modification.
729fn write_index_header<W: Write>(out: &mut W, count: u64) -> Result<()> {
730    super::pack_index::index_header().write_to(out, count)
731}
732
733/// Append one `(id, offset)` index entry to `out`. The encoding
734/// matches [`PackIndex::to_bytes`]: tagged id immediately followed by
735/// an 8-byte big-endian offset.
736fn 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/// Encode a `u64` as a non-canonical 10-byte LEB128 varint. The first
744/// 9 bytes always set the continuation bit (`0x80`), the 10th never
745/// does — so the decoder reads exactly 10 bytes regardless of the
746/// value. Used by the streaming path to reserve a fixed-width
747/// placeholder for `compressed_size` before stream-compressing the
748/// payload, then patch the placeholder with the actual size after.
749///
750/// `decode_varint` ignores the canonicalness of the encoding (it
751/// walks continuation bits without checking minimum-byte form), so
752/// the value round-trips exactly. Cost is up to 9 wasted bytes per
753/// entry, ~115 KB on a 13 K-entry import — negligible relative to
754/// the pack body.
755#[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        // Best-effort cleanup of bucket dir on abort. Errors here are
771        // suppressed because Drop can't propagate them.
772        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
779/// Map a `PackObjectId` to one of `TOTAL_BUCKETS` buckets. The variant
780/// (Hash vs StateId) picks the upper half; the first byte of the
781/// inner id picks the slot within the half.
782fn 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
791/// Decode `(id, offset)` records from a bucket file. The format
792/// matches `PackObjectId::encode_tagged` followed by a u64 BE offset,
793/// repeated. Unrecognized tags or truncated trailers fail loudly —
794/// we wrote the bytes, so any corruption is a bug, not user input.
795fn 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// ---------------------- Tests ----------------------
816
817#[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        // Spread `seed` across the high byte so different seeds end up
829        // in different hash-prefix buckets. We don't actually want
830        // collisions in the tests that check distribution.
831        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    /// Test rig: returns the builder, the bucket dir (for cleanup
849    /// inspection), and the index path the builder will write at
850    /// finalize. The index path lives in the temp dir so it gets
851    /// auto-cleaned with `tmp`.
852    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    /// Finalize the builder and return `(pack_bytes, index_bytes, stats)`.
869    /// The index bytes are read back from the file the builder wrote
870    /// to — verifying that the streaming index path actually produced
871    /// readable bytes.
872    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        // PackReader can parse the empty pack and reports zero objects.
889        let reader = PackReader::from_bytes(pack_data, index_data).unwrap();
890        assert!(reader.list_ids().is_empty());
891        // Bucket dir was removed.
892        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        // Stresses the bucket sort: 10K objects spread across
989        // 256 hash buckets averages 40 entries per bucket — well
990        // within in-memory sort capacity but covers every bucket.
991        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            // Use BLAKE3 over the index so first-byte distribution is
996            // pseudo-uniform across the 256 hash buckets.
997            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        // Spot-check ten across the range.
1009        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        // PackBuilder groups objects by `ObjectType` before encoding,
1059        // which changes the byte offsets relative to a streaming builder
1060        // that writes in added-order. The bytes of the two indices
1061        // therefore can't match exactly. What MUST match is the
1062        // **sort order of ids** — both builders ultimately call
1063        // `PackIndex::sort()` (or the bucket-equivalent), and any
1064        // reader binary-searches against that order.
1065        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        // Disable delta encoding so the classic builder produces a pack
1082        // shape comparable to the streaming one (which never deltas).
1083        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        // Same set of ids in the same sorted order — that's the
1107        // contract for binary search to work.
1108        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        // Spot-check that each id resolves to a payload that matches
1114        // the classic builder's output (equal bytes after decompression).
1115        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        // Flip one byte in the body. The trailer checksum must reject.
1135        let body_byte = 18; // past the 16-byte header
1136        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        // Header count is bytes 8..16 (big-endian).
1158        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        // Buckets exist and contain data.
1250        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            // Drop without finalize — Drop impl should clean up.
1280        }
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        // 4 MiB blob — well under the actual streaming target but big
1294        // enough to confirm we're not buffering the entire pack body in
1295        // RAM. The pack data on disk should be at least 4 MiB; the
1296        // builder's in-memory state is per-object only.
1297        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        // Pack on disk holds the whole compressed payload + headers
1323        // + trailer. Confirm it round-trips.
1324        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        // Confirms our sort-time peak memory bound. We accumulate
1335        // 1024 random hashes through a builder and check that no
1336        // single bucket holds more than ~3× the average. (BLAKE3 hash
1337        // first-byte distribution is uniform; this is mostly a
1338        // sanity check that we route to the right bucket and aren't
1339        // accidentally collapsing.)
1340        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        // Inspect bucket file sizes BEFORE finalize (which deletes them).
1357        b.pack_writer.as_mut().unwrap().flush().unwrap();
1358        let mut max_entries = 0usize;
1359        let entry_size = 33 + 8; // tagged-hash + u64 offset
1360        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        // Average is 1024 / 256 = 4 entries per bucket. Allow up to 16
1370        // (4× average) — uniformity isn't perfect on small samples.
1371        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        // Force the streaming path with a payload that compresses
1399        // well (long runs of identical bytes). Verifies:
1400        //  1. Streaming output decodes back to the original bytes.
1401        //  2. The compressed body is genuinely smaller than the
1402        //     uncompressed input (proving zstd ran), and
1403        //  3. The non-canonical 10-byte varint patched into the
1404        //     compressed_size slot decodes to the right value.
1405        let tmp = tempfile::TempDir::new().unwrap();
1406        let (mut b, _, idx_path) = fresh_builder(&tmp);
1407        // 64 KiB of zeros — compresses to a tiny zstd frame, well
1408        // above the default `min_size` so we hit the streaming branch.
1409        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        // Sanity for the seek-back scheme: for every value we'd want
1431        // to encode (small, mid, large), confirm the existing
1432        // `decode_varint` returns the same `value` from a 10-byte
1433        // padded encoding. If this ever fails the streaming path's
1434        // patched compressed_size would be misread by readers.
1435        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        // Smoke check: write a single 8 MiB payload, observe the
1453        // pack file size on disk during/after the add. The pack file
1454        // grows incrementally during the streaming compression — if
1455        // we were buffering an intermediate compressed `Vec<u8>` the
1456        // on-disk size would jump by ~8 MiB at finalize, not stay
1457        // bounded as the encoder pumps bytes through.
1458        //
1459        // We can't easily measure peak heap from inside Rust without
1460        // a custom allocator. What we *can* verify is that calling
1461        // `add` returns control with the pack file already at its
1462        // final body size, demonstrating the encoder wrote through
1463        // and didn't accumulate.
1464        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        // Pack file already on disk holds at least the entry header +
1486        // compressed payload (excluding the 32-byte trailer the builder
1487        // appends at finalize).
1488        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        // Mix of Hash and StateId in a non-sorted order on input.
1510        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        // PackReader's list_ids returns index order — should already be
1524        // sorted because we sort on finalize.
1525        let mut sorted = got.clone();
1526        sorted.sort();
1527        assert_eq!(got, sorted, "list_ids must come back sorted");
1528        // And every added id should appear.
1529        added.sort();
1530        got.sort();
1531        assert_eq!(got, added);
1532    }
1533}