gatling/src/parwrite.rs
//! Parallel-writer gatling — the shared extract/decode/**write** fan-out.
//!
//! The streaming [`crate::gatling`] engine and the fork-join
//! [`crate::gatling_forkjoin`] both funnel their decoded output through a
//! **single ordered sink** (one collector thread draining an mpsc, one writer).
//! For codecs whose output units are *independent files or independent regions
//! of one pre-sized stream*, that single writer is the last un-parallelised
//! stage — the serial drain that leaves cores idle while one thread copies and
//! writes gigabytes.
//!
//! This module removes it. It is the parallel-writer sibling of
//! [`gatling_for_each`](crate::gatling_forkjoin::gatling_for_each): the same
//! no-barrier, self-dispatching (atomic-cursor) worker pool built on
//! `std::thread::scope`, but each worker also **writes its own output** — there
//! is no collector and no writer thread on the critical path.
//!
//! Two shapes, both format-agnostic (the caller plugs in the decode):
//!
//! - [`extract_entries_unordered`] — N workers self-dispatch **independent
//! entries** (ZIP/JAR members). Each worker `pread`s its entry's compressed
//! bytes from the shared input (thread-safe positional reads, no shared file
//! offset), decodes, and writes its **own output file**. Entries are
//! independent, so the file writes never conflict and no ordering is needed.
//!
//! - [`write_segments_positional`] — N workers self-dispatch **independent
//! segments of one output stream** (concatenated gzip members). The caller
//! pre-computes each segment's byte offset (a prefix-sum of known output
//! sizes) and pre-sizes the output file; each worker decodes its segment and
//! `pwrite`s it at the segment's offset. Exact byte order is preserved with
//! no ordered collector and no serial concatenation.
//!
//! Both reuse each worker's scratch buffers across the units it claims
//! (zero-alloc hot loop after warm-up), never copy input into the pool (the
//! decode reads directly from the shared source), and are rayon-free: pure
//! `std::thread::scope` + one `AtomicUsize` cursor — the gatling soul.
use std::fs::{self, File};
use std::io::Write;
use std::os::unix::fs::FileExt;
use std::path::Path;
use std::sync::atomic::{AtomicUsize, Ordering};
/// Outcome of a parallel-writer run: how many units were written and how many
/// failed (a failed unit is logged to stderr and skipped, never aborting the
/// siblings).
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct WriteStats {
/// Units decoded and written successfully.
pub ok: usize,
/// Units that failed to read/decode/write (skipped, logged).
pub failed: usize,
}
/// One independent extraction unit for [`extract_entries_unordered`]: read
/// `comp_len` compressed bytes at `data_offset` from the shared input, decode,
/// and write the result to `out_dir.join(name)`.
pub struct EntryJob {
/// Byte offset of the entry's compressed payload in the shared input file.
pub data_offset: u64,
/// Compressed payload length in bytes.
pub comp_len: usize,
/// Uncompressed size (exact, from the archive central directory) — sizes the
/// decode output buffer so the hot loop never reallocates.
pub out_size: usize,
/// Method-0 STORE entry: payload is copied verbatim (no decode).
pub is_store: bool,
/// Relative output path (joined onto `out_dir`).
pub name: String,
}
/// Extract independent entries across a no-barrier, self-dispatching worker pool
/// — **each worker decodes AND writes its own output file**, with no ordered
/// collector and no single writer thread.
///
/// `input` is the shared archive file; workers read each entry's compressed
/// bytes with `read_exact_at` (`pread`), which is thread-safe and uses no shared
/// file offset, so N workers read disjoint regions concurrently. `decode` turns
/// one entry's compressed bytes into its uncompressed bytes: it is handed the
/// compressed slice, the entry's exact `out_size`, and a **reused** output
/// buffer to fill (cleared first). STORE entries (`is_store`) bypass `decode`
/// and are copied verbatim.
///
/// Returns the [`WriteStats`]. A unit that fails to read/decode/write is logged
/// to stderr and skipped; its siblings keep running.
pub fn extract_entries_unordered<D>(
input: &File,
jobs: &[EntryJob],
out_dir: &Path,
n_workers: usize,
decode: D,
) -> WriteStats
where
D: Fn(&[u8], usize, &mut Vec<u8>) -> Result<(), String> + Sync,
{
let n = jobs.len();
if n == 0 {
return WriteStats::default();
}
let workers = n_workers.max(1).min(n);
let next = AtomicUsize::new(0);
let ok = AtomicUsize::new(0);
let failed = AtomicUsize::new(0);
let next_ref = &next;
let ok_ref = &ok;
let failed_ref = &failed;
let decode_ref = &decode;
std::thread::scope(|s| {
for _ in 0..workers {
s.spawn(move || {
// Per-worker scratch, reused across every entry this worker
// claims — zero-alloc after the first (largest) entry.
let mut comp: Vec<u8> = Vec::new();
let mut out: Vec<u8> = Vec::new();
loop {
let i = next_ref.fetch_add(1, Ordering::Relaxed);
if i >= n {
break;
}
let job = &jobs[i];
// Read the compressed payload via pread (no shared offset).
if comp.len() < job.comp_len {
comp.resize(job.comp_len, 0);
}
if let Err(e) = input.read_exact_at(&mut comp[..job.comp_len], job.data_offset)
{
eprintln!("parwrite: read {}: {e}", job.name);
failed_ref.fetch_add(1, Ordering::Relaxed);
continue;
}
let comp_slice = &comp[..job.comp_len];
// Decode (or copy for STORE) into the reused output buffer.
out.clear();
if job.is_store {
out.extend_from_slice(comp_slice);
} else if let Err(e) = decode_ref(comp_slice, job.out_size, &mut out) {
eprintln!("parwrite: decode {}: {e}", job.name);
failed_ref.fetch_add(1, Ordering::Relaxed);
continue;
}
// Write our own output file (independent entry → no ordering).
let dest = out_dir.join(&job.name);
if let Some(parent) = dest.parent() {
if let Err(e) = fs::create_dir_all(parent) {
eprintln!("parwrite: mkdir {}: {e}", parent.display());
failed_ref.fetch_add(1, Ordering::Relaxed);
continue;
}
}
match File::create(&dest).and_then(|mut f| f.write_all(&out)) {
Ok(()) => {
ok_ref.fetch_add(1, Ordering::Relaxed);
}
Err(e) => {
eprintln!("parwrite: write {}: {e}", dest.display());
failed_ref.fetch_add(1, Ordering::Relaxed);
}
}
}
});
}
});
WriteStats {
ok: ok.load(Ordering::Relaxed),
failed: failed.load(Ordering::Relaxed),
}
}
/// Write independent segments of **one** output stream at pre-computed offsets,
/// across a no-barrier, self-dispatching worker pool — **each worker decodes
/// AND `pwrite`s its own segment**, with no ordered collector and no serial
/// concatenation.
///
/// The caller has already laid out the stream: `offsets[i]` is the byte offset
/// in `out` where segment `i`'s decoded bytes belong (a prefix-sum of the
/// segments' known output sizes), and `out` has been pre-sized to the total
/// (e.g. via `set_len`). Each worker claims the next segment index, calls
/// `decode(i, &mut buf)` to fill its reused buffer, and `write_all_at`
/// (`pwrite`) writes it at `offsets[i]` — positional writes to disjoint regions
/// never conflict, so exact stream byte-order is preserved without any ordering
/// on the critical path.
///
/// `decode` should verify the decoded length matches what the caller assumed
/// for `offsets` (returning `Err` on mismatch); a failed segment is logged and
/// skipped, and the caller can detect `failed > 0` to fall back to a safe path.
pub fn write_segments_positional<D>(
out: &File,
n_segments: usize,
offsets: &[u64],
n_workers: usize,
decode: D,
) -> WriteStats
where
D: Fn(usize, &mut Vec<u8>) -> Result<(), String> + Sync,
{
assert_eq!(
offsets.len(),
n_segments,
"offsets must have one entry per segment"
);
if n_segments == 0 {
return WriteStats::default();
}
let workers = n_workers.max(1).min(n_segments);
let next = AtomicUsize::new(0);
let ok = AtomicUsize::new(0);
let failed = AtomicUsize::new(0);
let next_ref = &next;
let ok_ref = &ok;
let failed_ref = &failed;
let decode_ref = &decode;
std::thread::scope(|s| {
for _ in 0..workers {
s.spawn(move || {
let mut buf: Vec<u8> = Vec::new();
loop {
let i = next_ref.fetch_add(1, Ordering::Relaxed);
if i >= n_segments {
break;
}
buf.clear();
if let Err(e) = decode_ref(i, &mut buf) {
eprintln!("parwrite: segment {i} decode: {e}");
failed_ref.fetch_add(1, Ordering::Relaxed);
continue;
}
match out.write_all_at(&buf, offsets[i]) {
Ok(()) => {
ok_ref.fetch_add(1, Ordering::Relaxed);
}
Err(e) => {
eprintln!("parwrite: segment {i} pwrite: {e}");
failed_ref.fetch_add(1, Ordering::Relaxed);
}
}
}
});
}
});
WriteStats {
ok: ok.load(Ordering::Relaxed),
failed: failed.load(Ordering::Relaxed),
}
}
/// Write already-decoded slices of **one** output stream at pre-computed offsets,
/// across a no-barrier, self-dispatching worker pool — **zero-copy parallel
/// `pwrite`**, no ordered collector and no serial concatenation.
///
/// The two-phase sibling of [`write_segments_positional`]: use it when the
/// segments are already materialised (e.g. members decoded in parallel by a
/// prior [`gatling_for_each`](crate::gatling_forkjoin::gatling_for_each) pass)
/// and their exact sizes — hence `offsets` — are known. `out` must be pre-sized
/// to the total (`set_len`). Each worker claims the next slice index and
/// `write_all_at` (`pwrite`)s `slices[i]` at `offsets[i]`; the slices are never
/// copied through the pool, and positional writes to disjoint regions preserve
/// exact stream byte-order with no ordering on the critical path.
pub fn write_slices_positional(
out: &File,
slices: &[&[u8]],
offsets: &[u64],
n_workers: usize,
) -> WriteStats {
assert_eq!(
offsets.len(),
slices.len(),
"offsets must have one entry per slice"
);
let n = slices.len();
if n == 0 {
return WriteStats::default();
}
let workers = n_workers.max(1).min(n);
let next = AtomicUsize::new(0);
let ok = AtomicUsize::new(0);
let failed = AtomicUsize::new(0);
let next_ref = &next;
let ok_ref = &ok;
let failed_ref = &failed;
std::thread::scope(|s| {
for _ in 0..workers {
s.spawn(move || {
loop {
let i = next_ref.fetch_add(1, Ordering::Relaxed);
if i >= n {
break;
}
match out.write_all_at(slices[i], offsets[i]) {
Ok(()) => {
ok_ref.fetch_add(1, Ordering::Relaxed);
}
Err(e) => {
eprintln!("parwrite: slice {i} pwrite: {e}");
failed_ref.fetch_add(1, Ordering::Relaxed);
}
}
}
});
}
});
WriteStats {
ok: ok.load(Ordering::Relaxed),
failed: failed.load(Ordering::Relaxed),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Read;
#[test]
fn extract_writes_every_entry_concurrently() {
// Build an "archive": three payloads concatenated in one temp file.
let dir = std::env::temp_dir().join(format!("parwrite-test-{}", std::process::id()));
let _ = fs::remove_dir_all(&dir);
fs::create_dir_all(&dir).unwrap();
let arc_path = dir.join("archive.bin");
let payloads: Vec<Vec<u8>> = (0..3)
.map(|k| (0..1000).map(|i| (i as u8).wrapping_add(k as u8)).collect())
.collect();
let mut arc = Vec::new();
let mut jobs = Vec::new();
for (k, p) in payloads.iter().enumerate() {
jobs.push(EntryJob {
data_offset: arc.len() as u64,
comp_len: p.len(),
out_size: p.len(),
is_store: true, // verbatim copy path
name: format!("out_{k}.bin"),
});
arc.extend_from_slice(p);
}
fs::write(&arc_path, &arc).unwrap();
let f = File::open(&arc_path).unwrap();
let out_dir = dir.join("out");
let stats = extract_entries_unordered(&f, &jobs, &out_dir, 4, |_c, _s, _o| Ok(()));
assert_eq!(stats.ok, 3);
assert_eq!(stats.failed, 0);
for (k, p) in payloads.iter().enumerate() {
let got = fs::read(out_dir.join(format!("out_{k}.bin"))).unwrap();
assert_eq!(&got, p, "entry {k} bytes must round-trip");
}
let _ = fs::remove_dir_all(&dir);
}
#[test]
fn positional_writes_reassemble_in_order() {
let dir = std::env::temp_dir().join(format!("parwrite-pos-{}", std::process::id()));
let _ = fs::remove_dir_all(&dir);
fs::create_dir_all(&dir).unwrap();
let out_path = dir.join("stream.bin");
// Four segments of differing sizes; workers may finish out of order but
// the positional writes must reassemble the exact concatenation.
let segs: Vec<Vec<u8>> = vec![
vec![1u8; 100],
vec![2u8; 250],
vec![3u8; 50],
vec![4u8; 400],
];
let mut offsets = Vec::new();
let mut total = 0u64;
for s in &segs {
offsets.push(total);
total += s.len() as u64;
}
let out = File::create(&out_path).unwrap();
out.set_len(total).unwrap();
let stats = write_segments_positional(&out, segs.len(), &offsets, 4, |i, buf| {
buf.extend_from_slice(&segs[i]);
Ok(())
});
assert_eq!(stats.ok, segs.len());
assert_eq!(stats.failed, 0);
drop(out);
let mut got = Vec::new();
File::open(&out_path)
.unwrap()
.read_to_end(&mut got)
.unwrap();
let mut want = Vec::new();
for s in &segs {
want.extend_from_slice(s);
}
assert_eq!(
got, want,
"positional segments must reassemble in stream order"
);
let _ = fs::remove_dir_all(&dir);
}
}