From e7686e102d04fd7ca61178e64bf2ec1809893a87 Mon Sep 17 00:00:00 2001 From: Jeremy Howard Date: Mon, 24 Aug 2026 07:31:39 +1000 Subject: [PATCH] Fix parallel scheduling across concatenated bzip2 streams --- Cargo.toml | 2 +- DEV.md | 55 ++++++++- README.md | 32 ++++- python/fastbz2/__init__.py | 18 ++- src/bin/fastbz2.rs | 10 +- src/decode.rs | 232 ++++++++++++++++++++++++++++--------- src/decoder.rs | 23 +++- tests/wiki_perf.rs | 43 +++++++ 8 files changed, 340 insertions(+), 75 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 93aae18..6ff64bf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,7 @@ rayon = "1.12.0" tempfile = "3.27.0" [dev-dependencies] -crabz2 = "0.4.0" +crabz2 = { version = "0.4.0", features = ["parallel"] } libbz2-rs-sys = { version = "0.2.5", default-features = false, features = ["std"] } [features] diff --git a/DEV.md b/DEV.md index ca7eb40..38a32ba 100644 --- a/DEV.md +++ b/DEV.md @@ -26,6 +26,8 @@ The current scanner deliberately does not treat 48-bit marker matches or later ` The decoder remains independent of files, threads, Python, and the CLI. Parallel scanning/decoding and indexed seeking are layered over it. Native workers never call Python. Large offsets use explicit 64-bit bit/byte types, and speculative block-marker hits are accepted only when they form an exact stream chain with valid block and combined stream CRCs. +Parallel decoding uses a rolling candidate queue rather than stopping at stream boundaries or waiting for fixed batches. Workers reserve the maximum possible decoded block size before starting; once a block finishes, that conservative reservation shrinks to its actual output size and is released when ordered validation consumes or rejects it. Thus the `memory_limit` bounds speculative decoded output while short multistream inputs can keep the worker pool busy. The 1 GiB default admits one worst-case block per worker on the primary 18-core machine. + The production 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. Add narrowly scoped unsafe or architecture-specific SIMD only after profiling; `libbz2-rs-sys` remains the dev-only differential oracle. ## Commands @@ -51,7 +53,7 @@ The same test binary contains a warmed end-to-end performance regression gate ca Legacy randomized blocks produced by bzip2 versions before 0.9.5 are intentionally unsupported. Supporting that obsolete format would add complexity to the production decoder for data that is not realistically encountered today. -### Local SimpleWiki benchmark +### Local Wikipedia benchmarks `tests/wiki_perf.rs` contains a release-only local benchmark that is skipped by default. Its git-ignored fixture is a single bzip2 stream containing the first 84,423,012 decoded bytes (about 5%) of SimpleWiki, stored at `meta/simplewiki-first-5pct.xml.bz2`. Run it with: @@ -69,6 +71,57 @@ cargo test --release --test wiki_perf simplewiki_full -- --ignored --nocapture The full test streams decoded bytes into a counting sink and validates every block and stream CRC. Stream headers found during the structural scan remain speculative: only the exact byte-aligned header following a validated end-of-stream marker starts another stream. +The 1,000-stream enwiki comparison has a separate ignored test for each implementation and mode so a changed decoder can be measured without rerunning unchanged baselines. Each test reads the compressed fixture before starting its single timed decode, with no warmups or repeats: + +```bash +cargo test --release --test wiki_perf enwiki_first_1000_fastbz2_parallel -- --ignored --exact --nocapture +cargo test --release --test wiki_perf enwiki_first_1000_crabz2_parallel -- --ignored --exact --nocapture +cargo test --release --test wiki_perf enwiki_first_1000_fastbz2_serial -- --ignored --exact --nocapture +cargo test --release --test wiki_perf enwiki_first_1000_crabz2_serial -- --ignored --exact --nocapture +``` + +It compares fastbz2 and crabz2 in parallel and serial modes. Set `FASTBZ2_THREADS` to give both parallel decoders an explicit thread count. Each implementation validates the bzip2 CRCs; the benchmark also checks the exact decoded length. + +The decoded lengths and the 5% BLAKE3 in `tests/wiki_perf.rs` are acceptance values, not parameters used by the decoder. They prevent a truncated decode from appearing artificially fast without reading a separate multi-gigabyte reference during each timed run. Regenerating a fixture requires independently validating it and updating the corresponding acceptance value. + +#### Regenerating the fixtures + +Set `wiki` to the checkout containing the Wikimedia dumps: + +```bash +wiki=/path/to/wiki2md +``` + +Recreate the SimpleWiki fixtures from its validated compressed and decoded files: + +```bash +head -c 84423012 "$wiki/data/simplewiki-latest-pages-articles.xml" | bzip2 -9c > meta/simplewiki-first-5pct.xml.bz2 +ln -s "$wiki/data/simplewiki-latest-pages-articles.xml.bz2" meta/simplewiki-full.xml.bz2 +``` + +Run the 5% test to obtain and verify its decoded length and BLAKE3. Obtain the full decoded length with `stat`; update `FIVE_PERCENT_LEN`, `FIVE_PERCENT_BLAKE3`, or `FULL_LEN` only when intentionally changing a fixture. + +The enwiki dump is multistream. Extract its unique compressed stream offsets from the official index, whose first page-bearing stream begins after an unindexed initial stream at byte zero: + +```bash +bzcat "$wiki/data/enwiki-latest-pages-articles-multistream-index.txt.bz2" \ + | awk -F: '!seen[$1]++ {print $1}' > "$wiki/meta/enwiki-multistream-offsets.txt" +boundary=$(sed -n '1000p' "$wiki/meta/enwiki-multistream-offsets.txt") +head -c "$boundary" "$wiki/data/enwiki-latest-pages-articles-multistream.xml.bz2" \ + > "$wiki/data/enwiki-first-1000-streams.xml.bz2" +ln -s "$wiki/data/enwiki-first-1000-streams.xml.bz2" meta/enwiki-first-1000-streams.xml.bz2 +``` + +Line 1000 is the start of stream 1001 because byte zero is stream 1 and is absent from the page index. Thus `[0, boundary)` contains exactly 1,000 complete bzip2 streams. For the 2026-08-01 dump, `boundary` is `654362682` and the decoded bzip2 payload length is `2715335085`. + +To create the separately useful well-formed parser fixture, append only the XML root close after decoding; those 13 bytes are deliberately excluded from `ENWIKI_1000_LEN`: + +```bash +fastbz2 decode "$wiki/data/enwiki-first-1000-streams.xml.bz2" -o "$wiki/data/enwiki-first-1000-streams.xml" +printf '\n' >> "$wiki/data/enwiki-first-1000-streams.xml" +xmllint --stream --noout "$wiki/data/enwiki-first-1000-streams.xml" +``` + ## Platforms CI tests and builds Linux on x86-64 and ARM64, and macOS on ARM64. macOS Intel remains best-effort and should not add implementation complexity. Keep the core portable: no required mmap, custom allocator, `io_uring`, assembly, or native-endian parsing. Platform-specific positional I/O belongs behind a small source abstraction. diff --git a/README.md b/README.md index a474253..73cb826 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,37 @@ Fast parallel and indexed bzip2 decompression for Rust and Python. The performance floor is end-to-end decompression within 20% of the maintained pure-Rust `libbz2-rs-sys` decoder on a representative corpus. Portable, SIMD-friendly Rust comes first; architecture-specific SIMD is added only when profiles justify it. The much larger Simple English Wikipedia dump is used for local throughput measurements. -The implementation includes a safe structural scanner, an in-repo decoder with a tuned 12-bit Huffman lookup table, CRC-validated block decoding, bounded parallel scheduling, persistent indexes, a native CLI, and a seekable Python file API. Marker scans remain speculative until decoding establishes an exact stream chain and validates block and combined-stream CRCs. +The implementation includes a safe structural scanner, an in-repo decoder with a tuned 12-bit Huffman lookup table, CRC-validated block decoding, memory-bounded rolling parallel scheduling, persistent indexes, a native CLI, and a seekable Python file API. Marker scans remain speculative until decoding establishes an exact stream chain and validates block and combined-stream CRCs. + +## Performance + +These are single local release-mode runs on the primary Apple Silicon development machine, using 18 workers for parallel rows. They are observations rather than statistically aggregated benchmarks. Modes are stated per row because a streaming sink, an in-memory `Vec`, and a CLI pipeline have different allocation and I/O costs; compare rows using the same method most directly. + +Full Simple English Wikipedia (`338 MB` compressed, `1,688,460,257` bytes decoded): + +| Decoder | Mode | Seconds | +|---|---:|---:| +| fastbz2 | parallel, 18 threads, streaming sink | 2.244 | +| crabz2 0.4.0 | parallel | 4.460 | +| bzip2 | serial CLI | 20.310 | +| pbzip2 1.1.13 | CLI | 20.240 | +| libbz2-rs 0.2.5 | serial, in process | 20.700 | +| fastbz2 | serial, in process | 21.279 | + +The first 1,000 streams of English Wikipedia (`654,362,682` bytes compressed, `2,715,335,085` bytes decoded, 99,853 pages) exercise scheduling across many short concatenated streams: + +| Decoder | Mode | Seconds | +|---|---:|---:| +| crabz2 0.4.0 | parallel, in process | 3.815 | +| fastbz2 | parallel, 18 threads, in process | 3.881 | +| fastbz2 | serial, in process | 37.198 | +| crabz2 0.4.0 | serial, in process | 40.602 | +| pbzip2 1.1.13 | 18-thread CLI + byte comparison | 88.080 | +| bzip2 | serial CLI + byte comparison | 92.960 | + +The CLI rows in the second table stream 2.5 GB through `cmp` against the validated XML, so their absolute times are not directly comparable with the in-process rows. DEV “Local Wikipedia benchmarks” documents exact fixture generation and commands. + +Homebrew `pbzip2` 1.1.13 could not safely decompress the complete 26,668,484,995-byte English Wikipedia multistream dump on this machine. It segfaulted, and repeated attempts produced divergent and truncated plaintext. Its successful 1,000-stream result above does not establish full-file reliability. ```python import bz2 diff --git a/python/fastbz2/__init__.py b/python/fastbz2/__init__.py index f7a82c9..e8a12df 100644 --- a/python/fastbz2/__init__.py +++ b/python/fastbz2/__init__.py @@ -1,11 +1,10 @@ from collections import namedtuple -import io -import os +import io, os from pathlib import Path from ._core import BadBzip2File, _IndexedReader, __version__, _build_index, _decompress, _scan, _test, bz2_crc32 -DEFAULT_MEMORY_LIMIT = 512 * 1024 * 1024 +DEFAULT_MEMORY_LIMIT = 1024 * 1024 * 1024 DEFAULT_CACHE_LIMIT = 64 * 1024 * 1024 StreamHeaderCandidate = namedtuple("StreamHeaderCandidate", "byte_offset block_size_100k") @@ -26,17 +25,16 @@ def scan(data: bytes) -> ScanResult: return ScanResult(streams, blocks, stream_ends) def decompress(data: bytes, *, threads=0, memory_limit=DEFAULT_MEMORY_LIMIT) -> bytes: - """Decompress and fully CRC-validate one or more concatenated bzip2 streams.""" + "Decompress and fully CRC-validate one or more concatenated bzip2 streams." return _decompress(data, threads, memory_limit) class IndexedBzip2File(io.RawIOBase): - """Seekable binary reader backed by a validated bzip2 block index.""" + "Seekable binary reader backed by a validated bzip2 block index." def __init__(self, source, *, threads=0, index=None, memory_limit=DEFAULT_MEMORY_LIMIT, cache_limit=DEFAULT_CACHE_LIMIT): super().__init__() if isinstance(source, (bytes, bytearray, memoryview)): - if index is not None and not isinstance(index, (bytes, bytearray, memoryview)): - index = Path(index).read_bytes() + if index is not None and not isinstance(index, (bytes, bytearray, memoryview)): index = Path(index).read_bytes() self._reader = _IndexedReader.from_bytes(bytes(source), threads, memory_limit, index, cache_limit) else: if index is not None and isinstance(index, (bytes, bytearray, memoryview)): @@ -78,11 +76,11 @@ def close(self): super().close() def open(source, *, threads=0, index=None, memory_limit=DEFAULT_MEMORY_LIMIT, cache_limit=DEFAULT_CACHE_LIMIT): - """Open a path or bytes object as a seekable bzip2 binary file.""" + "Open a path or bytes object as a seekable bzip2 binary file." return IndexedBzip2File(source, threads=threads, index=index, memory_limit=memory_limit, cache_limit=cache_limit) def build_index(source, path=None, *, threads=0, memory_limit=DEFAULT_MEMORY_LIMIT) -> bytes: - """Fully validate *source* and return its source-bound binary block index.""" + "Fully validate *source* and return its source-bound binary block index." if isinstance(source, (bytes, bytearray, memoryview)): reader = _IndexedReader.from_bytes(bytes(source), threads, memory_limit, None, DEFAULT_CACHE_LIMIT) encoded = reader.index_bytes() @@ -91,7 +89,7 @@ def build_index(source, path=None, *, threads=0, memory_limit=DEFAULT_MEMORY_LIM return encoded def test(source, *, threads=0, memory_limit=DEFAULT_MEMORY_LIMIT): - """Fully decode and CRC-validate *source*, returning ``None`` on success.""" + "Fully decode and CRC-validate *source*, returning ``None`` on success." if isinstance(source, (bytes, bytearray, memoryview)): _IndexedReader.from_bytes(bytes(source), threads, memory_limit, None, DEFAULT_CACHE_LIMIT) else: _test(os.fspath(source), threads, memory_limit) diff --git a/src/bin/fastbz2.rs b/src/bin/fastbz2.rs index 96a5952..f1908af 100644 --- a/src/bin/fastbz2.rs +++ b/src/bin/fastbz2.rs @@ -27,7 +27,7 @@ enum Command { stdout: bool, #[arg(short = 'P', long, default_value_t = 0)] threads: usize, - #[arg(long, default_value = "512M", value_parser = parse_size)] + #[arg(long, default_value = "1G", value_parser = parse_size)] memory_limit: usize, #[arg(short, long)] force: bool, @@ -37,7 +37,7 @@ enum Command { input: String, #[arg(short = 'P', long, default_value_t = 0)] threads: usize, - #[arg(long, default_value = "512M", value_parser = parse_size)] + #[arg(long, default_value = "1G", value_parser = parse_size)] memory_limit: usize, }, /// Build a validated, source-bound block index. @@ -47,7 +47,7 @@ enum Command { output: Option, #[arg(short = 'P', long, default_value_t = 0)] threads: usize, - #[arg(long, default_value = "512M", value_parser = parse_size)] + #[arg(long, default_value = "1G", value_parser = parse_size)] memory_limit: usize, #[arg(short, long)] force: bool, @@ -57,7 +57,7 @@ enum Command { input: PathBuf, #[arg(short = 'P', long, default_value_t = 0)] threads: usize, - #[arg(long, default_value = "512M", value_parser = parse_size)] + #[arg(long, default_value = "1G", value_parser = parse_size)] memory_limit: usize, }, } @@ -80,7 +80,7 @@ fn run(cli: Cli) -> fastbz2::Result<()> { if stdout || output.is_none() { return decode_stdin(&mut io::stdout().lock()); } - return atomic_write(output.as_ref().unwrap(), force, |writer| decode_stdin(writer)); + return atomic_write(output.as_ref().unwrap(), force, decode_stdin); } let input_path = Path::new(&input); let source = Source::open(input_path)?; diff --git a/src/decode.rs b/src/decode.rs index f6e312b..6983006 100644 --- a/src/decode.rs +++ b/src/decode.rs @@ -1,19 +1,22 @@ -use std::{collections::HashMap, io::Write, sync::Arc, thread}; +use std::{ + collections::HashMap, + io::Write, + sync::{Arc, Condvar, Mutex, mpsc}, + thread, +}; -use rayon::{ThreadPool, ThreadPoolBuilder, prelude::*}; +use rayon::{ThreadPool, ThreadPoolBuilder}; use crate::format::scan_with_pool; -use crate::{ - BlockCandidate, BlockIndex, DecodeError, EndCandidate, Error, Index, MAX_DECODED_BLOCK, Result, StreamIndex, combine_stream_crc, decode_block, decoder, -}; +use crate::{BlockCandidate, BlockIndex, DecodeError, EndCandidate, Error, Index, MAX_DECODED_BLOCK, Result, StreamIndex, combine_stream_crc, decoder}; -pub const DEFAULT_MEMORY_LIMIT: usize = 512 * 1024 * 1024; +pub const DEFAULT_MEMORY_LIMIT: usize = 1024 * 1024 * 1024; #[derive(Clone, Copy, Debug)] pub struct DecodeOptions { /// Zero selects the process's available parallelism. pub threads: usize, - /// Maximum decoded bytes held in the reorder window. + /// Maximum decoded bytes reserved for in-flight and completed speculative blocks. pub memory_limit: usize, } @@ -92,6 +95,36 @@ pub fn decode_to_writer(data: &[u8], output: &mut impl Write, options: DecodeOpt return Err(Error::InvalidConfiguration("input contains too many speculative markers".into())); } + let Some(pool) = pool else { + let mut candidates = SerialCandidates { data, markers: &markers }; + return assemble(data, output, &markers, &mut candidates); + }; + let jobs: Vec<_> = markers + .iter() + .enumerate() + .filter_map(|(marker_index, marker)| match marker { + Marker::Block(block) => Some(CandidateJob { marker_index, start_bit: block.bit_offset, expected_crc: block.expected_crc }), + Marker::End(_) => None, + }) + .collect(); + let work = WorkQueue::new(options.memory_limit); + let (sender, receiver) = mpsc::channel(); + + thread::scope(|scope| { + let worker = scope.spawn(|| { + pool.broadcast(|_| worker_loop(data, &jobs, &work, &sender)); + }); + let result = { + let mut candidates = ParallelCandidates { receiver, ready: HashMap::new(), work: &work }; + assemble(data, output, &markers, &mut candidates) + }; + work.cancel(); + worker.join().map_err(|_| Error::InvalidConfiguration("parallel decoder worker panicked".into()))?; + result + }) +} + +fn assemble(data: &[u8], output: &mut impl Write, markers: &[Marker], candidates: &mut impl Candidates) -> Result { let mut blocks = Vec::new(); let mut streams = Vec::new(); let mut decoded_offset = 0_u64; @@ -104,11 +137,7 @@ pub fn decode_to_writer(data: &[u8], output: &mut impl Write, options: DecodeOpt let stream_decoded_start = decoded_offset; let mut combined_crc = 0_u32; let mut current_bit = header_byte.checked_mul(8).and_then(|bit| bit.checked_add(32)).ok_or_else(offset_overflow)?; - let mut marker_index = marker_at(&markers, current_bit)?; - let max_block = (usize::from(level) * 100_000 / 5 * 259 + 4).max(1); - let batch_size = (options.memory_limit / max_block).max(1).min(threads.saturating_mul(2).max(1)); - let mut ready: HashMap>> = HashMap::new(); - + let mut marker_index = marker_at(markers, current_bit)?; loop { match &markers[marker_index] { Marker::End(end) => { @@ -130,26 +159,14 @@ pub fn decode_to_writer(data: &[u8], output: &mut impl Write, options: DecodeOpt break; } Marker::Block(block) => { - if ready.is_empty() { - ready = decode_batch(data, &markers, marker_index, batch_size, level, pool.as_deref()); + let decoded = candidates.take(marker_index)?; + if decoded.block_len > usize::from(level) * 100_000 { + return Err(Error::Decode { bit_offset: block.bit_offset, source: DecodeError::BlockOverflow }); } - let mut end_index = marker_index + 1; - let decoded = match ready.remove(&marker_index) { - Some(Ok(decoded)) => decoded, - Some(Err(first_error)) => { - ready.clear(); - match retry_merged(data, &markers, marker_index, level) { - Ok((decoded, found_end)) => { - end_index = found_end; - decoded - } - Err(_) => return Err(first_error), - } - } - None => return Err(required_marker(current_bit)), - }; - output.write_all(&decoded)?; - let decoded_len = decoded.len() as u64; + let end_index = marker_at(markers, decoded.end_bit)?; + candidates.discard_before(end_index); + output.write_all(&decoded.output)?; + let decoded_len = decoded.output.len() as u64; blocks.push(BlockIndex { compressed_start_bit: block.bit_offset, compressed_end_bit: markers[end_index].bit_offset(), @@ -162,7 +179,6 @@ pub fn decode_to_writer(data: &[u8], output: &mut impl Write, options: DecodeOpt combined_crc = combine_stream_crc(combined_crc, block.expected_crc); current_bit = markers[end_index].bit_offset(); marker_index = end_index; - ready.retain(|&index, _| index >= marker_index); } } } @@ -178,32 +194,129 @@ pub fn decode_to_writer(data: &[u8], output: &mut impl Write, options: DecodeOpt Ok(Index::new(data, decoded_offset, streams, blocks)) } -fn decode_batch(data: &[u8], markers: &[Marker], start: usize, limit: usize, level: u8, pool: Option<&ThreadPool>) -> HashMap>> { - let end = (start + limit).min(markers.len().saturating_sub(1)); - let jobs: Vec<_> = (start..end) - .take_while(|&index| matches!(markers[index], Marker::Block(_))) - .map(|index| { - let Marker::Block(block) = &markers[index] else { unreachable!() }; - (index, block.bit_offset, markers[index + 1].bit_offset(), block.expected_crc) - }) - .collect(); - let decode = || jobs.par_iter().map(|&(index, start_bit, end_bit, crc)| (index, decode_block(data, start_bit, end_bit, level, crc))).collect(); - match pool { - Some(pool) => pool.install(decode), - None => jobs.into_iter().map(|(index, start_bit, end_bit, crc)| (index, decode_block(data, start_bit, end_bit, level, crc))).collect(), +trait Candidates { + fn take(&mut self, marker_index: usize) -> Result; + fn discard_before(&mut self, marker_index: usize); +} + +struct SerialCandidates<'a> { + data: &'a [u8], + markers: &'a [Marker], +} + +impl Candidates for SerialCandidates<'_> { + fn take(&mut self, marker_index: usize) -> Result { + let Marker::Block(block) = &self.markers[marker_index] else { return Err(required_marker(self.markers[marker_index].bit_offset())) }; + decoder::decode_candidate(self.data, block.bit_offset, block.expected_crc) + } + + fn discard_before(&mut self, _marker_index: usize) {} +} + +#[derive(Clone, Copy)] +struct CandidateJob { + marker_index: usize, + start_bit: u64, + expected_crc: u32, +} + +struct WorkState { + next: usize, + reserved: usize, + cancelled: bool, +} + +struct WorkQueue { + limit: usize, + state: Mutex, + wake: Condvar, +} + +impl WorkQueue { + fn new(limit: usize) -> Self { + Self { limit, state: Mutex::new(WorkState { next: 0, reserved: 0, cancelled: false }), wake: Condvar::new() } + } + + fn next(&self, job_count: usize) -> Option { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + while state.reserved > self.limit - MAX_DECODED_BLOCK && state.next < job_count && !state.cancelled { + state = self.wake.wait(state).unwrap_or_else(|error| error.into_inner()); + } + if state.cancelled || state.next >= job_count { + return None; + } + let next = state.next; + state.next += 1; + state.reserved += MAX_DECODED_BLOCK; + Some(next) + } + + fn complete(&self, decoded_len: usize) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.reserved -= MAX_DECODED_BLOCK - decoded_len; + self.wake.notify_all(); + } + + fn retire(&self, decoded_len: usize) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.reserved -= decoded_len; + self.wake.notify_all(); + } + + fn cancel(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.cancelled = true; + self.wake.notify_all(); } } -fn retry_merged(data: &[u8], markers: &[Marker], start: usize, level: u8) -> Result<(Vec, usize)> { - let Marker::Block(block) = &markers[start] else { return Err(required_marker(markers[start].bit_offset())) }; - let mut last_error = None; - for (end, marker) in markers.iter().enumerate().skip(start + 2) { - match decode_block(data, block.bit_offset, marker.bit_offset(), level, block.expected_crc) { - Ok(decoded) => return Ok((decoded, end)), - Err(error) => last_error = Some(error), +type CandidateResult = (usize, Result); + +fn worker_loop(data: &[u8], jobs: &[CandidateJob], work: &WorkQueue, sender: &mpsc::Sender) { + while let Some(index) = work.next(jobs.len()) { + let job = jobs[index]; + let result = decoder::decode_candidate(data, job.start_bit, job.expected_crc); + let decoded_len = candidate_len(&result); + work.complete(decoded_len); + if sender.send((job.marker_index, result)).is_err() { + work.retire(decoded_len); + return; + } + } +} + +fn candidate_len(result: &Result) -> usize { + result.as_ref().map_or(0, |decoded| decoded.output.len()) +} + +struct ParallelCandidates<'a> { + receiver: mpsc::Receiver, + ready: HashMap>, + work: &'a WorkQueue, +} + +impl Candidates for ParallelCandidates<'_> { + fn take(&mut self, marker_index: usize) -> Result { + while !self.ready.contains_key(&marker_index) { + let (index, result) = self.receiver.recv().map_err(|_| Error::InvalidConfiguration("parallel decoder stopped early".into()))?; + if index < marker_index { + self.work.retire(candidate_len(&result)); + } else { + self.ready.insert(index, result); + } + } + let result = self.ready.remove(&marker_index).unwrap(); + self.work.retire(candidate_len(&result)); + result + } + + fn discard_before(&mut self, marker_index: usize) { + let stale: Vec<_> = self.ready.keys().copied().filter(|&index| index < marker_index).collect(); + for index in stale { + let result = self.ready.remove(&index).unwrap(); + self.work.retire(candidate_len(&result)); } } - Err(last_error.unwrap_or_else(|| required_marker(block.bit_offset))) } fn thread_pool(threads: usize) -> Result>> { @@ -292,6 +405,19 @@ mod tests { assert!(index.blocks.len() >= 3); } + #[test] + fn bounded_scheduler_crosses_many_short_streams() { + let mut compressed = Vec::new(); + let mut expected = Vec::new(); + for stream in 0..32 { + let plain = patterned(4_000 + stream * 31); + compressed.extend_from_slice(&compress(&plain, Level::BEST)); + expected.extend_from_slice(&plain); + } + let options = DecodeOptions { threads: 4, memory_limit: MAX_DECODED_BLOCK * 2 }; + assert_eq!(decompress(&compressed, options).unwrap(), expected); + } + #[test] fn rejects_stream_crc_corruption() { let mut compressed = compress(b"integrity matters", Level::BEST); diff --git a/src/decoder.rs b/src/decoder.rs index fc5b2e2..1a28de2 100644 --- a/src/decoder.rs +++ b/src/decoder.rs @@ -191,12 +191,18 @@ struct Decoder { tt: Vec, } +pub(crate) struct DecodedCandidate { + pub output: Vec, + pub end_bit: u64, + pub block_len: usize, +} + impl Decoder { fn new() -> Self { Self { tt: Vec::new() } } - fn block(&mut self, bits: &mut Bits<'_>, level: u8, expected_crc: Option) -> Result<(Vec, u32)> { + fn block(&mut self, bits: &mut Bits<'_>, level: u8, expected_crc: Option) -> Result<(Vec, u32, usize)> { let block_offset = bits.position().saturating_sub(48); let stored_crc = bits.read(32)?; if expected_crc.is_some_and(|expected| expected != stored_crc) { @@ -366,8 +372,17 @@ impl Decoder { if bz2_crc32(&output) != stored_crc { return Err(decode_error(block_offset, DecodeError::CrcMismatch)); } - Ok((output, stored_crc)) + Ok((output, stored_crc, block_len)) + } +} + +pub(crate) fn decode_candidate(data: &[u8], start_bit: u64, expected_crc: u32) -> Result { + let mut bits = Bits::at(data, start_bit)?; + if bits.magic()? != BLOCK_MAGIC { + 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 }) } pub(crate) fn decode_block(data: &[u8], start_bit: u64, end_bit: u64, level: u8, expected_crc: u32) -> Result> { @@ -378,7 +393,7 @@ pub(crate) fn decode_block(data: &[u8], start_bit: u64, end_bit: u64, level: u8, if bits.magic()? != BLOCK_MAGIC { return Err(decode_error(start_bit, DecodeError::InvalidMagic)); } - let (output, _) = Decoder::new().block(&mut bits, level, Some(expected_crc))?; + let (output, _, _) = Decoder::new().block(&mut bits, level, Some(expected_crc))?; if bits.position() != end_bit { return Err(decode_error(bits.position(), DecodeError::InvalidBlock)); } @@ -406,7 +421,7 @@ pub(crate) fn decode_serial(data: &[u8], output: &mut impl Write) -> Result<()> let marker_offset = bits.position(); match bits.magic()? { BLOCK_MAGIC => { - let (block, crc) = decoder.block(&mut bits, level, None)?; + let (block, crc, _) = decoder.block(&mut bits, level, None)?; output.write_all(&block)?; combined_crc = combine_stream_crc(combined_crc, crc); } diff --git a/tests/wiki_perf.rs b/tests/wiki_perf.rs index 4c65fd8..7ddb44e 100644 --- a/tests/wiki_perf.rs +++ b/tests/wiki_perf.rs @@ -10,6 +10,7 @@ use fastbz2::{DecodeOptions, Source, decompress, decompress_to_writer}; const FIVE_PERCENT_LEN: usize = 84_423_012; const FIVE_PERCENT_BLAKE3: &str = "69f41f28dc8ac74509d368c6aaec02f3cdf891c9da4ccf8caf625687dcd61908"; const FULL_LEN: u64 = 1_688_460_257; +const ENWIKI_1000_LEN: usize = 2_715_335_085; #[derive(Default)] struct CountingSink(u64); @@ -72,3 +73,45 @@ fn simplewiki_full() { let elapsed = timed_fastbz2(&path, options.threads); eprintln!("fastbz2 ({} threads): {elapsed:.3?}", options.resolved_threads()); } + +fn timed_vec(name: &str, decode: impl FnOnce() -> Vec) { + let start = Instant::now(); + let decoded = decode(); + let elapsed = start.elapsed(); + assert_eq!(decoded.len(), ENWIKI_1000_LEN); + eprintln!("{name}: {elapsed:.3?}"); +} + +fn enwiki_fixture() -> (Vec, usize) { + let encoded = fs::read(corpus_path("enwiki-first-1000-streams.xml.bz2")).unwrap(); + (encoded, requested_threads()) +} + +#[test] +#[ignore = "local enwiki multistream performance comparison"] +fn enwiki_first_1000_fastbz2_parallel() { + let (encoded, threads) = enwiki_fixture(); + let options = DecodeOptions { threads, ..DecodeOptions::default() }; + timed_vec(&format!("fastbz2 parallel ({} threads)", options.resolved_threads()), || decompress(&encoded, options).unwrap()); +} + +#[test] +#[ignore = "local enwiki multistream performance comparison"] +fn enwiki_first_1000_crabz2_parallel() { + let (encoded, threads) = enwiki_fixture(); + timed_vec("crabz2 parallel", || crabz2::decompress_parallel(&encoded, if threads == 0 { None } else { Some(threads) }).unwrap()); +} + +#[test] +#[ignore = "local enwiki multistream performance comparison"] +fn enwiki_first_1000_fastbz2_serial() { + let (encoded, _) = enwiki_fixture(); + timed_vec("fastbz2 serial", || decompress(&encoded, DecodeOptions { threads: 1, ..DecodeOptions::default() }).unwrap()); +} + +#[test] +#[ignore = "local enwiki multistream performance comparison"] +fn enwiki_first_1000_crabz2_serial() { + let (encoded, _) = enwiki_fixture(); + timed_vec("crabz2 serial", || crabz2::decompress(&encoded).unwrap()); +}