oci-image 0.1.0

//! Streaming parallel gzip decode for streams with full-flush points.
//!
//! A feeder thread reads the file, cuts the DEFLATE stream at validated full
//! flushes (`lgz::deflate_scan`) into segments of at least [`MIN_SEGMENT`],
//! and hands them to decoder threads (`lgz::chunk::decode_segment`). The
//! reader returns the decoded segments in order while later ones are still
//! being decoded, so the tar reader starts on the first segment at once and
//! memory stays at a few segments per thread. The CRC32 and size in the gzip
//! trailer are checked over the whole output; a mismatch is an error at the
//! end of the stream.

use std::collections::BTreeMap;
use std::fs::File;
use std::io::{self, Read};
use std::sync::mpsc::{self, Receiver, SyncSender};
use std::sync::{Arc, Mutex};
use std::thread;

/// Smallest segment handed to a decoder thread.
pub(crate) const MIN_SEGMENT: usize = 1 << 20;
/// Bytes read from the file at a time.
const READ: usize = 4 << 20;
/// Bytes past a candidate flush the validation probe needs.
const PROBE_MARGIN: usize = 64 << 10;

type Segment = (usize, io::Result<Vec<u8>>);

pub(crate) struct ParallelGz {
    results: Receiver<Segment>,
    credits: SyncSender<()>,
    trailer: Arc<Mutex<Option<(usize, u32, u32)>>>,
    pending: BTreeMap<usize, io::Result<Vec<u8>>>,
    next: usize,
    cur: Vec<u8>,
    pos: usize,
    crc: flate2::Crc,
    finished: bool,
}

impl ParallelGz {
    /// `file` is positioned at the start of the gzip member.
    pub(crate) fn new(file: File, threads: usize) -> Self {
        let threads = threads.max(1);
        let (job_tx, job_rx) = mpsc::sync_channel::<(usize, Vec<u8>)>(threads);
        let (res_tx, results) = mpsc::channel::<Segment>();
        // Segments cut but not yet read: one per thread and two spare.
        let window = threads + 2;
        let (credits, credit_rx) = mpsc::sync_channel::<()>(window);
        for _ in 0..window {
            credits.send(()).expect("room for the initial credits");
        }
        let job_rx = Arc::new(Mutex::new(job_rx));
        for _ in 0..threads {
            let (rx, tx) = (job_rx.clone(), res_tx.clone());
            thread::spawn(move || loop {
                let job = rx.lock().unwrap().recv();
                let Ok((i, seg)) = job else { return };
                let out = lgz::chunk::decode_segment(&seg)
                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()));
                if tx.send((i, out)).is_err() {
                    return;
                }
            });
        }
        let trailer = Arc::new(Mutex::new(None));
        let t = trailer.clone();
        thread::spawn(move || {
            if let Err((i, e)) = feed(file, &job_tx, &credit_rx, &t) {
                let _ = res_tx.send((i, Err(e)));
            }
        });
        Self {
            results,
            credits,
            trailer,
            pending: BTreeMap::new(),
            next: 0,
            cur: Vec::new(),
            pos: 0,
            crc: flate2::Crc::new(),
            finished: false,
        }
    }

    /// The next segment in order, or `None` after the last.
    fn next_segment(&mut self) -> io::Result<Option<Vec<u8>>> {
        loop {
            if let Some(r) = self.pending.remove(&self.next) {
                self.next += 1;
                let _ = self.credits.try_send(());
                return r.map(Some);
            }
            if let Some((n, _, _)) = *self.trailer.lock().unwrap() {
                if self.next == n {
                    return Ok(None);
                }
            }
            match self.results.recv() {
                Ok((i, r)) => {
                    self.pending.insert(i, r);
                }
                Err(_) => {
                    let done = self
                        .trailer
                        .lock()
                        .unwrap()
                        .is_some_and(|t| t.0 == self.next);
                    if done {
                        return Ok(None);
                    }
                    return Err(io::Error::new(
                        io::ErrorKind::UnexpectedEof,
                        "parallel gzip decode stopped early",
                    ));
                }
            }
        }
    }
}

