znippy-compress/src/stream_packer.rs
//! Streaming compressor (`compress_stream`) on the canonical no-barrier Gatling
//! engine: [`znippy_zoomies::gatling::ordered::run_ordered_sink`].
//!
//! The caller stays the **producer**: entries arrive in memory via a channel; the
//! producer pulls one entry at a time and splits it into rounds (small entry = one
//! whole round; big entry = slice-size rounds), referencing the entry's
//! `Arc<Vec<u8>>` — zero-copy, no slot buffers needed since the data is already in
//! RAM. The engine fans rounds across N **map** workers that stamp BLAKE3 over the
//! original bytes and compress (or store raw). The **sink** receives outputs in
//! strict producer order and pwrites each at the next archive offset, recording its
//! `BlobMeta` incrementally — no full-archive RAM buffering. The finalizer groups
//! by (pkg_type, repo) into a v0.7 multi-index — arrow-ipc from metadata only.
use anyhow::{Result, anyhow};
use crossbeam_channel::{Receiver, Sender, bounded};
use std::cell::RefCell;
use std::fs::File;
use std::os::unix::fs::FileExt;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::thread;
use znippy_zoomies::gatling::ordered::{OrderedSink, run_ordered_sink};
use znippy_common::CompressionReport;
use znippy_common::codec::CompressCtx;
use znippy_common::common_config::CONFIG;
use znippy_common::index::{
build_arrow_metadata_for_config, build_metadata_batch,
compose_index_schema,
};
use znippy_common::meta::{BlobMeta, ChunkMeta};
use znippy_common::precompressed::{SNIFF_PREFIX_LEN, SkipPolicy};
use znippy_common::{ArchiveMetaSink, ArrowIpcSink, GroupKey};
/// Entries bigger than this are cut into slice-size rounds; smaller stay whole.
const SLICE_SIZE: usize = 8 * 1024 * 1024;
/// An entry to be compressed into the archive.
pub struct ArchiveEntry {
pub relative_path: String,
pub data: Vec<u8>,
/// Package type discriminator. None means "untyped / default group".
/// When all entries share the same (pkg_type, repo), the archive is written
/// in v0.6 format. Multiple distinct pairs produce a v0.7 multi-index archive.
pub pkg_type: Option<i8>,
/// Repository label for this entry. None is treated as "".
pub repo: Option<String>,
}
impl ArchiveEntry {
pub fn new(relative_path: impl Into<String>, data: Vec<u8>) -> Self {
Self { relative_path: relative_path.into(), data, pkg_type: None, repo: None }
}
}
impl Default for ArchiveEntry {
fn default() -> Self {
Self { relative_path: String::new(), data: Vec::new(), pkg_type: None, repo: None }
}
}
/// A handle to the streaming compressor.
pub struct StreamCompressor {
tx: Option<Sender<ArchiveEntry>>,
join_handle: Option<thread::JoinHandle<Result<CompressionReport>>>,
}
impl StreamCompressor {
pub fn sender(&self) -> &Sender<ArchiveEntry> {
self.tx.as_ref().expect("sender already consumed")
}
pub fn finish(mut self) -> Result<CompressionReport> {
drop(self.tx.take());
self.join_handle
.take()
.expect("already finished")
.join()
.map_err(|e| anyhow!("Compression thread panicked: {:?}", e))?
}
}
pub fn compress_stream(output: &PathBuf, no_skip: bool) -> Result<StreamCompressor> {
compress_stream_with_sink(output, no_skip, None)
}
/// [`compress_stream`] with an explicit [`SkipPolicy`] instead of a bare
/// `no_skip` flag.
///
/// The policy covers the whole stream, because that is how a caller's knowledge
/// arrives — gunnar sealing a repository's packfiles knows it for the whole seal,
/// not entry by entry. [`SkipPolicy::already_compressed`] stores every entry raw
/// without inspecting a byte.
pub fn compress_stream_with_policy(
output: &PathBuf,
policy: SkipPolicy,
) -> Result<StreamCompressor> {
compress_stream_with_sink_and_policy(output, policy, None)
}
/// Like [`compress_stream`] but lets the caller choose the metadata backend via a
/// [`MetaSinkFactory`] (e.g. an Iceberg sink). `None` uses the default inline
/// `ArrowIpcSink`. Mirrors `compress_dir`'s sink injection for the in-memory path.
pub fn compress_stream_with_sink(
output: &PathBuf,
no_skip: bool,
sink_factory: Option<znippy_common::MetaSinkFactory>,
) -> Result<StreamCompressor> {
compress_stream_with_sink_and_policy(output, SkipPolicy::from_no_skip(no_skip), sink_factory)
}
/// The one implementation behind [`compress_stream`],
/// [`compress_stream_with_sink`] and [`compress_stream_with_policy`] — both a
/// caller-chosen metadata backend and a caller-supplied [`SkipPolicy`].
pub fn compress_stream_with_sink_and_policy(
output: &PathBuf,
policy: SkipPolicy,
sink_factory: Option<znippy_common::MetaSinkFactory>,
) -> Result<StreamCompressor> {
// PERF: bounded entry channel caps in-flight `ArchiveEntry { data: Vec<u8> }`
// payloads so producers get backpressure instead of the unbounded queue
// ballooning memory when the reader/compressors fall behind.
let num_workers = CONFIG.max_core_in_flight.max(1);
let (tx_entry, rx_entry): (Sender<ArchiveEntry>, Receiver<ArchiveEntry>) =
bounded(num_workers * 4);
let output = output.clone();
let join_handle = thread::spawn(move || -> Result<CompressionReport> {
run_pipeline(rx_entry, &output, policy, sink_factory)
});
Ok(StreamCompressor { tx: Some(tx_entry), join_handle: Some(join_handle) })
}
/// Per-file metadata + byte accounting collected by the producer, consumed by the
/// finalizer. Small (path strings + per-file discriminators) — the file *data*
/// streams through the rounds and is freed incrementally, never accumulated here.
#[derive(Default)]
struct Registry {
paths: Vec<String>,
pkg_types: Vec<Option<i8>>,
repos: Vec<Option<String>>,
uf: u64,
ub: u64,
cf: u64,
cb: u64,
}
/// The gatling **label**: out-of-band metadata for one chunk (everything the sink
/// needs to place it and record its `ChunkMeta`, carried alongside the bytes).
struct ChunkLabel {
file_index: u64,
fdata_offset: u64,
chunk_seq: u32,
skip: bool,
}
/// The gatling **input** for one chunk: a zero-copy slice of the entry's `Arc` data.
struct ChunkInput {
data: Arc<Vec<u8>>,
start: usize,
len: usize,
}
/// What the sink pwrites for one chunk.
enum OutPayload {
/// Compressed output owned by the map worker.
Buf(Vec<u8>),
/// Skip / incompressible path: zero-copy straight from the entry's data (the
/// `Arc` keeps it alive until the sink has pwritten it).
Skip { data: Arc<Vec<u8>>, start: usize, len: usize },
}
/// The gatling **output** for one chunk: payload to pwrite + the `ChunkMeta` fields.
struct ChunkOut {
payload: OutPayload,
on_disk_len: usize,
file_index: u64,
fdata_offset: u64,
chunk_seq: u32,
checksum: [u8; 32],
compressed: bool,
uncompressed_size: u64,
}
thread_local! {
/// One OpenZL context + one reusable compress-scratch buffer **per map worker
/// thread** (the gatling workers are persistent, so this is created once per
/// worker and reused across every chunk it handles). Keeps the hot path
/// allocation-free — exactly what the recycled free-buffer pool did before.
static COMPRESS_TLS: RefCell<Option<(CompressCtx, Vec<u8>)>> =
const { RefCell::new(None) };
}
/// One chunk's in-flight cursor while the producer splits an entry into rounds.
struct CurEntry {
data: Arc<Vec<u8>>,
total: usize,
off: usize,
seq: u32,
skip: bool,
file_index: u64,
small: bool,
}
/// The ordered streaming sink: pwrites each in-order chunk at the next archive
/// offset and records its `BlobMeta`. Runs on the pipeline (calling) thread, so it
/// owns its mutable archive state directly — no `Send`/lock dance.
struct ArchiveSink {
file: Arc<File>,
cursor: u64,
blobs: Vec<BlobMeta>,
}
impl OrderedSink<Result<ChunkOut>> for ArchiveSink {
fn emit(&mut self, _seq: u64, output: Result<ChunkOut>) -> Result<()> {
let job = output?;
let off = self.cursor;
self.cursor += job.on_disk_len as u64;
match job.payload {
OutPayload::Buf(buf) => {
self.file.write_all_at(&buf[..job.on_disk_len], off)?;
}
OutPayload::Skip { data, start, len } => {
self.file.write_all_at(&data[start..start + len], off)?;
}
}
self.blobs.push(BlobMeta {
chunk_meta: ChunkMeta {
fdata_offset: job.fdata_offset,
file_index: job.file_index,
chunk_seq: job.chunk_seq,
checksum: job.checksum,
compressed: job.compressed,
uncompressed_size: job.uncompressed_size,
compressed_size: job.on_disk_len as u64,
},
blob_offset: off,
blob_size: job.on_disk_len as u64,
});
Ok(())
}
}
fn run_pipeline(
rx_entry: Receiver<ArchiveEntry>,
output: &PathBuf,
policy: SkipPolicy,
sink_factory: Option<znippy_common::MetaSinkFactory>,
) -> Result<CompressionReport> {
let output_path = output.with_extension("znippy");
let file = Arc::new(File::create(&output_path)?);
let num_workers = CONFIG.max_core_in_flight.max(1);
// In-flight / reorder-buffer bound — matches the old bounded-channel depth so
// the producer is back-pressured the same way (memory bounded by `cap`, not by
// stream length). The entry channel above already bounds the upstream sender.
let cap = num_workers * 4;
let level = CONFIG.compression_level;
// Shared per-file registry: the producer appends one record per entry (a single
// coarse-grained lock per entry, not per round); the finalizer reads it back
// after the engine has joined every thread.
let registry: Arc<Mutex<Registry>> = Arc::new(Mutex::new(Registry::default()));
// ── PRODUCER: pull entries, split into rounds, yield (label, bytes) lazily ──
let producer = {
let registry = Arc::clone(®istry);
let mut cur: Option<CurEntry> = None;
move || -> Option<(ChunkLabel, ChunkInput)> {
loop {
if let Some(c) = cur.as_mut() {
if c.off < c.total {
let len = if c.small { c.total } else { SLICE_SIZE.min(c.total - c.off) };
let label = ChunkLabel {
file_index: c.file_index,
fdata_offset: c.off as u64,
chunk_seq: c.seq,
skip: c.skip,
};
let input = ChunkInput { data: Arc::clone(&c.data), start: c.off, len };
c.off += len;
c.seq += 1;
return Some((label, input));
}
cur = None;
}
match rx_entry.recv() {
Ok(entry) => {
// Path fast path first (it covers the whole entry, every
// chunk of it), then a magic-byte probe over the head of
// the real bytes — which is the only thing that can reach
// an entry carrying no extension, or one whose name lies.
let path = Path::new(&entry.relative_path);
let skip = policy.skip_by_path(path)
|| policy.skip_by_bytes(
&entry.data[..entry.data.len().min(SNIFF_PREFIX_LEN)],
);
let data_len = entry.data.len() as u64;
let file_index;
{
let mut reg = registry.lock().expect("registry lock");
file_index = reg.paths.len() as u64;
if skip {
reg.uf += 1;
reg.ub += data_len;
} else {
reg.cf += 1;
reg.cb += data_len;
}
reg.pkg_types.push(entry.pkg_type);
reg.repos.push(entry.repo);
reg.paths.push(entry.relative_path);
}
let data = Arc::new(entry.data);
let total = data.len();
if total == 0 {
// Empty entry → one zero-length round so it appears in
// the index. (No further rounds for this entry.)
return Some((
ChunkLabel { file_index, fdata_offset: 0, chunk_seq: 0, skip },
ChunkInput { data, start: 0, len: 0 },
));
}
let small = total <= SLICE_SIZE;
cur = Some(CurEntry { data, total, off: 0, seq: 0, skip, file_index, small });
// Loop to emit this entry's first round.
}
Err(_) => return None, // sender hung up: end of stream
}
}
}
};
// ── MAP: BLAKE3 over the ORIGINAL bytes, then compress (or store raw) ────────
let map = move |label: ChunkLabel, input: ChunkInput| -> Result<ChunkOut> {
let src = &input.data[input.start..input.start + input.len];
let checksum = *blake3::hash(src).as_bytes(); // ORIGINAL bytes, pre-compression
let uncompressed_size = input.len as u64;
let (payload, on_disk_len, compressed) = if label.skip {
(
OutPayload::Skip {
data: Arc::clone(&input.data),
start: input.start,
len: input.len,
},
input.len,
false,
)
} else {
COMPRESS_TLS.with(|cell| -> Result<(OutPayload, usize, bool)> {
let mut guard = cell.borrow_mut();
if guard.is_none() {
*guard = Some((CompressCtx::new(level)?, Vec::new()));
}
let (cctx, scratch) = guard.as_mut().unwrap();
let n = cctx.compress_into(src, scratch)?;
if n >= input.len {
// Incompressible: storing raw is no bigger and skips the decode
// cost. Zero-copy straight from the entry's `Arc` — the scratch
// buffer keeps its capacity for the next chunk.
Ok((
OutPayload::Skip {
data: Arc::clone(&input.data),
start: input.start,
len: input.len,
},
input.len,
false,
))
} else {
// ZERO-ALLOC: hand the already-filled scratch buffer to the
// sink via `take` — `compress_into` truncated it to exactly `n`,
// so no alloc-a-new-buffer-and-memcpy. The next chunk on this
// worker gets a fresh scratch (the CompressCtx, the costly part,
// is retained in TLS).
Ok((OutPayload::Buf(std::mem::take(scratch)), n, true))
}
})?
};
Ok(ChunkOut {
payload,
on_disk_len,
file_index: label.file_index,
fdata_offset: label.fdata_offset,
chunk_seq: label.chunk_seq,
checksum,
compressed,
uncompressed_size,
})
};
// ── SINK: ordered, streaming pwrite + incremental BlobMeta ───────────────────
let mut sink = ArchiveSink { file: Arc::clone(&file), cursor: 0, blobs: Vec::new() };
let sink_result = run_ordered_sink(producer, num_workers, cap, map, &mut sink);
// test-matrix emit: HONEST verdict of the streaming big-file ordered-sink run.
#[cfg(feature = "testmatrix")]
crate::functional_status(
"znippy-compress/stream_packer",
"run_ordered_sink",
sink_result.is_ok(),
&format!(
"workers={num_workers} cap={cap} blobs={} ok={}",
sink.blobs.len(),
sink_result.is_ok()
),
);
sink_result?;
let mut all_blobs = sink.blobs;
let blob_bytes = sink.cursor; // blob region starts at 0
all_blobs.sort_by_key(|b| (b.chunk_meta.file_index, b.chunk_meta.chunk_seq));
let total_chunks = all_blobs.len() as u64;
// Recover the registry (the producer has been joined, so we are the sole owner).
let reg = std::mem::take(&mut *registry.lock().expect("registry lock"));
let (uf, ub, cf, cb) = (reg.uf, reg.ub, reg.cf, reg.cb);
// ── FINALIZER: group by (pkg_type, repo) into a v0.7 multi-index ─────────────
let file_keys: Vec<(i8, String)> = reg
.pkg_types
.iter()
.zip(reg.repos.iter())
.map(|(p, r)| (p.unwrap_or(0), r.clone().unwrap_or_default()))
.collect();
let mut groups: std::collections::BTreeMap<(i8, String), Vec<usize>> =
std::collections::BTreeMap::new();
for (i, blob) in all_blobs.iter().enumerate() {
let key = file_keys[blob.chunk_meta.file_index as usize].clone();
groups.entry(key).or_default().push(i);
}
let meta_map = build_arrow_metadata_for_config(&CONFIG);
let mut sink: Box<dyn ArchiveMetaSink> = match sink_factory {
Some(make) => make(Arc::clone(&file), blob_bytes),
None => Box::new(ArrowIpcSink::new(Arc::clone(&file), blob_bytes)),
};
let group_count = groups.len();
for ((pkg_type, repo), blob_indices) in &groups {
let group_blobs: Vec<_> = blob_indices.iter().map(|&i| all_blobs[i].clone()).collect();
let batch = build_metadata_batch(&group_blobs, |fi| reg.paths[fi as usize].clone(), &[], &[])
.map_err(|e| anyhow!("build sub-index batch: {e}"))?;
let schema_with_meta = arrow::datatypes::Schema::new_with_metadata(
compose_index_schema(&[]).fields().to_vec(),
meta_map.clone(),
);
sink.push_subindex(
&schema_with_meta,
std::slice::from_ref(&batch),
GroupKey {
pkg_type: *pkg_type,
repo: repo.clone(),
module_name: String::new(),
},
)?;
}
let manifest_offset = blob_bytes; // logging only; sink owns real placement
let total_bytes_out = sink.finish()?;
let total_files = uf + cf;
log::info!(
"[stream] gatling archive: {} group(s), {} blob bytes, manifest at {}",
group_count,
blob_bytes,
manifest_offset
);
Ok(CompressionReport {
total_files,
compressed_files: cf,
uncompressed_files: uf,
chunks: total_chunks,
total_dirs: 0,
total_bytes_in: cb + ub,
total_bytes_out,
compressed_bytes: cb,
uncompressed_bytes: ub,
compression_ratio: if cb > 0 && total_bytes_out > ub {
(cb as f32 / (total_bytes_out - ub) as f32) * 100.0
} else {
0.0
},
// The stream API counts uf/cf only for entries actually received with
// their bytes in hand (the caller has already read the file), and sets
// `total_files = uf + cf`, so no enumerated-but-dropped files exist here.
files_failed: 0,
})
}