use std::io::Read;
use std::sync::{mpsc, Arc, Mutex};
use anyhow::Result;
#[path = "gatling_io.rs"]
pub mod io;
#[path = "gatling_ordered.rs"]
pub mod ordered;
#[path = "revolver/mod.rs"]
pub mod revolver;
pub struct Split<S> {
pub segments: Vec<S>,
pub consumed: usize,
}
pub trait Codec: Sync {
type Seg: Send + Copy + 'static;
fn split(&self, data: &[u8], n_workers: usize, is_last: bool) -> Option<Split<Self::Seg>>;
fn decode(&self, data: &[u8], seg: &Self::Seg) -> Vec<u8>;
}
pub trait Sink: Send {
fn safe_end(&self, assembled: &[u8], is_last: bool) -> usize;
fn process(&mut self, bytes: &[u8]) -> Result<()>;
}
pub trait TypedCodec: Sync {
type Seg: Send + Copy + 'static;
type Output: Send + 'static;
fn split(&self, data: &[u8], n_workers: usize, is_last: bool) -> Option<Split<Self::Seg>>;
fn transform(&self, data: &[u8], seg: &Self::Seg) -> Self::Output;
fn finish_worker(&self) -> Option<Self::Output> { None }
}
pub trait TypedSink<T>: Send {
fn process(&mut self, output: T, is_last: bool) -> Result<()>;
fn finish(&mut self) -> Result<()> { Ok(()) }
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum SlotFill {
#[default]
Incremental,
Big,
}
pub struct Config {
pub chunk_size: usize,
pub carry_headroom: usize,
pub ring_slots: usize,
pub initial_carry: Vec<u8>,
pub slot_fill: SlotFill,
}
struct WorkItem<S> {
chunk_id: u64,
seg_id: usize,
data_ptr: *const u8,
data_len: usize,
seg: S,
}
unsafe impl<S: Send> Send for WorkItem<S> {}
struct SegResult {
chunk_id: u64,
seg_id: usize,
output: Vec<u8>,
}
struct InFlight {
chunk_id: u64,
slot: Vec<u8>,
n_segments: usize,
results: Vec<Option<Vec<u8>>>,
done: usize,
is_last: bool,
}
enum Msg {
New(InFlight),
Result(SegResult),
}
fn fill_slot<R: Read>(
src: &mut R,
slot: &mut Vec<u8>,
carry_headroom: usize,
chunk_size: usize,
fill: SlotFill,
pending: &mut Option<u8>,
) -> (usize, bool) {
let prefill = pending.take();
let got = match fill {
SlotFill::Big => {
let slot_size = carry_headroom + chunk_size;
if slot.len() != slot_size {
slot.resize(slot_size, 0);
}
let mut got = 0usize;
if let Some(b) = prefill {
slot[carry_headroom] = b;
got = 1;
}
while got < chunk_size {
match src.read(&mut slot[carry_headroom + got..slot_size]) {
Ok(0) => break,
Ok(k) => got += k,
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(_) => break,
}
}
got
}
SlotFill::Incremental => {
const READ_BLOCK: usize = 8 * 1024 * 1024;
slot.clear();
slot.resize(carry_headroom, 0);
let mut got = 0usize;
if let Some(b) = prefill {
slot.push(b);
got = 1;
}
while got < chunk_size {
let want = READ_BLOCK.min(chunk_size - got);
let base = slot.len();
slot.resize(base + want, 0);
match src.read(&mut slot[base..base + want]) {
Ok(0) => {
slot.truncate(base);
break;
}
Ok(k) => {
slot.truncate(base + k);
got += k;
}
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
slot.truncate(base);
continue;
}
Err(_) => {
slot.truncate(base);
break;
}
}
}
got
}
};
let is_last = if got < chunk_size {
true
} else {
let mut probe = [0u8; 1];
loop {
match src.read(&mut probe) {
Ok(0) => break true,
Ok(_) => {
*pending = Some(probe[0]);
break false;
}
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
Err(_) => break true,
}
}
};
(got, is_last)
}
pub fn run<C: Codec, S: Sink>(
reader: impl Read + Send,
codec: C,
sink: &mut S,
n_workers: usize,
cfg: Config,
) -> Result<()> {
let n_workers = n_workers.max(1);
let slot_size = cfg.carry_headroom + cfg.chunk_size;
assert!(
cfg.initial_carry.len() <= cfg.carry_headroom,
"initial_carry {} exceeds carry_headroom {}",
cfg.initial_carry.len(),
cfg.carry_headroom
);
let codec_ref = &codec;
let sink_ref: &mut S = sink;
std::thread::scope(|s| -> Result<()> {
let (slot_return_tx, slot_return_rx) = mpsc::sync_channel::<Vec<u8>>(cfg.ring_slots);
let (filled_tx, filled_rx) =
mpsc::sync_channel::<(Vec<u8>, usize, bool)>(cfg.ring_slots);
let chunk_size = cfg.chunk_size;
let carry_headroom = cfg.carry_headroom;
let ring_slots = cfg.ring_slots;
let slot_fill = cfg.slot_fill;
s.spawn(move || {
use std::sync::mpsc::TryRecvError;
let mut src = reader;
let mut allocated = 0usize;
let mut pending: Option<u8> = None;
loop {
let mut slot = match slot_return_rx.try_recv() {
Ok(sl) => sl,
Err(TryRecvError::Empty) if allocated < ring_slots => {
allocated += 1;
Vec::with_capacity(slot_size)
}
Err(TryRecvError::Empty) => match slot_return_rx.recv() {
Ok(sl) => sl,
Err(_) => break,
},
Err(TryRecvError::Disconnected) => break,
};
let (got, is_last) =
fill_slot(&mut src, &mut slot, carry_headroom, chunk_size, slot_fill, &mut pending);
if filled_tx.send((slot, got, is_last)).is_err() {
break;
}
if is_last {
break;
}
}
});
let (work_tx, work_rx) = mpsc::sync_channel::<WorkItem<C::Seg>>(n_workers * 2);
let (collector_tx, collector_rx) = mpsc::sync_channel::<Msg>(n_workers * 4);
let work_rx = Arc::new(Mutex::new(work_rx));
for _ in 0..n_workers {
let work_rx = Arc::clone(&work_rx);
let collector_tx = collector_tx.clone();
s.spawn(move || {
loop {
let item = {
let rx = work_rx.lock().expect("work_rx lock");
match rx.recv() {
Ok(item) => item,
Err(_) => break,
}
};
let data = unsafe {
std::slice::from_raw_parts(item.data_ptr, item.data_len)
};
let output = codec_ref.decode(data, &item.seg);
let _ = collector_tx.send(Msg::Result(SegResult {
chunk_id: item.chunk_id,
seg_id: item.seg_id,
output,
}));
}
});
}
drop(work_rx);
let inflight_tx = collector_tx.clone();
drop(collector_tx);
let slot_return_for_collector = slot_return_tx.clone();
let collector = s.spawn(move || -> Result<()> {
let mut in_flight: Vec<InFlight> = Vec::new();
let mut next_id: u64 = 0;
let mut carry: Vec<u8> = Vec::new();
while let Ok(msg) = collector_rx.recv() {
match msg {
Msg::New(slot) => in_flight.push(slot),
Msg::Result(r) => {
for slot in in_flight.iter_mut() {
if slot.chunk_id == r.chunk_id {
slot.results[r.seg_id] = Some(r.output);
slot.done += 1;
break;
}
}
}
}
loop {
let Some(idx) = in_flight.iter().position(|s| s.chunk_id == next_id) else {
break;
};
if in_flight[idx].done < in_flight[idx].n_segments {
break;
}
let mut done = in_flight.remove(idx);
let is_last = done.is_last;
let mut buf = std::mem::take(&mut carry);
for seg in done.results.drain(..) {
if let Some(bytes) = seg {
buf.extend_from_slice(&bytes);
}
}
let safe = sink_ref.safe_end(&buf, is_last);
if safe == 0 {
carry = buf;
} else {
carry = buf[safe..].to_vec();
buf.truncate(safe);
sink_ref.process(&buf)?;
}
slot_return_for_collector.send(done.slot).ok();
next_id += 1;
}
}
if !carry.is_empty() {
sink_ref.process(&carry)?;
}
Ok(())
});
let carry_floor = cfg.initial_carry.len();
let mut carry: Vec<u8> = cfg.initial_carry;
let mut chunk_id: u64 = 0;
let max_carry = cfg.chunk_size.saturating_mul(16).max(1 << 30);
for (mut slot, read_len, is_last) in filled_rx.iter() {
if carry.len() > max_carry {
return Err(anyhow::anyhow!(
"gatling: no split boundary found across {} MiB of accumulated carry — \
the input stream is truncated or corrupt (e.g. an incomplete / zero-padded download)",
carry.len() >> 20
));
}
if read_len == 0 && carry.len() <= carry_floor {
slot_return_tx.send(slot).ok();
break;
}
let carry_len = carry.len();
let (buf, data_start, data_end) = if carry_len <= carry_headroom {
let data_start = carry_headroom - carry_len;
slot[data_start..carry_headroom].copy_from_slice(&carry);
(slot, data_start, carry_headroom + read_len)
} else {
let mut combined = Vec::with_capacity(carry_len + read_len);
combined.extend_from_slice(&carry);
combined.extend_from_slice(&slot[carry_headroom..carry_headroom + read_len]);
slot_return_tx.send(slot).ok();
let end = combined.len();
(combined, 0, end)
};
let data = &buf[data_start..data_end];
let split = match codec_ref.split(data, n_workers, is_last) {
Some(sp) => sp,
None => {
carry.clear();
carry.extend_from_slice(data);
slot_return_tx.send(buf).ok();
if is_last {
break;
}
continue;
}
};
carry.clear();
carry.extend_from_slice(&data[split.consumed..]);
let n_segments = split.segments.len();
let inflight = InFlight {
chunk_id,
slot: buf,
n_segments,
results: (0..n_segments).map(|_| None).collect(),
done: 0,
is_last,
};
let data_ptr = inflight.slot[data_start..].as_ptr();
let data_len = data_end - data_start;
inflight_tx
.send(Msg::New(inflight))
.map_err(|_| anyhow::anyhow!("collector closed"))?;
for (seg_id, seg) in split.segments.into_iter().enumerate() {
work_tx
.send(WorkItem { chunk_id, seg_id, data_ptr, data_len, seg })
.map_err(|_| anyhow::anyhow!("work channel closed"))?;
}
chunk_id += 1;
if is_last {
break;
}
}
drop(work_tx);
drop(inflight_tx);
drop(slot_return_tx);
collector.join().expect("collector panicked")?;
Ok(())
})
}
struct TypedWorkItem<S> {
chunk_id: u64,
seg_id: usize,
data_ptr: *const u8,
data_len: usize,
seg: S,
}
unsafe impl<S: Send> Send for TypedWorkItem<S> {}
struct TypedSegResult<T> {
chunk_id: u64,
seg_id: usize,
output: T,
}
struct TypedInFlight<T> {
chunk_id: u64,
slot: Vec<u8>,
n_segments: usize,
results: Vec<Option<T>>,
done: usize,
is_last: bool,
}
enum TypedMsg<T> {
New(TypedInFlight<T>),
Result(TypedSegResult<T>),
Tail(T),
}
pub fn run_typed<C: TypedCodec, S: TypedSink<C::Output>>(
reader: impl Read + Send,
codec: C,
sink: &mut S,
n_workers: usize,
cfg: Config,
) -> Result<()> {
let n_workers = n_workers.max(1);
let slot_size = cfg.carry_headroom + cfg.chunk_size;
assert!(
cfg.initial_carry.len() <= cfg.carry_headroom,
"initial_carry {} exceeds carry_headroom {}",
cfg.initial_carry.len(),
cfg.carry_headroom
);
let codec_ref = &codec;
let sink_ref: &mut S = sink;
std::thread::scope(|s| -> Result<()> {
let (slot_return_tx, slot_return_rx) = mpsc::sync_channel::<Vec<u8>>(cfg.ring_slots);
let (filled_tx, filled_rx) =
mpsc::sync_channel::<(Vec<u8>, usize, bool)>(cfg.ring_slots);
let chunk_size = cfg.chunk_size;
let carry_headroom = cfg.carry_headroom;
let ring_slots = cfg.ring_slots;
let slot_fill = cfg.slot_fill;
s.spawn(move || {
use std::sync::mpsc::TryRecvError;
let mut src = reader;
let mut allocated = 0usize;
let mut pending: Option<u8> = None;
loop {
let mut slot = match slot_return_rx.try_recv() {
Ok(sl) => sl,
Err(TryRecvError::Empty) if allocated < ring_slots => {
allocated += 1;
Vec::with_capacity(slot_size)
}
Err(TryRecvError::Empty) => match slot_return_rx.recv() {
Ok(sl) => sl,
Err(_) => break,
},
Err(TryRecvError::Disconnected) => break,
};
let (got, is_last) =
fill_slot(&mut src, &mut slot, carry_headroom, chunk_size, slot_fill, &mut pending);
if filled_tx.send((slot, got, is_last)).is_err() {
break;
}
if is_last {
break;
}
}
});
let (work_tx, work_rx) =
mpsc::sync_channel::<TypedWorkItem<C::Seg>>(n_workers * 2);
let (collector_tx, collector_rx) =
mpsc::sync_channel::<TypedMsg<C::Output>>(n_workers * 4);
let work_rx = Arc::new(Mutex::new(work_rx));
for _ in 0..n_workers {
let work_rx = Arc::clone(&work_rx);
let collector_tx = collector_tx.clone();
s.spawn(move || {
loop {
let item = {
let rx = work_rx.lock().expect("work_rx lock");
match rx.recv() {
Ok(item) => item,
Err(_) => break,
}
};
let data = unsafe {
std::slice::from_raw_parts(item.data_ptr, item.data_len)
};
let output = codec_ref.transform(data, &item.seg);
let _ = collector_tx.send(TypedMsg::Result(TypedSegResult {
chunk_id: item.chunk_id,
seg_id: item.seg_id,
output,
}));
}
if let Some(tail) = codec_ref.finish_worker() {
let _ = collector_tx.send(TypedMsg::Tail(tail));
}
});
}
drop(work_rx);
let inflight_tx = collector_tx.clone();
drop(collector_tx);
let slot_return_for_collector = slot_return_tx.clone();
let collector = s.spawn(move || -> Result<()> {
let mut in_flight: Vec<TypedInFlight<C::Output>> = Vec::new();
let mut next_id: u64 = 0;
let mut tails: Vec<C::Output> = Vec::new();
while let Ok(msg) = collector_rx.recv() {
match msg {
TypedMsg::New(slot) => in_flight.push(slot),
TypedMsg::Result(r) => {
for slot in in_flight.iter_mut() {
if slot.chunk_id == r.chunk_id {
slot.results[r.seg_id] = Some(r.output);
slot.done += 1;
break;
}
}
}
TypedMsg::Tail(out) => tails.push(out),
}
loop {
let Some(idx) = in_flight.iter().position(|s| s.chunk_id == next_id) else {
break;
};
if in_flight[idx].done < in_flight[idx].n_segments {
break;
}
let mut done = in_flight.remove(idx);
let is_last = done.is_last;
let n = done.results.len();
for (i, output) in done.results.drain(..).enumerate() {
if let Some(val) = output {
let last_segment = is_last && i == n - 1;
sink_ref.process(val, last_segment)?;
}
}
slot_return_for_collector.send(done.slot).ok();
next_id += 1;
}
}
for tail in tails {
sink_ref.process(tail, false)?;
}
sink_ref.finish()?;
Ok(())
});
let carry_floor = cfg.initial_carry.len();
let mut carry: Vec<u8> = cfg.initial_carry;
let mut chunk_id: u64 = 0;
let max_carry = chunk_size.saturating_mul(16).max(1 << 30);
for (mut slot, read_len, is_last) in filled_rx.iter() {
if carry.len() > max_carry {
return Err(anyhow::anyhow!(
"gatling: no split boundary found across {} MiB of accumulated carry — \
the input stream is truncated or corrupt (e.g. an incomplete / zero-padded download)",
carry.len() >> 20
));
}
if read_len == 0 && carry.len() <= carry_floor {
slot_return_tx.send(slot).ok();
break;
}
let carry_len = carry.len();
let (buf, data_start, data_end) = if carry_len <= carry_headroom {
let data_start = carry_headroom - carry_len;
slot[data_start..carry_headroom].copy_from_slice(&carry);
(slot, data_start, carry_headroom + read_len)
} else {
let mut combined = Vec::with_capacity(carry_len + read_len);
combined.extend_from_slice(&carry);
combined.extend_from_slice(&slot[carry_headroom..carry_headroom + read_len]);
slot_return_tx.send(slot).ok();
let end = combined.len();
(combined, 0, end)
};
let data = &buf[data_start..data_end];
let split = match codec_ref.split(data, n_workers, is_last) {
Some(sp) => sp,
None => {
carry.clear();
carry.extend_from_slice(data);
slot_return_tx.send(buf).ok();
if is_last {
break;
}
continue;
}
};
carry.clear();
carry.extend_from_slice(&data[split.consumed..]);
let n_segments = split.segments.len();
let inflight = TypedInFlight {
chunk_id,
slot: buf,
n_segments,
results: (0..n_segments).map(|_| None).collect(),
done: 0,
is_last,
};
let data_ptr = inflight.slot[data_start..].as_ptr();
let data_len = data_end - data_start;
inflight_tx
.send(TypedMsg::New(inflight))
.map_err(|_| anyhow::anyhow!("collector closed"))?;
for (seg_id, seg) in split.segments.into_iter().enumerate() {
work_tx
.send(TypedWorkItem { chunk_id, seg_id, data_ptr, data_len, seg })
.map_err(|_| anyhow::anyhow!("work channel closed"))?;
}
chunk_id += 1;
if is_last {
break;
}
}
drop(work_tx);
drop(inflight_tx);
drop(slot_return_tx);
collector.join().expect("collector panicked")?;
Ok(())
})
}