lgz/src/stream.rs
//! **One-pass streaming gzip decode**: `Read → Write`, chunk by chunk.
//!
//! ```text
//! reader thread ──chunk──▶ split (all cores at once, i × chunk / cores, first
//! border each: flush point or bit-level block start)
//! ▶ decode (all cores; markers for the unknown window)
//! ▶ in order: false starts, windows, markers resolved
//! on all cores, CRC pieces combined per member
//! ──bytes──▶ writer thread
//! ```
//!
//! The input is never read whole: one chunk is being read while the previous
//! one is decoded and the one before that is written. What carries from one
//! chunk to the next is the undecoded tail (from the exact bit where the last
//! decoded segment ended), the last 32 KB of output, and the member state
//! (inside DEFLATE, at a trailer, at a header). Multi-member gzip (`cat a.gz
//! b.gz`, pigz, bgzf) streams the same way: each member's trailer is checked
//! where it ends.
//!
//! Small inputs (below [`SERIAL_BELOW`]) and one thread take plain streaming
//! zlib-rs (flate2), so they are never slower than it.
//!
//! An error says how many bytes were already written: those bytes are not
//! verified, since a later trailer may still fail.
use std::io::{Read, Write};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc;
use crate::bitsplit::{self, MarkedSeg, SegEnd, WINDOW};
/// Compressed bytes read per chunk.
pub const DEFAULT_CHUNK: usize = 32 << 20;
/// Inputs smaller than this (compressed bytes) decode on one thread with
/// flate2: below it the parallel start-up costs more than it gains.
pub const SERIAL_BELOW: usize = 12 << 20;
/// Counters from a [`decode`] run.
#[derive(Debug, Default, Clone, Copy)]
pub struct Stats {
/// Segments decoded.
pub segments: usize,
/// Segments that needed the cross-border marker resolution.
pub marker_segments: usize,
/// Chunks read.
pub chunks: usize,
/// Segments re-decoded in order because their start was false.
pub redone: usize,
/// Segments dropped because the previous one already covered them.
pub skipped: usize,
/// gzip members whose trailer was checked.
pub members: usize,
/// Symbols decoded by the marker decoder.
pub head_symbols: usize,
/// The input took the one-thread flate2 path.
pub serial: bool,
/// Bytes written.
pub written: u64,
}
/// A failed decode, and how many (unverified) bytes were written before it.
#[derive(Debug)]
pub struct StreamError {
pub msg: &'static str,
pub written: u64,
}
impl std::fmt::Display for StreamError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"lgz: {} ({} bytes were already written and are not verified)",
self.msg, self.written
)
}
}
impl std::error::Error for StreamError {}
struct CountingWriter<W> {
w: W,
n: Arc<AtomicU64>,
}
impl<W: Write> Write for CountingWriter<W> {
fn write(&mut self, b: &[u8]) -> std::io::Result<usize> {
let k = self.w.write(b)?;
self.n.fetch_add(k as u64, Ordering::Relaxed);
Ok(k)
}
fn flush(&mut self) -> std::io::Result<()> {
self.w.flush()
}
}
/// Fill `buf` (cleared, capacity `want`) with up to `want` bytes; true = EOF.
fn read_full<R: Read>(r: &mut R, buf: &mut Vec<u8>, want: usize) -> std::io::Result<bool> {
buf.clear();
buf.reserve(want);
// SAFETY: bytes past `got` are never read; the length is set back to
// what `read` actually filled.
#[allow(clippy::uninit_vec)]
unsafe {
buf.set_len(want)
};
let mut got = 0;
let res = loop {
if got == want {
break Ok(false);
}
match r.read(&mut buf[got..]) {
Ok(0) => break Ok(true),
Ok(k) => got += k,
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
Err(e) => break Err(e),
}
};
buf.truncate(got);
res
}
/// The last `WINDOW` bytes of `prev ++ a ++ b`.
fn next_window(prev: &[u8], a: &[u8], b: &[u8]) -> Vec<u8> {
let mut w = Vec::with_capacity(WINDOW);
let need_b = b.len().min(WINDOW);
let need_a = a.len().min(WINDOW - need_b);
let need_p = prev.len().min(WINDOW - need_b - need_a);
w.extend_from_slice(&prev[prev.len() - need_p..]);
w.extend_from_slice(&a[a.len() - need_a..]);
w.extend_from_slice(&b[b.len() - need_b..]);
w
}
/// Where the carried input stands.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Mode {
/// at a gzip header (or the end of the stream)
Header,
/// inside DEFLATE; `fresh` = at a member's first block (no window)
Block { fresh: bool },
/// at a member's 8-byte trailer
Trailer,
}
/// Running CRC32 + length of the current member.
struct Member {
crc: flate2::Crc,
len: u64,
}
impl Member {
fn feed(&mut self, c: &flate2::Crc) {
if c.amount() > 0 {
self.crc.combine(c);
self.len += c.amount() as u64;
}
}
fn close(&mut self, crc: u32, isize: u32, members: &mut usize) -> Result<(), &'static str> {
if self.crc.sum() != crc || self.len as u32 != isize {
return Err("gzip checksum mismatch");
}
*self = Member {
crc: flate2::Crc::new(),
len: 0,
};
*members += 1;
Ok(())
}
}
/// One chunk of compressed input, shared by its jobs.
struct ChunkData {
/// Absolute byte offset of `bytes[0]` in the input.
base: usize,
bytes: Data,
is_last: bool,
/// Absolute bit starts of this chunk's segments (job starts), ascending.
starts: Vec<usize>,
/// Absolute bit where a decode that ran past every start ends (at the
/// first block end at or after it): the window's end on a mapped input,
/// so a false last start never runs a worker on to the end of the input.
soft: usize,
}
/// A chunk's compressed bytes: read into its own buffer (a pipe), or the
/// whole mapped input, read in place (zero copy).
enum Data {
Owned(Vec<u8>),
/// SAFETY: points into the input slice given to `decode_slice`, which
/// outlives every chunk (all are dropped before it returns).
Mapped(*const u8, usize),
}
// SAFETY: `Mapped` is a read-only view of a slice that outlives the decode.
unsafe impl Send for Data {}
unsafe impl Sync for Data {}
impl std::ops::Deref for Data {
type Target = [u8];
fn deref(&self) -> &[u8] {
match self {
Data::Owned(v) => v,
// SAFETY: see `Mapped`
Data::Mapped(p, n) => unsafe { std::slice::from_raw_parts(*p, *n) },
}
}
}
impl Data {
fn owned_mut(&mut self) -> &mut Vec<u8> {
match self {
Data::Owned(v) => v,
Data::Mapped(..) => unreachable!("a mapped chunk is never extended"),
}
}
}
impl ChunkData {
fn rel(&self, abs_bit: usize) -> usize {
abs_bit - self.base * 8
}
fn covers(&self, abs_bit: usize) -> bool {
abs_bit >= self.base * 8 && abs_bit <= (self.base + self.bytes.len()) * 8
}
/// Later job starts after `abs_bit`, relative to this chunk.
fn stops_after(&self, abs_bit: usize) -> Vec<usize> {
self.starts
.iter()
.filter(|&&s| s > abs_bit)
.map(|&s| self.rel(s))
.collect()
}
}
/// One segment to decode.
struct Job {
c: Arc<ChunkData>,
i: usize,
fresh: bool,
}
type Done = (Arc<ChunkData>, usize, Result<MarkedSeg, &'static str>);
/// Decode the gzip stream `r` into `w` on `threads` workers, `chunk`
/// compressed bytes at a time. Every member's trailer is checked.
///
/// One gatling ordered pipeline: the producer reads a chunk and splits it (all
/// cores at once from i × chunk / cores, first border each) and hands every
/// segment on as a job; workers decode jobs as they come, across chunk
/// boundaries with no barrier; the in-order sink checks that each segment
/// starts where the previous one ended (re-decoding past a false start), resolves
/// window markers, checks member CRCs and writes.
pub fn decode<R: Read + Send, W: Write + Send>(
r: R,
w: W,
threads: usize,
chunk: usize,
) -> Result<Stats, StreamError> {
let chunk = chunk.max(1 << 20) as u64;
decode_windows(r, w, threads, std::iter::repeat(chunk), None)
}
/// lgz's read-window plan: segments of 1 to 50 MB, about 4 per core over the
/// input ([`gatling::window_plan`]).
pub const PLAN: gatling::window_plan::Plan = gatling::window_plan::Plan {
min_seg: 1 << 20,
max_seg: 50 << 20,
k: 4,
};
/// [`PLAN`], with `LGZ_K` and `LGZ_MIN_SEG_KB` overrides for benching.
pub fn plan() -> gatling::window_plan::Plan {
let env = |k: &str| std::env::var(k).ok().and_then(|v| v.parse::<u64>().ok());
gatling::window_plan::Plan {
k: env("LGZ_K").unwrap_or(PLAN.k),
min_seg: env("LGZ_MIN_SEG_KB")
.map(|v| v << 10)
.unwrap_or(PLAN.min_seg),
..PLAN
}
}
/// [`decode`] with read windows from the shared plan: `size` is the input's
/// size if known (a file, an OCI layer), `None` on a pipe.
pub fn decode_sized<R: Read + Send, W: Write + Send>(
r: R,
w: W,
threads: usize,
size: Option<u64>,
) -> Result<Stats, StreamError> {
let plan = plan();
let wins = gatling::window_plan::windows(size, threads, plan);
// past a known size, keep reading in steady windows (a size that lied)
let steady = plan.read_window(size, threads);
decode_windows(r, w, threads, wins.chain(std::iter::repeat(steady)), size)
}
fn decode_windows<R: Read + Send, W: Write + Send>(
mut r: R,
w: W,
threads: usize,
mut wins: impl Iterator<Item = u64> + Send,
size: Option<u64>,
) -> Result<Stats, StreamError> {
let written = Arc::new(AtomicU64::new(0));
let fail = |msg: &'static str, written: &Arc<AtomicU64>| StreamError {
msg,
written: written.load(Ordering::Relaxed),
};
if threads <= 1 || size.is_some_and(|s| s < SERIAL_BELOW as u64) {
return serial(Vec::new(), r, w, written);
}
let mut first = Vec::new();
let first_want = wins.next().unwrap_or(1 << 20) as usize;
let eof =
read_full(&mut r, &mut first, first_want).map_err(|_| fail("read failed", &written))?;
if eof && first.len() < SERIAL_BELOW {
return serial(first, r, w, written);
}
let threads = threads.max(2);
let mut sink = Sink {
w: CountingWriter {
w,
n: written.clone(),
},
at: 0,
mode: Mode::Header,
window: Vec::new(),
member: Member {
crc: flate2::Crc::new(),
len: 0,
},
cur: None,
prev: None,
end: None,
lag: None,
st: Stats::default(),
err: None,
released: 0,
release: None,
pos_out: None,
};
let mut prod_err: Option<&'static str> = None;
let res = std::thread::scope(|sc| {
// reader thread: the next read overlaps the decode
let (ctx, crx) = mpsc::sync_channel::<(Vec<u8>, bool)>(0);
let reader = (!eof).then(|| {
sc.spawn(move || {
loop {
let mut buf = Vec::new();
let want = wins.next().unwrap_or(1 << 20) as usize;
let eof = match read_full(&mut r, &mut buf, want) {
Ok(e) => e,
Err(_) => return Err(()),
};
if ctx.send((buf, eof)).is_err() || eof {
return Ok(());
}
}
})
});
let mut pending: std::collections::VecDeque<Job> = std::collections::VecDeque::new();
let mut next_read = Some((first, eof));
let mut carry: Vec<u8> = Vec::new();
let mut carry_base = 0usize;
let mut start: Option<(usize, bool)> = None; // (abs bit, fresh)
let mut finished = false;
let pe = &mut prod_err;
let producer = move || -> Option<((), Job)> {
loop {
if let Some(j) = pending.pop_front() {
return Some(((), j));
}
if finished {
return None;
}
let (buf, is_last) = match next_read.take() {
Some(x) => x,
None => match crx.recv() {
Ok(x) => x,
Err(_) => {
*pe = Some("read failed");
finished = true;
return None;
}
},
};
let base = carry_base;
let mut bytes = std::mem::take(&mut carry);
bytes.extend_from_slice(&buf);
drop(buf);
let (s0, fresh) = match start {
Some(x) => x,
None => match bitsplit::gzip_header_len(&bytes) {
Ok(Some(n)) => ((base + n) * 8, true),
Ok(None) if !is_last => {
carry = bytes;
continue;
}
_ => {
*pe = Some("not a gzip stream");
finished = true;
return None;
}
},
};
let rel: Vec<usize> =
bitsplit::split_chunk_bits(&bytes, s0 - base * 8, threads).starts;
let starts: Vec<usize> = rel.iter().map(|&b| b + base * 8).collect();
// the last segment of a mid-stream chunk is the carry
let k = if is_last {
starts.len()
} else {
starts.len() - 1
};
if k == 0 {
carry = bytes; // no border yet: read more
start = Some((s0, fresh));
continue;
}
// the carry keeps one segment more than needed, so the sink can
// re-decode from the previous start if the carry start is false
let keep_from = starts[k - 1] / 8;
if !is_last {
carry = bytes[keep_from - base..].to_vec();
carry_base = keep_from;
start = Some((starts[k], false));
}
let c = Arc::new(ChunkData {
base,
bytes: Data::Owned(bytes),
is_last,
starts,
soft: usize::MAX,
});
for i in 0..k {
pending.push_back(Job {
c: c.clone(),
i,
fresh: fresh && i == 0,
});
}
if is_last {
finished = true;
}
}
};
let map = |_: (), j: Job| -> Done {
let s = j.c.starts[j.i];
let stops = j.c.stops_after(s);
let soft = if j.c.soft == usize::MAX {
usize::MAX
} else {
j.c.rel(j.c.soft)
};
let r = bitsplit::decode_from(
&j.c.bytes,
j.c.rel(s),
&stops,
if j.fresh { Some(&[]) } else { None },
soft,
)
.map(|mut seg| {
seg.end_bit += j.c.base * 8;
seg
})
.map_err(|e| e.0);
(j.c, j.i, r)
};
let res = gatling::gatling::ordered::run_ordered_sink(
producer,
threads,
threads + 2,
map,
&mut sink,
);
let rr = reader
.map(|h| h.join().unwrap_or(Err(())))
.unwrap_or(Ok(()));
(res, rr)
});
let (res, rr) = res;
if let Some(m) = sink.err.or(prod_err) {
return Err(fail(m, &written));
}
res.map_err(|_| fail("pipeline failed", &written))?;
rr.map_err(|_| fail("read failed", &written))?;
sink.finish().map_err(|m| fail(m, &written))?;
sink.w.flush().map_err(|_| fail("write failed", &written))?;
let mut st = sink.st;
st.written = written.load(Ordering::Relaxed);
Ok(st)
}
/// The in-order stage.
struct Sink<'r, W: Write> {
w: CountingWriter<W>,
/// Absolute bit where the next byte of output comes from.
at: usize,
mode: Mode,
window: Vec<u8>,
member: Member,
/// The latest two chunks seen (a re-decode may need either).
cur: Option<Arc<ChunkData>>,
prev: Option<Arc<ChunkData>>,
/// Absolute byte length of the whole input, once the last chunk is seen.
end: Option<usize>,
/// Contiguous input from `at` on, kept while a re-decode waits for more
/// bytes (several split starts in a row were false): each new chunk's
/// bytes are appended until the re-decode reaches a real start.
lag: Option<ChunkData>,
st: Stats,
err: Option<&'static str>,
/// Input bytes already handed to `release`.
released: usize,
release: Option<&'r (dyn Fn(usize) + Sync)>,
/// Output to a regular file: each segment goes to a writer pool at its
/// known offset (pwrite), so no single thread copies the whole output.
pos_out: Option<PosOut>,
}
/// Positional output: the next segment's file offset, and the writer pool.
struct PosOut {
off: u64,
tx: mpsc::SyncSender<(u64, Vec<u8>, MarkedSeg)>,
n: Arc<AtomicU64>,
}
enum Step {
Progress,
/// the bytes needed are not in any chunk seen yet
Wait,
}
impl<W: Write> Sink<'_, W> {
fn note(&mut self, c: &Arc<ChunkData>) {
let new = self.cur.as_ref().is_none_or(|x| !Arc::ptr_eq(x, c));
if new {
self.st.chunks += 1;
self.prev = self.cur.replace(c.clone());
if c.is_last {
self.end = Some(c.base + c.bytes.len());
}
if let Some(l) = &mut self.lag {
// chunks overlap (each starts inside the previous one): append
// only the part past what the lag already holds
let lend = l.base + l.bytes.len();
let cend = c.base + c.bytes.len();
if c.base <= lend && cend > lend {
l.bytes
.owned_mut()
.extend_from_slice(&c.bytes[lend - c.base..]);
}
let last = l.starts.last().copied().unwrap_or(0);
l.starts.extend(c.starts.iter().filter(|&&x| x > last));
l.is_last = c.is_last;
}
}
}
/// The seen chunk whose bytes hold `abs_bit`, latest first.
fn data_for(&self, abs_bit: usize) -> Option<Arc<ChunkData>> {
[&self.cur, &self.prev]
.into_iter()
.flatten()
.find(|c| c.covers(abs_bit))
.cloned()
}
fn accept(&mut self, c: &ChunkData, seg: MarkedSeg) -> Result<(), &'static str> {
let _ = c;
let mut h = bitsplit::pool::take_u8(seg.head.len());
h.resize(seg.head.len(), 0);
bitsplit::resolve_into(&seg.head, &self.window, &mut h).map_err(|e| e.0)?;
// head pieces and marks interleave, then the tail's
let mut from = 0;
for m in &seg.marks[..seg.head_marks] {
let mut p = flate2::Crc::new();
p.update(&h[from..m.at]);
self.member.feed(&p);
self.member.close(m.crc, m.isize, &mut self.st.members)?;
from = m.at;
}
let mut p = flate2::Crc::new();
p.update(&h[from..]);
self.member.feed(&p);
for (k, piece) in seg.tail_crcs.iter().enumerate() {
self.member.feed(piece);
if let Some(m) = seg.marks.get(seg.head_marks + k) {
self.member.close(m.crc, m.isize, &mut self.st.members)?;
}
}
let t = seg.tail_bytes();
if t.len() >= WINDOW {
// the window is reused in place
self.window.clear();
self.window.extend_from_slice(&t[t.len() - WINDOW..]);
} else {
self.window = next_window(&self.window, &h, t);
}
let (hl, tl) = (h.len(), t.len());
let end_bit = seg.end_bit;
let end = seg.end;
let had = seg.had_markers;
let head_syms = seg.head.len();
match &mut self.pos_out {
Some(p) => {
let off = p.off;
p.off += (hl + tl) as u64;
p.n.fetch_add((hl + tl) as u64, Ordering::Relaxed);
p.tx.send((off, h, seg)).map_err(|_| "write failed")?;
}
None => {
self.w.write_all(&h).map_err(|_| "write failed")?;
self.w
.write_all(seg.tail_bytes())
.map_err(|_| "write failed")?;
bitsplit::pool::give_u8(h);
bitsplit::pool::recycle(seg);
}
}
let seg = MarkedSeg {
end_bit,
end,
had_markers: had,
..MarkedSeg::default()
};
self.st.segments += 1;
self.st.head_symbols += head_syms;
if seg.had_markers {
self.st.marker_segments += 1;
}
self.at = seg.end_bit;
if let Some(rel) = self.release {
// input before `at` is never read again: give back whole 8 MB steps
let done = (self.at / 8) & !((8 << 20) - 1);
if done > self.released {
rel(done);
self.released = done;
}
}
self.mode = match seg.end {
SegEnd::Stop => Mode::Block { fresh: false },
SegEnd::Trailer => Mode::Trailer,
SegEnd::Header => Mode::Header,
};
Ok(())
}
/// Move `at` forward by itself: read a trailer or header, or re-decode
/// from `at` to the next job start.
fn advance(&mut self) -> Result<Step, &'static str> {
// the bytes from `at` on: the lag if it holds them, else a seen chunk
let lag = self.lag.take().filter(|l| l.covers(self.at));
let arc = if lag.is_none() {
match self.data_for(self.at) {
Some(x) => Some(x),
None => return Ok(Step::Wait),
}
} else {
None
};
let c: &ChunkData = lag.as_ref().unwrap_or_else(|| arc.as_ref().unwrap());
let r = self.advance_in(c);
// still waiting for bytes inside DEFLATE: keep (or start) the lag
let keep =
matches!(r, Ok(Step::Wait)) && matches!(self.mode, Mode::Block { .. }) && !c.is_last;
let started = (keep && lag.is_none()).then(|| {
let from = self.at / 8;
ChunkData {
base: from,
bytes: Data::Owned(c.bytes[from - c.base..].to_vec()),
soft: c.soft,
is_last: c.is_last,
starts: c.starts.iter().copied().filter(|&x| x > self.at).collect(),
}
});
if keep {
self.lag = lag.or(started);
}
r
}
fn advance_in(&mut self, c: &ChunkData) -> Result<Step, &'static str> {
let a = self.at / 8;
let end_here = c.base + c.bytes.len();
match self.mode {
Mode::Trailer => {
if a + 8 > end_here {
return if c.is_last {
Err("truncated gzip trailer")
} else {
Ok(Step::Wait)
};
}
let t = &c.bytes[a - c.base..a - c.base + 8];
let crc = u32::from_le_bytes(t[..4].try_into().unwrap());
let isz = u32::from_le_bytes(t[4..].try_into().unwrap());
self.member.close(crc, isz, &mut self.st.members)?;
self.at = (a + 8) * 8;
self.mode = Mode::Header;
}
Mode::Header => {
if Some(a) == self.end {
return Ok(Step::Wait); // the end of the stream
}
match bitsplit::gzip_header_len(&c.bytes[a - c.base..]).map_err(|e| e.0)? {
Some(n) => {
self.at = (a + n) * 8;
self.mode = Mode::Block { fresh: true };
}
None if c.is_last => return Err("truncated gzip header"),
None => return Ok(Step::Wait),
}
}
Mode::Block { fresh } => {
let stops = c.stops_after(self.at);
// the window before `at` is known here: no markers needed
let known: &[u8] = if fresh { &[] } else { &self.window };
match bitsplit::decode_from(
&c.bytes,
c.rel(self.at),
&stops,
Some(known),
if c.soft != usize::MAX && c.soft > self.at {
c.rel(c.soft)
} else {
usize::MAX
},
) {
Ok(mut seg) => {
seg.end_bit += c.base * 8;
self.st.redone += 1;
self.accept(c, seg)?;
}
// ran out of these bytes: wait for the next chunk
Err(_) if !c.is_last => return Ok(Step::Wait),
Err(e) => return Err(e.0),
}
}
}
Ok(Step::Progress)
}
fn take(
&mut self,
c: Arc<ChunkData>,
i: usize,
r: Result<MarkedSeg, &'static str>,
) -> Result<(), &'static str> {
self.note(&c);
let s = c.starts[i];
if std::env::var_os("LGZ_TRACE").is_some() {
eprintln!(
"sink: job base {} i {i} start {s} at {} mode {:?} ok {} end {:?}",
c.base,
self.at,
self.mode,
r.is_ok(),
r.as_ref().ok().map(|x| (x.end_bit, x.end))
);
}
let mut r = Some(r);
loop {
match self.mode {
Mode::Block { .. } if s < self.at => {
self.st.skipped += 1; // covered by an overshoot
return Ok(());
}
Mode::Block { .. } if s == self.at => match r.take().unwrap() {
Ok(seg) => return self.accept(&c, seg),
Err(_) => {
if let Step::Wait = self.advance()? {
return Ok(());
}
}
},
_ => {
// `at` is not this job's start: catch up by itself
if let Step::Wait = self.advance()? {
return Ok(());
}
}
}
}
}
fn finish(&mut self) -> Result<(), &'static str> {
// whatever is left after the last job: trailers, headers, a re-decode
loop {
if self.mode == Mode::Header && Some(self.at / 8) == self.end {
return if self.st.members == 0 {
Err("empty input")
} else {
Ok(())
};
}
if let Step::Wait = self.advance()? {
return Err("stream ended inside a gzip member");
}
}
}
}
impl<W: Write> gatling::gatling::ordered::OrderedSink<Done> for Sink<'_, W> {
fn emit(&mut self, _seq: u64, (c, i, r): Done) -> anyhow::Result<()> {
if self.err.is_some() {
return Err(anyhow::anyhow!("lgz"));
}
if let Err(e) = self.take(c, i, r) {
self.err = Some(e);
return Err(anyhow::anyhow!(e));
}
Ok(())
}
}
/// Output buffer of the one-thread slice path (plus the window in front).
const SLICE_OUT: usize = 32 << 20;
/// Decode a gzip input that is all in memory (an mmap in the CLI) into `w`.
///
/// One thread or a small input: **zero copy, zero allocation per window** —
/// the compressed bytes are read in place by bit offset, and one output
/// buffer, made once, is reused for every window (only the 32 KB window moves
/// to its front). More threads: the streaming pipeline over the slice.
pub fn decode_slice<W: Write + Send>(
input: &[u8],
w: W,
threads: usize,
) -> Result<Stats, StreamError> {
decode_slice_release(input, w, threads, &|_| {})
}
/// [`decode_slice`], calling `release(n)` as the first `n` input bytes are no
/// longer needed (the CLI drops those mapped pages: MADV_DONTNEED).
pub fn decode_slice_release<W: Write + Send>(
input: &[u8],
w: W,
threads: usize,
release: &(dyn Fn(usize) + Sync),
) -> Result<Stats, StreamError> {
if threads > 1 && input.len() >= SERIAL_BELOW {
return parallel_slice(input, w, threads, release, None);
}
serial_slice(input, w)
}
/// [`decode_slice_release`] into a regular file: segments are written at
/// their offsets by a small writer pool (pwrite), not copied out by one thread.
pub fn decode_slice_to_file(
input: &[u8],
file: &std::fs::File,
threads: usize,
release: &(dyn Fn(usize) + Sync),
) -> Result<Stats, StreamError> {
if threads > 1 && input.len() >= SERIAL_BELOW {
return parallel_slice(input, std::io::sink(), threads, release, Some(file));
}
serial_slice(input, file)
}
/// The pipeline over an in-memory input, zero copy: every chunk is the whole
/// input seen in place; a window is only the range its split starts fall in.
fn parallel_slice<W: Write + Send>(
input: &[u8],
w: W,
threads: usize,
release: &(dyn Fn(usize) + Sync),
file: Option<&std::fs::File>,
) -> Result<Stats, StreamError> {
let written = Arc::new(AtomicU64::new(0));
let fail = |msg: &'static str| StreamError {
msg,
written: written.load(Ordering::Relaxed),
};
let mut sink = Sink {
w: CountingWriter {
w,
n: written.clone(),
},
at: 0,
mode: Mode::Header,
window: Vec::new(),
member: Member {
crc: flate2::Crc::new(),
len: 0,
},
cur: None,
prev: None,
end: None,
lag: None,
st: Stats::default(),
err: None,
released: 0,
release: Some(release),
pos_out: None,
};
let s0 = match bitsplit::gzip_header_len(input).map_err(|e| fail(e.0))? {
Some(n) => n * 8,
None => return Err(fail("truncated gzip header")),
};
let plan = plan();
let mut wins = gatling::window_plan::windows(Some(input.len() as u64), threads, plan);
let mut pending: std::collections::VecDeque<Job> = std::collections::VecDeque::new();
let mut start = (s0, true);
let mut wend = 0usize;
let mut finished = false;
let producer = move || -> Option<((), Job)> {
loop {
if let Some(j) = pending.pop_front() {
return Some(((), j));
}
if finished {
return None;
}
wend = (wend + wins.next().unwrap_or(1 << 20) as usize).min(input.len());
let is_last = wend == input.len();
let (s, fresh) = start;
if s >= wend * 8 && !is_last {
continue; // this window is all behind the next start
}
let starts = bitsplit::split_chunk_bits(&input[..wend], s, threads).starts;
let k = if is_last {
starts.len()
} else {
starts.len() - 1
};
if k == 0 {
continue; // no border in this window: take the next one too
}
if !is_last {
start = (starts[k], false);
}
let c = Arc::new(ChunkData {
base: 0,
bytes: Data::Mapped(input.as_ptr(), input.len()),
is_last,
starts,
soft: if is_last { usize::MAX } else { wend * 8 },
});
for i in 0..k {
pending.push_back(Job {
c: c.clone(),
i,
fresh: fresh && i == 0,
});
}
if is_last {
finished = true;
}
}
};
let map = |_: (), j: Job| -> Done {
let s = j.c.starts[j.i];
let stops = j.c.stops_after(s);
let soft = if j.c.soft == usize::MAX {
usize::MAX
} else {
j.c.rel(j.c.soft)
};
let known: Option<&[u8]> = if j.fresh { Some(&[]) } else { None };
let r = bitsplit::decode_from(&j.c.bytes, j.c.rel(s), &stops, known, soft).map_err(|e| e.0);
(j.c, j.i, r)
};
let (res, wres) = std::thread::scope(|sc| {
let mut pool = Vec::new();
if let Some(f) = file {
let (tx, rx) = mpsc::sync_channel::<(u64, Vec<u8>, MarkedSeg)>(threads);
let rx = Arc::new(std::sync::Mutex::new(rx));
for _ in 0..4 {
let rx = rx.clone();
pool.push(sc.spawn(move || -> Result<(), ()> {
use std::os::unix::fs::FileExt;
loop {
let job = rx.lock().unwrap().recv();
let Ok((off, h, seg)) = job else {
return Ok(());
};
f.write_all_at(&h, off).map_err(|_| ())?;
f.write_all_at(seg.tail_bytes(), off + h.len() as u64)
.map_err(|_| ())?;
bitsplit::pool::give_u8(h);
bitsplit::pool::recycle(seg);
}
}));
}
sink.pos_out = Some(PosOut {
off: 0,
tx,
n: written.clone(),
});
}
let res = gatling::gatling::ordered::run_ordered_sink(
producer,
threads,
threads + 2,
map,
&mut sink,
);
let tail = sink.finish();
sink.pos_out = None; // closes the pool's channel
let wres: Result<(), ()> = pool
.into_iter()
.map(|h| h.join().unwrap_or(Err(())))
.collect();
(res.map(|_| tail), wres)
});
if let Some(m) = sink.err {
return Err(fail(m));
}
res.map_err(|_| fail("pipeline failed"))?.map_err(fail)?;
wres.map_err(|_| fail("write failed"))?;
sink.w.flush().map_err(|_| fail("write failed"))?;
let mut st = sink.st;
st.written = written.load(Ordering::Relaxed);
Ok(st)
}
fn serial_slice<W: Write>(input: &[u8], mut w: W) -> Result<Stats, StreamError> {
let mut written = 0u64;
let fail = |msg: &'static str, n: u64| StreamError { msg, written: n };
let mut st = Stats {
serial: true,
..Stats::default()
};
let head = linflate::OVERWRITE_HEADROOM;
let mut cap = SLICE_OUT;
let mut out: Vec<u8> = Vec::with_capacity(WINDOW + cap + 2 * head);
// SAFETY: linflate writes every byte before it is read (back-references
// are checked against the window plus what it wrote); only those bytes are
// ever read back here.
#[allow(clippy::uninit_vec)]
unsafe {
out.set_len(WINDOW + cap + 2 * head)
};
let end_bits = input.len() * 8;
let mut pos = match bitsplit::gzip_header_len(input).map_err(|e| fail(e.0, 0))? {
Some(n) => n * 8,
None => return Err(fail("truncated gzip header", 0)),
};
let mut win = 0usize; // window bytes at the front of `out`
let mut crc = flate2::Crc::new();
let mut len = 0u64;
// compressed bits per window: aim the output at `cap`, halve on overflow
let mut span = cap * 8 / 4;
loop {
let stop = if pos + span >= end_bits {
usize::MAX
} else {
pos + span
};
match linflate::inflate_bits(input, pos, stop, &mut out[..win + cap + head], win) {
Ok(r) => {
let new = &out[win..win + r.written];
crc.update(new);
len += r.written as u64;
w.write_all(new)
.map_err(|_| fail("write failed", written))?;
written += r.written as u64;
// the next window: the last 32 KB of (window ++ new)
let total = win + r.written;
let keep = total.min(WINDOW);
out.copy_within(total - keep..total, 0);
win = keep;
pos = r.end_bit;
st.segments += 1;
if r.final_block {
let a = pos.div_ceil(8);
if a + 8 > input.len() {
return Err(fail("truncated gzip trailer", written));
}
let tc = u32::from_le_bytes(input[a..a + 4].try_into().unwrap());
let ts = u32::from_le_bytes(input[a + 4..a + 8].try_into().unwrap());
if crc.sum() != tc || len as u32 != ts {
return Err(fail("gzip checksum mismatch", written));
}
st.members += 1;
crc.reset();
len = 0;
win = 0; // a new member references nothing before it
let h = a + 8;
if h == input.len() {
break;
}
match bitsplit::gzip_header_len(&input[h..]).map_err(|e| fail(e.0, written))? {
Some(n) => pos = (h + n) * 8,
None => return Err(fail("truncated gzip header", written)),
}
}
}
Err(linflate::InflateError::OutputOverflow) => {
if span > 64 * 1024 * 8 {
span /= 2;
} else {
// one block wider than the buffer: grow it, once per need
cap *= 2;
out.reserve(WINDOW + cap + 2 * head - out.len());
#[allow(clippy::uninit_vec)]
unsafe {
out.set_len(WINDOW + cap + 2 * head)
};
}
}
Err(_) => return Err(fail("gzip decode failed", written)),
}
}
w.flush().map_err(|_| fail("write failed", written))?;
st.written = written;
Ok(st)
}
fn serial<R: Read, W: Write>(
first: Vec<u8>,
r: R,
w: W,
written: Arc<AtomicU64>,
) -> Result<Stats, StreamError> {
let input = std::io::Cursor::new(first).chain(r);
let mut dec = flate2::read::MultiGzDecoder::new(input);
let mut w = CountingWriter {
w,
n: written.clone(),
};
// a 1 MiB buffer: io::copy's 8 KiB would cost more than zlib-rs itself
let res = (|| -> std::io::Result<()> {
let mut buf = vec![0u8; 1 << 20];
loop {
let k = dec.read(&mut buf)?;
if k == 0 {
break;
}
w.write_all(&buf[..k])?;
}
w.flush()
})();
let n = written.load(Ordering::Relaxed);
match res {
Ok(()) => Ok(Stats {
segments: 1,
serial: true,
written: n,
..Stats::default()
}),
Err(_) => Err(StreamError {
msg: "gzip decode failed",
written: n,
}),
}
}