impl Read for ParallelGz {
    fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
        while self.pos == self.cur.len() {
            if self.finished {
                return Ok(0);
            }
            match self.next_segment()? {
                Some(seg) => {
                    self.crc.update(&seg);
                    self.cur = seg;
                    self.pos = 0;
                }
                None => {
                    self.finished = true;
                    let (_, crc, size) = self.trailer.lock().unwrap().expect("set at the end");
                    if self.crc.sum() != crc || self.crc.amount() != size {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "gzip layer does not match its CRC32/size trailer",
                        ));
                    }
                    return Ok(0);
                }
            }
        }
        let n = out.len().min(self.cur.len() - self.pos);
        out[..n].copy_from_slice(&self.cur[self.pos..self.pos + n]);
        self.pos += n;
        Ok(n)
    }
}

/// Cut the stream into segments and send them, in order, as credits allow.
/// On the way out, record (segment count, crc, size) from the trailer.
fn feed(
    mut file: File,
    jobs: &SyncSender<(usize, Vec<u8>)>,
    credits: &Receiver<()>,
    trailer: &Mutex<Option<(usize, u32, u32)>>,
) -> Result<(), (usize, io::Error)> {
    let mut idx = 0;
    let fail = |idx: usize, e: io::Error| (idx, e);
    let mut buf: Vec<u8> = Vec::with_capacity(READ + MIN_SEGMENT);
    let mut eof = false;
    let read_more = |buf: &mut Vec<u8>, file: &mut File| -> io::Result<bool> {
        let old = buf.len();
        buf.resize(old + READ, 0);
        let mut n = 0;
        while n < READ {
            match file.read(&mut buf[old + n..])? {
                0 => break,
                k => n += k,
            }
        }
        buf.truncate(old + n);
        Ok(n < READ)
    };
    eof |= read_more(&mut buf, &mut file).map_err(|e| fail(0, e))?;
    let start = lgz::deflate_offset(&buf)
        .map_err(|e| fail(0, io::Error::new(io::ErrorKind::InvalidData, e.to_string())))?;
    buf.drain(..start);
    let mut scan_from = MIN_SEGMENT;
    loop {
        // The last 8 bytes of the file are the trailer, not DEFLATE data.
        let avail = if eof {
            buf.len().saturating_sub(8)
        } else {
            buf.len()
        };
        let limit = if eof {
            avail
        } else {
            avail.saturating_sub(PROBE_MARGIN)
        };
        let cut = if scan_from < limit {
            lgz::deflate_scan::find_next_valid_flush(&buf[..avail], scan_from)
                .map(|b| b.start)
                .filter(|&s| s <= limit && s > 0)
        } else {
            None
        };
        let (seg_end, last) = match cut {
            Some(s) => (s, false),
            None if eof => (avail, true),
            None => {
                scan_from = limit.max(MIN_SEGMENT);
                eof |= read_more(&mut buf, &mut file).map_err(|e| fail(idx, e))?;
                continue;
            }
        };
        if credits.recv().is_err() {
            return Ok(()); // the reader is gone
        }
        let seg = buf[..seg_end].to_vec();
        buf.drain(..seg_end);
        if jobs.send((idx, seg)).is_err() {
            return Ok(());
        }
        idx += 1;
        scan_from = MIN_SEGMENT;
        if last {
            if buf.len() != 8 {
                return Err(fail(
                    idx,
                    io::Error::new(io::ErrorKind::InvalidData, "gzip trailer missing"),
                ));
            }
            let crc = u32::from_le_bytes(buf[..4].try_into().unwrap());
            let size = u32::from_le_bytes(buf[4..8].try_into().unwrap());
            *trailer.lock().unwrap() = Some((idx, crc, size));
            return Ok(());
        }
    }
}