znippy-compress 0.9.6

//! 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, should_skip_compression,
};
use znippy_common::meta::{BlobMeta, ChunkMeta};
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)
}

/// 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> {
    // 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, no_skip, 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,
    no_skip: bool,
    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(&registry);
        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) => {
                        let skip =
                            !no_skip && should_skip_compression(Path::new(&entry.relative_path));
                        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
        },
    })
}