From 4fa666ea13f0d4a88314108bd3263fee39b9d718 Mon Sep 17 00:00:00 2001 From: Jeremy Howard Date: Tue, 25 Aug 2026 14:08:12 +1000 Subject: [PATCH] Add parallel streaming Reader and crates.io publishing --- .github/workflows/ci.yml | 43 +++++------ Cargo.toml | 3 +- DEV.md | 27 ++++++- README.md | 27 ++++++- src/bin/fbz.rs | 48 ++++-------- src/bin/fbz/tar_extract.rs | 73 +----------------- src/bin/fbz/zip_extract.rs | 4 + src/decode.rs | 55 +++++++++++++- src/decoder.rs | 14 +++- src/error.rs | 2 + src/format.rs | 34 ++++++--- src/lib.rs | 6 +- src/output.rs | 119 +++++++++++++++++++++++++++++- src/reader.rs | 147 +++++++++++++++++++++++++++++++++++++ src/stream.rs | 73 ++++++++++++++++++ tests/reader.rs | 106 ++++++++++++++++++++++++++ tests/reader_perf.rs | 53 +++++++++++++ 17 files changed, 685 insertions(+), 149 deletions(-) create mode 100644 src/reader.rs create mode 100644 src/stream.rs create mode 100644 tests/reader.rs create mode 100644 tests/reader_perf.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 59fea33..0eda8f8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,6 +17,7 @@ jobs: - uses: dtolnay/rust-toolchain@stable - run: cargo test --release - run: cargo check --all-features + - run: cargo package - uses: actions/setup-python@v7 with: python-version: '3.12' @@ -49,31 +50,9 @@ jobs: name: wheels-${{ matrix.os }} path: dist - sdist: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v7 - - uses: dtolnay/rust-toolchain@stable - - uses: actions/setup-python@v7 - with: - python-version: '3.14' - - run: mkdir -p target/wheel-data - - uses: PyO3/maturin-action@v1 - with: - command: sdist - args: -o dist - - run: | - python -m venv /tmp/fbz-sdist-test - /tmp/fbz-sdist-test/bin/pip install pytest dist/*.tar.gz - /tmp/fbz-sdist-test/bin/pytest tests/test_install.py - - uses: actions/upload-artifact@v7 - with: - name: wheels-sdist - path: dist - publish: if: startsWith(github.ref, 'refs/tags/v') - needs: [build, sdist] + needs: build runs-on: ubuntu-latest permissions: id-token: write @@ -91,3 +70,21 @@ jobs: - uses: pypa/gh-action-pypi-publish@release/v1 with: packages-dir: dist/ + + publish-crate: + if: startsWith(github.ref, 'refs/tags/v') + needs: test + runs-on: ubuntu-latest + permissions: + id-token: write + steps: + - uses: actions/checkout@v7 + - uses: dtolnay/rust-toolchain@stable + - id: auth + uses: rust-lang/crates-io-auth-action@v1 + - name: Publish crate when its version is new + run: | + version=$(cargo metadata --format-version 1 --no-deps | jq -r '.packages[] | select(.name=="fbz") | .version') + if curl -sf -A "fbz-ci (https://github.com/AnswerDotAI/fbz)" "https://crates.io/api/v1/crates/fbz/$version" >/dev/null; then echo "fbz $version is already published"; else cargo publish; fi + env: + CARGO_REGISTRY_TOKEN: ${{ steps.auth.outputs.token }} diff --git a/Cargo.toml b/Cargo.toml index ffcdad0..7ccc714 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,7 +7,8 @@ license = "Apache-2.0" description = "Compression-format research workbench with fast bzip2, gzip, LZ4, and ZIP decompression" repository = "https://github.com/AnswerDotAI/fbz" homepage = "https://github.com/AnswerDotAI/fbz" -documentation = "https://github.com/AnswerDotAI/fbz" +documentation = "https://docs.rs/fbz" +include = ["/src/**", "/Cargo.toml", "/README.md", "/LICENSE"] [lib] name = "fbz" diff --git a/DEV.md b/DEV.md index 113a32b..a570ac4 100644 --- a/DEV.md +++ b/DEV.md @@ -49,6 +49,10 @@ ZIP scheduling deliberately uses one parallelism level at a time. A sole DEFLATE The shared `pipeline.rs` scheduler provides ordered results, byte-budgeted admission, cancellation, and a staged priority queue. Bzip2 uses the rolling candidate path: workers reserve the maximum possible decoded block size, then shrink that reservation to actual retained output until ordered validation consumes or rejects it. Gzip uses the staged path: native workers alternate speculative DEFLATE decoding with higher-priority marker resolution, while the coordinator advances only the 32 KiB dependency windows and emits resolved chunks in order. LZ4 independent blocks use the simpler ordered-job path and charge each retained allocation against the budget. Decode results and outstanding resolution results have bounded horizons, preventing dependency stalls from causing unbounded memory. +The shared `OutputSink` boundary accepts owned decoder chunks. Direct writer APIs adapt those chunks to `Write::write_all` without a channel or allocation copy. `Reader` and streaming tar extraction instead use the same zero-capacity owned-chunk pipe, so completed bzip2 blocks and resolved parallel-gzip segments move into the consumer rather than being copied into an intermediate pipe buffer. The pipe adds at most the consumer's current chunk and the producer's next blocked chunk beyond the scheduler budget. Its receiver owns a cancellation flag; parallel bzip2 scanning checks that flag between bounded waves. + +The non-indexing parallel bzip2 path decodes and emits its first CRC-validated block before scanning the remainder of the compressed source. That candidate is then reused by normal ordered assembly rather than decoded twice. Index construction retains the full scan-first path because it produces no output. `Reader` sends decoder success or failure separately from the byte pipe and treats a missing terminal status as an error, so a worker panic or corrupt trailer cannot appear as EOF. Reader errors are sticky across subsequent calls, and drop disconnects the pipe before joining the decoder thread. + The bzip2 decoder is safe scalar Rust designed for LLVM auto-vectorisation. Huffman decoding uses a 4096-entry direct table for codes up to 12 bits and canonical fallback for longer codes. Gzip uses full canonical lookup tables packed into `u16`, a branch-free 64 KiB marker-resolution lookup for large chunks, and `crc32fast::Hasher::combine` so CRC scanning runs with the resolution workers rather than serially in the coordinator. Gzip byte/marker output and LZ4 share `history::extend_match`, whose doubling copies handle overlapping matches in logarithmically many operations; format-specific history validation remains at each call site. The only unsafe codec operation marks a just-initialized `Vec` result as initialized after writing every spare-capacity byte. Add architecture-specific SIMD or further cross-codec abstraction only after profiling; `libbz2-rs-sys`, `flate2`, and `lz4_flex` remain dev-only differential oracles. ## Commands @@ -58,6 +62,7 @@ cargo test cargo test --release cargo check --all-features cargo build --release --bins +cargo package python tools/stage_binaries.py uv pip install --reinstall --no-deps -e . pytest -q @@ -74,6 +79,21 @@ The normal release test path decodes selected valid and corrupt cases from the m The normal release path contains warmed end-to-end performance regression gates capped at 1.3 times each oracle, allowing for noise on shared runners. The gzip gates independently exercise a highly compressible LZ77-heavy shape and an incompressible literal-heavy shape against `flate2`; the bzip2 gate uses `libbz2-rs-sys`. Representative local acceptance remains 1.2 times the corresponding oracle. The ignored full-wiki gzip test applies that threshold to rapidgzip-rust. Keep the whole release test suite below five seconds on the primary development laptop; individual timed workloads should normally be about 0.1 seconds or less. +### Local Reader benchmark + +`tests/reader_perf.rs` compares the public `Read` adapter with the direct writer path on both 84.4 MiB SimpleWiki fixtures. Each row is one release-mode run, and both timings include opening the file. The Reader path necessarily copies into the caller's buffer but transfers decoder-owned chunks into its rendezvous pipe without another copy: + +```bash +cargo test --release --test reader_perf reader_writer_comparison -- --ignored --exact --nocapture +``` + +| Format | First byte | Reader | Direct writer | Reader/writer | +|---|---:|---:|---:|---:| +| bzip2 | 36.627 ms | 205.081 ms | 183.465 ms | 1.118x | +| gzip | 11.509 ms | 28.119 ms | 24.515 ms | 1.147x | + +The bzip2 first-byte measurement includes opening, mapping, and decoding its first block, but not the subsequent whole-input parallel marker scan. The gzip path parses only the current member header before beginning DEFLATE work. Keep this diagnostic single-run. + ### Local LZ4 benchmark reproduction and diagnostics `tests/lz4_perf.rs` decodes `meta/simplewiki-first-5pct.xml.bz2` before timing, then uses dev-only `lz4_flex` to create a standard Max4MB independent-block LZ4 frame with a content checksum. Fixture construction and both warm-ups are outside the measured intervals. One test then measures exactly one validation run of each CLI and records child-process memory without task-inspection permissions: @@ -276,13 +296,14 @@ CI tests and builds Linux on x86-64 and ARM64, and macOS on ARM64. macOS Intel r The canonical version lives in `Cargo.toml`. `pyproject.toml` gets the Python package version from Cargo via `dynamic = ["version"]`. -The thin PEP 517 backend delegates to Maturin after building and staging the native executable. This makes wheels built from the sdist contain the same native CLI as CI-built wheels. Release CI verifies that path from a fresh Python 3.14 environment before publishing. +The thin PEP 517 backend delegates to Maturin after building and staging the native executable, so repository and editable builds contain the same native CLI as CI-built wheels. Releases publish platform wheels rather than a Python sdist; the Rust source package is published separately to crates.io. ## Release 1. Run `cargo build --release --bins && python tools/stage_binaries.py`. 2. Run `uv pip install --reinstall --no-deps -e . && pytest -q` so the custom backend installs both the extension and native CLI. 3. Confirm the release version in `Cargo.toml` (`[package].version`). -4. Run `ship-release`. +4. For the first crates.io release only, run `cargo publish`, then configure the `ci.yml` trusted publisher for `AnswerDotAI/fbz`; crates.io requires the crate to exist before trusted publishing can be configured. +5. Run `ship-release`. -Fastship pushes the version tag for GitHub Actions, then bumps and pushes `Cargo.toml`. +Fastship pushes the version tag for GitHub Actions, then bumps and pushes `Cargo.toml`. Tagged CI publishes the crate through crates.io trusted publishing as well as building the GitHub release and PyPI packages. diff --git a/README.md b/README.md index a7894b1..1a9e3fc 100644 --- a/README.md +++ b/README.md @@ -16,11 +16,11 @@ pip install fbz Python 3.10 and later are supported. Prebuilt wheels target Linux on x86-64 and ARM64, and macOS on ARM64. macOS Intel is best-effort and can build from source. -The Rust crate is not yet published separately on crates.io. Install the CLI from the repository, or add the library as a Git dependency: +Install the native CLI or add the Rust library from crates.io: ```bash -cargo install --git https://github.com/AnswerDotAI/fbz -cargo add fbz --git https://github.com/AnswerDotAI/fbz +cargo install fbz +cargo add fbz ``` ## CLI @@ -157,6 +157,27 @@ let plain = fbz::lz4::decompress(&compressed_lz4)?; It accepts standard independent or linked blocks, stored blocks, all four standard block maxima, optional block/content checksums and sizes, concatenated frames, and skippable frames. External dictionaries and the obsolete legacy frame format are intentionally unsupported. +### Streaming reads + +`fbz::Reader` provides a normal `std::io::Read` over bzip2 or gzip files without a preliminary indexing or validation pass: + +```rust +use std::io::{BufReader, Read}; +use fbz::{DecodeOptions, Reader}; + +fn main() -> Result<(), Box> { + let reader = Reader::open("dump.xml.bz2", DecodeOptions::default())?; + let mut reader = BufReader::new(reader); + let mut header = [0; 4096]; + reader.read_exact(&mut header)?; + Ok(()) +} +``` + +Magic takes priority over the filename extension, with the extension used as a fallback for damaged headers. The decoder runs on an owned worker thread and transfers completed decoder allocations through a zero-capacity pipe; it neither materializes the plaintext nor writes an intermediate file. `DecodeOptions` controls decoder threads and speculative memory. Dropping early disconnects the pipe, cancels outstanding work, and joins the worker. + +Checksum errors discovered after output has begun are returned by a later `read()` call. Therefore only successful EOF establishes that the complete stream was valid; dropping early deliberately does not finish validation. Compressed tar inputs yield the decoded tar byte stream rather than extracting it. LZ4 and ZIP are not currently exposed through `Reader`: LZ4's frame-layout pass must first become incremental, while ZIP has no single decoded byte stream. + ## Performance These are single local release-mode CLI runs on the primary Apple Silicon development machine. The gzip comparison warms each executable with the 5% fixture, then measures exactly one full validation run; peak physical footprint comes from a separate sampled run because process inspection can perturb such a short workload. ZIP and tar likewise warm each CLI once and then measure one extraction. The LZ4 comparison uses its automatic four-worker limit. These are observations rather than statistical aggregates. In-process codec and library-oracle comparisons live in [DEV.md](DEV.md), not this user-facing section. diff --git a/src/bin/fbz.rs b/src/bin/fbz.rs index 25d147d..a8f908e 100644 --- a/src/bin/fbz.rs +++ b/src/bin/fbz.rs @@ -15,8 +15,8 @@ use std::{ use clap::{ArgGroup, Parser}; use fbz::{ - DecodeOptions, DecodeProgress, Error, Index, OutputSink, Source, WriterSink, build_index_with_progress, decode_to_writer_with_progress, - decompress_to_sink_with_progress, gzip, lz4, + DecodeOptions, DecodeProgress, Error, Format, Index, OutputSink, Source, WriterSink, build_index_with_progress, decode_stream_to_sink_with_progress, + decode_to_writer_with_progress, gzip, lz4, }; use serde_json::{Value, json}; use tempfile::NamedTempFile; @@ -85,14 +85,6 @@ struct Cli { json: bool, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum Format { - Bzip2, - Gzip, - Lz4, - Zip, -} - fn main() -> ExitCode { match run(Cli::parse()) { Ok(()) => ExitCode::SUCCESS, @@ -326,12 +318,7 @@ fn decode_data_to_sink( ) -> fbz::Result<()> { let mut output = LimitedOutput::new(output, max_output); let mut display = ProgressDisplay::new(label, data.len() as u64, quiet); - match select_format(label, data)? { - Format::Bzip2 => decompress_to_sink_with_progress(data, &mut output, options, |progress| display.update(progress)), - Format::Gzip => gzip::decompress_to_sink_with_options_and_progress(data, &mut output, options, |progress| display.update(progress)).map(|_| ()), - Format::Lz4 => lz4::decompress_to_sink_with_options_and_progress(data, &mut output, options, |progress| display.update(progress)).map(|_| ()), - Format::Zip => Err(invalid("ZIP archives extract to a directory and cannot be decoded to one output stream")), - } + decode_stream_to_sink_with_progress(select_format(label, data)?, data, &mut output, options, |progress| display.update(progress)) } fn build_gzip_report_data(data: &[u8], label: &str, options: DecodeOptions, max_output: Option, quiet: bool) -> fbz::Result { @@ -406,6 +393,10 @@ impl OutputSink for LimitedOutput { fn flush(&mut self) -> io::Result<()> { self.inner.flush() } + + fn is_cancelled(&self) -> bool { + self.inner.is_cancelled() + } } struct ProgressDisplay { @@ -540,22 +531,7 @@ fn default_output(input: &Path) -> PathBuf { } fn select_format(input: &str, data: &[u8]) -> fbz::Result { - if let Some((format, _)) = format_extension(Path::new(input)) { - return Ok(format); - } - if data.starts_with(b"BZh") { - Ok(Format::Bzip2) - } else if data.starts_with(&[0x1f, 0x8b]) { - Ok(Format::Gzip) - } else if data.starts_with(&[0x04, 0x22, 0x4d, 0x18]) - || data.get(..4).is_some_and(|magic| (0x184d_2a50..=0x184d_2a5f).contains(&u32::from_le_bytes(magic.try_into().unwrap()))) - { - Ok(Format::Lz4) - } else if data.starts_with(b"PK\x03\x04") || data.starts_with(b"PK\x05\x06") || data.starts_with(b"PK\x07\x08") { - Ok(Format::Zip) - } else { - Err(invalid(format!("cannot determine compression format for {input}; expected a bzip2, gzip, LZ4, or ZIP extension or magic"))) - } + Format::detect(input, data) } fn print_index(input: Option<&String>, index: &Index) { @@ -805,7 +781,13 @@ fn exit_status(error: &Error) -> u8 { Error::Io(source) if source.kind() == io::ErrorKind::InvalidData => 3, Error::Io(_) => 1, Error::InvalidConfiguration(_) => 2, - Error::InvalidStreamHeader | Error::InvalidGzip(_) | Error::InvalidLz4(_) | Error::InvalidZip(_) | Error::Decode { .. } | Error::InvalidIndex(_) => 3, + Error::InvalidStreamHeader + | Error::InvalidGzip(_) + | Error::InvalidLz4(_) + | Error::InvalidZip(_) + | Error::UnsupportedFormat(_) + | Error::Decode { .. } + | Error::InvalidIndex(_) => 3, _ => 4, } } diff --git a/src/bin/fbz/tar_extract.rs b/src/bin/fbz/tar_extract.rs index fc603fb..f548561 100644 --- a/src/bin/fbz/tar_extract.rs +++ b/src/bin/fbz/tar_extract.rs @@ -1,76 +1,9 @@ -use std::{ - cmp, - io::{self, Read}, - path::Path, - sync::mpsc::{Receiver, SyncSender, sync_channel}, - thread, -}; +use std::{io, path::Path, thread}; -use fbz::{Error, OutputSink, Result}; +use fbz::{Error, PipeWriter, Result, output_pipe}; use super::archive_extract; -struct Chunk { - bytes: Vec, - offset: usize, -} - -pub(super) struct PipeWriter { - sender: SyncSender, -} - -impl PipeWriter { - fn send(&self, bytes: Vec, offset: usize) -> io::Result<()> { - let suffix = bytes.get(offset..).ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "owned chunk start exceeds its length"))?; - if suffix.is_empty() { - return Ok(()); - } - self.sender.send(Chunk { bytes, offset }).map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "tar extractor stopped reading")) - } -} - -impl OutputSink for PipeWriter { - fn write_borrowed(&mut self, buffer: &[u8]) -> io::Result<()> { - self.send(buffer.to_vec(), 0) - } - - fn write_owned_from(&mut self, buffer: Vec, start: usize) -> io::Result<()> { - self.send(buffer, start) - } -} - -struct PipeReader { - receiver: Receiver, - chunk: Vec, - offset: usize, -} - -impl Read for PipeReader { - fn read(&mut self, buffer: &mut [u8]) -> io::Result { - if buffer.is_empty() { - return Ok(0); - } - if self.offset == self.chunk.len() { - match self.receiver.recv() { - Ok(chunk) => { - self.chunk = chunk.bytes; - self.offset = chunk.offset; - } - Err(_) => return Ok(0), - } - } - let count = cmp::min(buffer.len(), self.chunk.len() - self.offset); - buffer[..count].copy_from_slice(&self.chunk[self.offset..self.offset + count]); - self.offset += count; - Ok(count) - } -} - -fn pipe() -> (PipeWriter, PipeReader) { - let (sender, receiver) = sync_channel(0); - (PipeWriter { sender }, PipeReader { receiver, chunk: Vec::new(), offset: 0 }) -} - fn broken_pipe(error: &Error) -> bool { matches!(error, Error::Io(source) if source.kind() == io::ErrorKind::BrokenPipe) } @@ -81,7 +14,7 @@ where { let staging = archive_extract::staging(destination)?; thread::scope(|scope| { - let (mut writer, mut reader) = pipe(); + let (mut writer, mut reader) = output_pipe(); let decoder = scope.spawn(move || decode(&mut writer)); let extracted = { let mut archive = tar::Archive::new(&mut reader); diff --git a/src/bin/fbz/zip_extract.rs b/src/bin/fbz/zip_extract.rs index 4b796cf..ca39c0f 100644 --- a/src/bin/fbz/zip_extract.rs +++ b/src/bin/fbz/zip_extract.rs @@ -210,6 +210,10 @@ impl OutputSink for ExpectedOutput { fn flush(&mut self) -> io::Result<()> { self.inner.flush() } + + fn is_cancelled(&self) -> bool { + self.inner.is_cancelled() + } } fn entry_error(entry: &Entry, error: Error) -> Error { diff --git a/src/decode.rs b/src/decode.rs index 06d181c..9720427 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -90,7 +90,18 @@ pub fn decompress_to_sink_with_progress( progress(DecodeProgress { compressed_bytes, decoded_bytes }); }); } - decode_to_sink_impl(data, output, options, &mut progress).map(|_| ()) + let prefetched = if data.get(4..10) == Some(&[0x31, 0x41, 0x59, 0x26, 0x53, 0x59]) { + let level = parse_header(data, 0)?; + let (expected_crc, mut decoded) = decoder::decode_first_candidate(data)?; + if decoded.block_len > usize::from(level) * 100_000 { + return Err(Error::Decode { bit_offset: 32, source: DecodeError::BlockOverflow }); + } + output.write_owned_from(std::mem::take(&mut decoded.output), 0)?; + Some(PrefetchedCandidate { bit_offset: 32, expected_crc, decoded }) + } else { + None + }; + decode_to_sink_impl_with_prefetched(data, output, options, &mut progress, prefetched).map(|_| ()) } pub fn build_index(data: &[u8], options: DecodeOptions) -> Result { @@ -115,10 +126,26 @@ pub fn decode_to_writer_with_progress(data: &[u8], output: &mut impl Write, opti } fn decode_to_sink_impl(data: &[u8], output: &mut impl OutputSink, options: DecodeOptions, progress: &mut impl FnMut(DecodeProgress)) -> Result { + decode_to_sink_impl_with_prefetched(data, output, options, progress, None) +} + +struct PrefetchedCandidate { + bit_offset: u64, + expected_crc: u32, + decoded: decoder::DecodedCandidate, +} + +fn decode_to_sink_impl_with_prefetched( + data: &[u8], + output: &mut impl OutputSink, + options: DecodeOptions, + progress: &mut impl FnMut(DecodeProgress), + prefetched: Option, +) -> Result { let options = options.validate()?; let threads = options.resolved_threads(); let pool = thread_pool(threads)?; - let scanned = scan_with_pool(data, pool.as_deref())?; + let scanned = scan_with_pool(data, pool.as_deref(), || output.is_cancelled())?; let mut markers: Vec<_> = scanned.blocks.into_iter().map(Marker::Block).collect(); markers.extend(scanned.stream_ends.into_iter().map(Marker::End)); markers.sort_unstable_by_key(Marker::bit_offset); @@ -128,7 +155,19 @@ fn decode_to_sink_impl(data: &[u8], output: &mut impl OutputSink, options: Decod return Err(Error::InvalidConfiguration("input contains too many speculative markers".into())); } + let prefetched = prefetched + .map(|prefetched| { + let marker_index = marker_at(&markers, prefetched.bit_offset)?; + let Marker::Block(block) = &markers[marker_index] else { return Err(required_marker(prefetched.bit_offset)) }; + if block.expected_crc != prefetched.expected_crc { + return Err(Error::Decode { bit_offset: prefetched.bit_offset, source: DecodeError::InvalidBlock }); + } + Ok((marker_index, prefetched.decoded)) + }) + .transpose()?; + let Some(pool) = pool else { + debug_assert!(prefetched.is_none()); let mut candidates = SerialCandidates { data, markers: &markers }; return assemble(data, output, &markers, &mut candidates, progress); }; @@ -136,6 +175,7 @@ fn decode_to_sink_impl(data: &[u8], output: &mut impl OutputSink, options: Decod .iter() .enumerate() .filter_map(|(marker_index, marker)| match marker { + Marker::Block(_) if prefetched.as_ref().is_some_and(|(prefetched_index, _)| *prefetched_index == marker_index) => None, Marker::Block(block) => Some(Job { key: marker_index, reservation: MAX_DECODED_BLOCK, @@ -151,7 +191,7 @@ fn decode_to_sink_impl(data: &[u8], output: &mut impl OutputSink, options: Decod |job| decoder::decode_candidate(data, job.start_bit, job.expected_crc), candidate_len, |results| { - let mut candidates = ParallelCandidates { results }; + let mut candidates = ParallelCandidates { results, prefetched }; assemble(data, output, &markers, &mut candidates, progress) }, ) @@ -204,7 +244,7 @@ fn assemble( } let end_index = marker_at(markers, decoded.end_bit)?; candidates.discard_before(end_index); - let decoded_len = decoded.output.len() as u64; + let decoded_len = decoded.decoded_len as u64; output.write_owned_from(decoded.output, 0)?; blocks.push(BlockIndex { compressed_start_bit: block.bit_offset, @@ -266,14 +306,21 @@ fn candidate_len(result: &Result) -> usize { struct ParallelCandidates<'results, 'pipeline> { results: &'results mut OrderedResults<'pipeline, Result>, + prefetched: Option<(usize, decoder::DecodedCandidate)>, } impl Candidates for ParallelCandidates<'_, '_> { fn take(&mut self, marker_index: usize) -> Result { + if self.prefetched.as_ref().is_some_and(|(prefetched_index, _)| *prefetched_index == marker_index) { + return Ok(self.prefetched.take().unwrap().1); + } self.results.take(marker_index)? } fn discard_before(&mut self, marker_index: usize) { + if self.prefetched.as_ref().is_some_and(|(prefetched_index, _)| *prefetched_index < marker_index) { + self.prefetched.take(); + } self.results.discard_before(marker_index); } } diff --git a/src/decoder.rs b/src/decoder.rs index 2cc0e92..7d97d35 100644 --- a/src/decoder.rs +++ b/src/decoder.rs @@ -191,6 +191,7 @@ struct Decoder { pub(crate) struct DecodedCandidate { pub output: Vec, + pub decoded_len: usize, pub end_bit: u64, pub block_len: usize, } @@ -380,7 +381,18 @@ pub(crate) fn decode_candidate(data: &[u8], start_bit: u64, expected_crc: u32) - return Err(decode_error(start_bit, DecodeError::InvalidMagic)); } let (output, _, block_len) = Decoder::new().block(&mut bits, 9, Some(expected_crc))?; - Ok(DecodedCandidate { output, end_bit: bits.position(), block_len }) + let decoded_len = output.len(); + Ok(DecodedCandidate { output, decoded_len, end_bit: bits.position(), block_len }) +} + +pub(crate) fn decode_first_candidate(data: &[u8]) -> Result<(u32, DecodedCandidate)> { + let mut bits = Bits::at(data, 32)?; + if bits.magic()? != BLOCK_MAGIC { + return Err(decode_error(32, DecodeError::InvalidMagic)); + } + let (output, expected_crc, block_len) = Decoder::new().block(&mut bits, 9, None)?; + let decoded_len = output.len(); + Ok((expected_crc, DecodedCandidate { output, decoded_len, end_bit: bits.position(), block_len })) } pub(crate) fn decode_block(data: &[u8], start_bit: u64, end_bit: u64, level: u8, expected_crc: u32) -> Result> { diff --git a/src/error.rs b/src/error.rs index 3dc82ec..2704fee 100644 --- a/src/error.rs +++ b/src/error.rs @@ -40,6 +40,7 @@ pub enum Error { InvalidGzip(String), InvalidLz4(String), InvalidZip(String), + UnsupportedFormat(String), Decode { bit_offset: u64, source: DecodeError }, InvalidIndex(String), InvalidConfiguration(String), @@ -60,6 +61,7 @@ impl fmt::Display for Error { Self::InvalidGzip(message) => write!(f, "invalid gzip stream: {message}"), Self::InvalidLz4(message) => write!(f, "invalid LZ4 frame: {message}"), Self::InvalidZip(message) => write!(f, "invalid ZIP archive: {message}"), + Self::UnsupportedFormat(message) => f.write_str(message), Self::Decode { bit_offset, source } => write!(f, "bzip2 decode error at bit {bit_offset}: {source}"), Self::InvalidIndex(message) => write!(f, "invalid fbz index: {message}"), Self::InvalidConfiguration(message) => write!(f, "invalid configuration: {message}"), diff --git a/src/format.rs b/src/format.rs index e1f4608..ba271b3 100644 --- a/src/format.rs +++ b/src/format.rs @@ -1,3 +1,5 @@ +use std::io; + use crate::{BitReader, Error, Result}; use rayon::{ThreadPool, prelude::*}; @@ -49,21 +51,35 @@ pub fn scan(data: &[u8]) -> Result { Ok(scan_range(data, 0, data.len())) } -pub(crate) fn scan_with_pool(data: &[u8], pool: Option<&ThreadPool>) -> Result { +pub(crate) fn scan_with_pool(data: &[u8], pool: Option<&ThreadPool>, mut cancelled: impl FnMut() -> bool) -> Result { if !is_stream_header(data, 0) { return Err(Error::InvalidStreamHeader); } - let Some(pool) = pool.filter(|_| data.len() > SCAN_CHUNK) else { return Ok(scan_range(data, 0, data.len())) }; + let cancellation_error = || Error::Io(io::Error::new(io::ErrorKind::BrokenPipe, "output reader stopped reading")); + if cancelled() { + return Err(cancellation_error()); + } + let Some(pool) = pool.filter(|_| data.len() > SCAN_CHUNK) else { + let result = scan_range(data, 0, data.len()); + return if cancelled() { Err(cancellation_error()) } else { Ok(result) }; + }; let chunks = data.len().div_ceil(SCAN_CHUNK); - let partial: Vec<_> = - pool.install(|| (0..chunks).into_par_iter().map(|chunk| scan_range(data, chunk * SCAN_CHUNK, ((chunk + 1) * SCAN_CHUNK).min(data.len()))).collect()); let mut result = ScanResult { streams: Vec::new(), blocks: Vec::new(), stream_ends: Vec::new() }; - for mut chunk in partial { - result.streams.append(&mut chunk.streams); - result.blocks.append(&mut chunk.blocks); - result.stream_ends.append(&mut chunk.stream_ends); + let wave = pool.current_num_threads().max(1); + for first in (0..chunks).step_by(wave) { + if cancelled() { + return Err(cancellation_error()); + } + let end = first.saturating_add(wave).min(chunks); + let partial: Vec<_> = pool + .install(|| (first..end).into_par_iter().map(|chunk| scan_range(data, chunk * SCAN_CHUNK, ((chunk + 1) * SCAN_CHUNK).min(data.len()))).collect()); + for mut chunk in partial { + result.streams.append(&mut chunk.streams); + result.blocks.append(&mut chunk.blocks); + result.stream_ends.append(&mut chunk.stream_ends); + } } - Ok(result) + if cancelled() { Err(cancellation_error()) } else { Ok(result) } } fn scan_range(data: &[u8], start: usize, end: usize) -> ScanResult { diff --git a/src/lib.rs b/src/lib.rs index 9a1fd1b..d01737f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,7 +15,9 @@ mod indexed; pub mod lz4; mod output; mod pipeline; +mod reader; mod source; +mod stream; pub use bitreader::BitReader; pub use block::{MAX_DECODED_BLOCK, MAX_ENCODED_BLOCK, decode_block}; @@ -28,8 +30,10 @@ pub use error::{DecodeError, Error, Result}; pub use format::{BLOCK_MAGIC, BlockCandidate, END_MAGIC, EndCandidate, ScanResult, StreamHeaderCandidate, scan}; pub use index::{BlockIndex, Index, StreamIndex}; pub use indexed::{DEFAULT_CACHE_LIMIT, IndexedReader}; -pub use output::{OutputSink, WriterSink}; +pub use output::{OutputSink, PipeReader, PipeWriter, WriterSink, output_pipe}; +pub use reader::Reader; pub use source::Source; +pub use stream::{Format, decode_stream_to_sink_with_progress}; #[cfg(feature = "python")] mod python { diff --git a/src/output.rs b/src/output.rs index 35fbffe..14571bb 100644 --- a/src/output.rs +++ b/src/output.rs @@ -1,4 +1,12 @@ -use std::io::{self, Write}; +use std::{ + cmp, + io::{self, Read, Write}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + mpsc::{Receiver, SyncSender, sync_channel}, + }, +}; /// Receives decoded bytes and can take ownership of decoder chunks. /// @@ -15,6 +23,11 @@ pub trait OutputSink { fn flush(&mut self) -> io::Result<()> { Ok(()) } + + /// Report whether the downstream consumer has stopped accepting output. + fn is_cancelled(&self) -> bool { + false + } } impl OutputSink for &mut S { @@ -29,6 +42,10 @@ impl OutputSink for &mut S { fn flush(&mut self) -> io::Result<()> { (**self).flush() } + + fn is_cancelled(&self) -> bool { + (**self).is_cancelled() + } } /// Adapts any `std::io::Write` destination to an `OutputSink`. @@ -51,3 +68,103 @@ impl OutputSink for WriterSink { self.writer.flush() } } + +struct Chunk { + bytes: Vec, + offset: usize, +} + +/// The producing side of a bounded owned-chunk output pipe. +#[doc(hidden)] +pub struct PipeWriter { + sender: SyncSender, + cancelled: Arc, +} + +impl PipeWriter { + fn send(&self, bytes: Vec, offset: usize) -> io::Result<()> { + let suffix = bytes.get(offset..).ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "owned chunk start exceeds its length"))?; + if suffix.is_empty() { + return Ok(()); + } + self.sender.send(Chunk { bytes, offset }).map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "output reader stopped reading")) + } +} + +impl OutputSink for PipeWriter { + fn write_borrowed(&mut self, bytes: &[u8]) -> io::Result<()> { + self.send(bytes.to_vec(), 0) + } + + fn write_owned_from(&mut self, bytes: Vec, start: usize) -> io::Result<()> { + self.send(bytes, start) + } + + fn is_cancelled(&self) -> bool { + self.cancelled.load(Ordering::Relaxed) + } +} + +/// The consuming side of a bounded owned-chunk output pipe. +#[doc(hidden)] +pub struct PipeReader { + receiver: Receiver, + chunk: Vec, + offset: usize, + cancelled: Arc, +} + +impl Drop for PipeReader { + fn drop(&mut self) { + self.cancelled.store(true, Ordering::Relaxed); + } +} + +impl Read for PipeReader { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if buffer.is_empty() { + return Ok(0); + } + if self.offset == self.chunk.len() { + match self.receiver.recv() { + Ok(chunk) => { + self.chunk = chunk.bytes; + self.offset = chunk.offset; + } + Err(_) => return Ok(0), + } + } + let count = cmp::min(buffer.len(), self.chunk.len() - self.offset); + buffer[..count].copy_from_slice(&self.chunk[self.offset..self.offset + count]); + self.offset += count; + Ok(count) + } +} + +/// Create a zero-capacity pipe that transfers owned decoder chunks. +#[doc(hidden)] +pub fn output_pipe() -> (PipeWriter, PipeReader) { + let (sender, receiver) = sync_channel(0); + let cancelled = Arc::new(AtomicBool::new(false)); + (PipeWriter { sender, cancelled: Arc::clone(&cancelled) }, PipeReader { receiver, chunk: Vec::new(), offset: 0, cancelled }) +} + +#[cfg(test)] +mod tests { + use std::thread; + + use super::*; + + #[test] + fn pipe_transfers_owned_allocation() { + let (mut writer, mut reader) = output_pipe(); + let bytes = vec![1, 2, 3, 4]; + let pointer = bytes.as_ptr() as usize; + let worker = thread::spawn(move || writer.write_owned_from(bytes, 1).unwrap()); + let mut output = [0; 3]; + assert_eq!(reader.read(&mut output).unwrap(), 3); + assert_eq!(output, [2, 3, 4]); + assert_eq!(reader.chunk.as_ptr() as usize, pointer); + worker.join().unwrap(); + } +} diff --git a/src/reader.rs b/src/reader.rs new file mode 100644 index 0000000..ea74377 --- /dev/null +++ b/src/reader.rs @@ -0,0 +1,147 @@ +use std::{ + io::{self, Read}, + path::Path, + sync::mpsc::{Receiver, sync_channel}, + thread::{self, JoinHandle}, +}; + +use crate::{DecodeOptions, Error, Format, PipeReader, Result, Source, decode_stream_to_sink_with_progress, output_pipe}; + +#[derive(Clone)] +struct StoredError { + kind: io::ErrorKind, + message: String, +} + +impl StoredError { + fn new(error: io::Error) -> Self { + Self { kind: error.kind(), message: error.to_string() } + } + + fn io_error(&self) -> io::Error { + io::Error::new(self.kind, self.message.clone()) + } +} + +enum State { + Reading, + Eof, + Failed(StoredError), +} + +/// A streaming, parallel bzip2 or gzip decoder for file-backed streams. +/// +/// Successful EOF means the complete stream and its checksums were validated. +/// A corrupt trailer can therefore produce decoded bytes before a later call to +/// `read` returns an error. Dropping before EOF cancels decoding without +/// completing validation. +pub struct Reader { + pipe: Option, + result: Option>>, + worker: Option>, + state: State, +} + +impl Reader { + /// Open a bzip2 or gzip file and start its decoder coordinator. + /// + /// This opens and memory-maps only the compressed source. Format and option + /// errors are returned here; compressed-data errors are returned later by + /// `Read::read` as the stream is consumed. + pub fn open(path: impl AsRef, options: DecodeOptions) -> Result { + let options = options.validate()?; + let path = path.as_ref(); + let source = Source::open(path)?; + let format = Format::detect(path, source.as_slice())?; + if !matches!(format, Format::Bzip2 | Format::Gzip) { + return Err(Error::UnsupportedFormat("fbz::Reader currently supports bzip2 and gzip streams".into())); + } + + let (mut output, pipe) = output_pipe(); + let (result_sender, result) = sync_channel(1); + let worker = thread::Builder::new() + .name("fbz-reader".into()) + .spawn(move || { + let decoded = decode_stream_to_sink_with_progress(format, source.as_slice(), &mut output, options, |_| {}); + drop(output); + let _ = result_sender.send(decoded); + }) + .map_err(Error::from)?; + Ok(Self { pipe: Some(pipe), result: Some(result), worker: Some(worker), state: State::Reading }) + } + + fn join_worker(&mut self) -> io::Result<()> { + let Some(worker) = self.worker.take() else { return Ok(()) }; + worker.join().map_err(|_| io::Error::other("fbz decoder worker panicked")) + } + + fn fail(&mut self, error: io::Error) -> io::Error { + let error = StoredError::new(error); + let returned = error.io_error(); + self.state = State::Failed(error); + returned + } + + fn finish(&mut self, result: Result<()>) -> io::Result<()> { + self.pipe.take(); + self.result.take(); + let joined = self.join_worker(); + match result { + Err(error) => Err(self.fail(read_error(error))), + Ok(()) => match joined { + Ok(()) => { + self.state = State::Eof; + Ok(()) + } + Err(error) => Err(self.fail(error)), + }, + } + } + + fn disconnected(&mut self) -> io::Error { + self.pipe.take(); + self.result.take(); + let error = self.join_worker().err().unwrap_or_else(|| io::Error::other("fbz decoder stopped without reporting completion")); + self.fail(error) + } +} + +impl Read for Reader { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if buffer.is_empty() { + return Ok(0); + } + loop { + match &self.state { + State::Eof => return Ok(0), + State::Failed(error) => return Err(error.io_error()), + State::Reading => {} + } + let count = self.pipe.as_mut().expect("reading state must own its pipe").read(buffer)?; + if count != 0 { + return Ok(count); + } + let result = match self.result.as_ref().expect("reading state must own its result receiver").recv() { + Ok(result) => result, + Err(_) => return Err(self.disconnected()), + }; + self.finish(result)?; + } + } +} + +impl Drop for Reader { + fn drop(&mut self) { + self.pipe.take(); + self.result.take(); + let _ = self.join_worker(); + } +} + +fn read_error(error: Error) -> io::Error { + match error { + Error::Io(error) => error, + Error::InvalidConfiguration(message) => io::Error::new(io::ErrorKind::InvalidInput, message), + error => io::Error::new(io::ErrorKind::InvalidData, error), + } +} diff --git a/src/stream.rs b/src/stream.rs new file mode 100644 index 0000000..a9f2981 --- /dev/null +++ b/src/stream.rs @@ -0,0 +1,73 @@ +use std::path::Path; + +use crate::{DecodeOptions, DecodeProgress, Error, OutputSink, Result, gzip, lz4}; + +/// A compression or archive format recognized by fbz. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Format { + Bzip2, + Gzip, + Lz4, + Zip, +} + +impl Format { + /// Detect a format from stream magic, falling back to the path extension. + pub fn detect(path: impl AsRef, data: &[u8]) -> Result { + if let Some(format) = Self::from_magic(data) { + return Ok(format); + } + if let Some(format) = Self::from_path(path.as_ref()) { + return Ok(format); + } + Err(Error::UnsupportedFormat(format!( + "cannot determine compression format for {}; expected a bzip2, gzip, LZ4, or ZIP extension or magic", + path.as_ref().display() + ))) + } + + /// Infer a format from a recognized filename extension. + pub fn from_path(path: &Path) -> Option { + match path.extension()?.to_str()?.to_ascii_lowercase().as_str() { + "bz2" | "bzip2" | "tbz" | "tbz2" => Some(Self::Bzip2), + "gz" | "gzip" | "tgz" => Some(Self::Gzip), + "lz4" => Some(Self::Lz4), + "zip" => Some(Self::Zip), + _ => None, + } + } + + /// Infer a format from its leading magic bytes. + pub fn from_magic(data: &[u8]) -> Option { + if data.starts_with(b"BZh") { + Some(Self::Bzip2) + } else if data.starts_with(&[0x1f, 0x8b]) { + Some(Self::Gzip) + } else if data.starts_with(&[0x04, 0x22, 0x4d, 0x18]) + || data.get(..4).is_some_and(|magic| (0x184d_2a50..=0x184d_2a5f).contains(&u32::from_le_bytes(magic.try_into().unwrap()))) + { + Some(Self::Lz4) + } else if data.starts_with(b"PK\x03\x04") || data.starts_with(b"PK\x05\x06") || data.starts_with(b"PK\x07\x08") { + Some(Self::Zip) + } else { + None + } + } +} + +/// Decode a single-stream format through the shared owned-chunk output path. +#[doc(hidden)] +pub fn decode_stream_to_sink_with_progress( + format: Format, + data: &[u8], + output: &mut impl OutputSink, + options: DecodeOptions, + progress: impl FnMut(DecodeProgress), +) -> Result<()> { + match format { + Format::Bzip2 => crate::decompress_to_sink_with_progress(data, output, options, progress), + Format::Gzip => gzip::decompress_to_sink_with_options_and_progress(data, output, options, progress).map(|_| ()), + Format::Lz4 => lz4::decompress_to_sink_with_options_and_progress(data, output, options, progress).map(|_| ()), + Format::Zip => Err(Error::UnsupportedFormat("ZIP archives do not have one decoded byte stream".into())), + } +} diff --git a/tests/reader.rs b/tests/reader.rs new file mode 100644 index 0000000..b9fa3d2 --- /dev/null +++ b/tests/reader.rs @@ -0,0 +1,106 @@ +use std::{ + fs, + io::{Read, Write}, + path::Path, +}; + +use crabz2::{Level, compress}; +use fbz::{DecodeOptions, Error, Reader}; +use flate2::{Compression, write::GzEncoder}; + +fn gzip(data: &[u8]) -> Vec { + let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6)); + encoder.write_all(data).unwrap(); + encoder.finish().unwrap() +} + +fn read(path: &Path, options: DecodeOptions) -> std::io::Result> { + let mut reader = Reader::open(path, options).unwrap(); + let mut output = Vec::new(); + reader.read_to_end(&mut output)?; + Ok(output) +} + +#[test] +fn reader_detects_stream_formats_without_prevalidation() { + let directory = tempfile::tempdir().unwrap(); + let plain = b"streaming reader format detection ".repeat(20_000); + let cases = [("bzip-magic.data", compress(&plain, Level::FASTEST)), ("gzip-magic.bz2", gzip(&plain))]; + for (name, encoded) in cases { + let path = directory.path().join(name); + fs::write(&path, encoded).unwrap(); + assert_eq!(read(&path, DecodeOptions { threads: 2, ..DecodeOptions::default() }).unwrap(), plain); + } + + let corrupt = directory.path().join("corrupt.gz"); + fs::write(&corrupt, b"not a gzip stream").unwrap(); + let mut reader = Reader::open(corrupt, DecodeOptions::default()).unwrap(); + assert_eq!(reader.read(&mut [0; 1]).unwrap_err().kind(), std::io::ErrorKind::InvalidData); +} + +#[test] +fn late_checksum_failure_is_an_error_after_the_decoded_prefix() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("late-error.gz"); + let plain = b"validated only at the trailer ".repeat(40_000); + let mut encoded = gzip(&plain); + let trailer_crc = encoded.len() - 8; + encoded[trailer_crc] ^= 1; + fs::write(&path, encoded).unwrap(); + + let mut reader = Reader::open(path, DecodeOptions { threads: 2, ..DecodeOptions::default() }).unwrap(); + let mut output = Vec::new(); + let error = reader.read_to_end(&mut output).unwrap_err(); + assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); + assert_eq!(output, plain); + assert_eq!(reader.read(&mut [0; 1]).unwrap_err().kind(), std::io::ErrorKind::InvalidData); +} + +#[test] +fn bzip2_failure_does_not_turn_a_valid_prefix_into_eof() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("concatenated.bz2"); + let first = b"complete first bzip2 stream ".repeat(10_000); + let mut encoded = compress(&first, Level::FASTEST); + let second_start = encoded.len(); + encoded.extend(compress(b"corrupt second stream", Level::FASTEST)); + encoded[second_start + 10] ^= 1; + fs::write(&path, encoded).unwrap(); + + let mut reader = Reader::open(path, DecodeOptions { threads: 2, ..DecodeOptions::default() }).unwrap(); + let mut output = Vec::new(); + assert_eq!(reader.read_to_end(&mut output).unwrap_err().kind(), std::io::ErrorKind::InvalidData); + assert!(output.starts_with(&first)); +} + +#[test] +fn dropping_under_backpressure_joins_the_decoder() { + let directory = tempfile::tempdir().unwrap(); + let inputs = [("early-drop.gz", gzip(&vec![7; 4 * 1024 * 1024]), 1), ("early-drop.bz2", compress(&vec![9; 2 * 1024 * 1024], Level::FASTEST), 4)]; + for (name, encoded, threads) in inputs { + let path = directory.path().join(name); + fs::write(&path, encoded).unwrap(); + let mut reader = Reader::open(path, DecodeOptions { threads, ..DecodeOptions::default() }).unwrap(); + assert_eq!(reader.read(&mut [0; 1]).unwrap(), 1); + drop(reader); + } +} + +#[test] +fn reader_rejects_non_stream_archives_and_bad_options_at_open() { + fn assert_send() {} + assert_send::(); + + let directory = tempfile::tempdir().unwrap(); + let zip = directory.path().join("archive.zip"); + fs::write(&zip, b"PK\x05\x06\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0").unwrap(); + assert!(matches!(Reader::open(&zip, DecodeOptions::default()), Err(Error::UnsupportedFormat(_)))); + + let lz4 = directory.path().join("data.lz4"); + fs::write(&lz4, [0x04, 0x22, 0x4d, 0x18]).unwrap(); + assert!(matches!(Reader::open(&lz4, DecodeOptions::default()), Err(Error::UnsupportedFormat(_)))); + + let gzip_path = directory.path().join("data.gz"); + fs::write(&gzip_path, gzip(b"data")).unwrap(); + assert!(matches!(Reader::open(gzip_path, DecodeOptions { memory_limit: 1, ..DecodeOptions::default() }), Err(Error::InvalidConfiguration(_)))); +} diff --git a/tests/reader_perf.rs b/tests/reader_perf.rs new file mode 100644 index 0000000..4775464 --- /dev/null +++ b/tests/reader_perf.rs @@ -0,0 +1,53 @@ +use std::{ + io::{self, Read}, + path::{Path, PathBuf}, + time::{Duration, Instant}, +}; + +use fbz::{DecodeOptions, Reader, Source, decompress_to_writer, gzip}; + +const DECODED_LEN: u64 = 84_423_012; + +fn fixture(extension: &str) -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join(format!("meta/simplewiki-first-5pct.xml.{extension}")) +} + +fn reader_run(path: &Path) -> (Duration, u64) { + let started = Instant::now(); + let mut reader = Reader::open(path, DecodeOptions::default()).unwrap(); + let decoded = io::copy(&mut reader, &mut io::sink()).unwrap(); + (started.elapsed(), decoded) +} + +fn writer_run(path: &Path, gzip_input: bool) -> Duration { + let started = Instant::now(); + let source = Source::open(path).unwrap(); + if gzip_input { + gzip::decompress_to_writer_with_options(source.as_slice(), &mut io::sink(), DecodeOptions::default()).unwrap(); + } else { + decompress_to_writer(source.as_slice(), &mut io::sink(), DecodeOptions::default()).unwrap(); + } + started.elapsed() +} + +#[test] +#[ignore = "local single-run Reader latency and writer-path comparison"] +fn reader_writer_comparison() { + for (extension, gzip_input) in [("bz2", false), ("gz", true)] { + let path = fixture(extension); + let first_started = Instant::now(); + let mut reader = Reader::open(&path, DecodeOptions::default()).unwrap(); + let mut first = [0]; + reader.read_exact(&mut first).unwrap(); + let first_byte = first_started.elapsed(); + drop(reader); + + let (reader_time, decoded) = reader_run(&path); + let writer_time = writer_run(&path, gzip_input); + assert_eq!(decoded, DECODED_LEN); + println!( + "{extension}: first byte {first_byte:?}, Reader {reader_time:?}, writer {writer_time:?}, ratio {:.3}x", + reader_time.as_secs_f64() / writer_time.as_secs_f64() + ); + } +}