gatling/gatling.rs
1//! Gatling — generic no-barrier worker-pool streaming engine.
2//!
3//! Extracted from katana-osm's `xml_to_pbf_bz2` (the pipeline behind its planet /
4//! Sweden benchmarks). The engine carries **no** codec / format knowledge — the
5//! caller plugs in a codec (how to split and transform) and a sink (how to consume
6//! the output in order).
7//!
8//! # Two modes
9//!
10//! ## Byte mode ([`run`])
11//! Use when workers produce raw bytes that need reassembly (e.g. bz2 decode → XML).
12//! - [`Codec`] splits input into segments and decodes each → `Vec<u8>`.
13//! - [`Sink`] receives assembled, boundary-aligned byte slices.
14//! - Collector concatenates decoded segments, finds a safe boundary via
15//! [`Sink::safe_end`], then calls [`Sink::process`].
16//!
17//! ```text
18//! Reader → Main(split) → N Workers(decode→bytes) → Collector(assemble+safe_end) → Sink
19//! ```
20//!
21//! ## Typed mode ([`run_typed`])
22//! Use when workers produce fully-processed output that needs no reassembly
23//! (e.g. VTD-parse raw XML → parsed records, or any transform where `split` already
24//! guarantees logical boundaries between segments).
25//! - [`TypedCodec`] splits input and transforms each segment → `Self::Output`.
26//! - [`TypedSink`] receives each segment's output individually, in stream order.
27//! - Collector is trivial: just delivers outputs in order, no assembly, no carry.
28//!
29//! ```text
30//! Reader → Main(split) → N Workers(transform→T) → Collector(forward in order) → TypedSink
31//! ```
32//!
33//! # Properties (both modes)
34//! - **Zero-copy input:** bytes stay in the slot; workers get a raw pointer.
35//! - **No barrier:** slot N+1 is split and dispatched while slot N is still processing.
36//! - **Carry in headroom:** unconsumed tail from `split` is prepended to next chunk.
37
38use std::io::Read;
39use std::sync::{mpsc, Arc, Mutex};
40
41use anyhow::Result;
42
43/// Async-I/O sibling of this engine — a no-barrier, bounded, ordered async task
44/// pool for **I/O-bound** work (S3 PUT/GET fan-out, etc.), as opposed to the
45/// **CPU-bound** thread pool [`run`]/[`run_typed`] of this module. See
46/// [`io`](mod@io) for the substrate-difference rationale.
47#[path = "gatling_io.rs"]
48pub mod io;
49
50/// Item-oriented sibling of this engine — an **ordered streaming-sink** gatling
51/// for a labelled-discrete-entry pipeline (e.g. znippy-compress). The caller stays
52/// the producer (yields `(Label, Input)` items, pulled lazily under backpressure);
53/// N workers run a `map`; outputs stream to a [`ordered::OrderedSink`] in producer
54/// order through a **bounded** reorder buffer. See [`ordered`](mod@ordered).
55#[path = "gatling_ordered.rs"]
56pub mod ordered;
57
58/// Zero-allocation **slot pool** variant — the buffer substrate the streaming
59/// [`run`]/[`run_typed`] engine fans over: one reader fills pre-allocated
60/// fixed-size slots, N workers borrow each as `&[u8]` (zero-copy) via
61/// [`revolver::get_chunk_slice`]. Folded in from the former top-level
62/// `chunk_revolver` module (still re-exported crate-root as `chunk_revolver` for
63/// API compatibility). See [`revolver`](mod@revolver).
64#[path = "revolver/mod.rs"]
65pub mod revolver;
66
67/// Result of splitting a compressed buffer into independent decode units.
68pub struct Split<S> {
69 /// One descriptor per decode unit, in stream order.
70 pub segments: Vec<S>,
71 /// Number of bytes consumed from the input slice; the remainder becomes carry.
72 pub consumed: usize,
73}
74
75/// Decompressor plug-in. `Seg` is an opaque, `Copy` per-unit descriptor (e.g. a
76/// bz2 bit range) computed by `split` and handed back to `decode`.
77pub trait Codec: Sync {
78 /// Per-decode-unit descriptor. Must be cheap to copy and `Send`.
79 type Seg: Send + Copy + 'static;
80
81 /// Find independent decode units in `data`. Return `None` when no complete
82 /// unit is available yet (the caller grows the carry and retries next chunk).
83 /// `is_last` is true for the final chunk of the stream.
84 fn split(&self, data: &[u8], n_workers: usize, is_last: bool) -> Option<Split<Self::Seg>>;
85
86 /// Decode one unit of `data` into its uncompressed bytes.
87 fn decode(&self, data: &[u8], seg: &Self::Seg) -> Vec<u8>;
88}
89
90/// Consumer plug-in. Methods run on the engine's single collector thread, so a
91/// `Sink` may safely own mutable state and fan work out to its own writer thread.
92pub trait Sink: Send {
93 /// Largest prefix length of `assembled` (decoded bytes) that ends on a safe
94 /// boundary (e.g. the end of an XML element). Return `0` if none yet; return
95 /// `assembled.len()` when `is_last` to flush everything.
96 fn safe_end(&self, assembled: &[u8], is_last: bool) -> usize;
97
98 /// Process one contiguous decoded run, already cut at a safe boundary.
99 fn process(&mut self, bytes: &[u8]) -> Result<()>;
100}
101
102// ── Typed mode traits ───────────────────────────────────────────────────────
103
104/// Transform plug-in for typed mode. Workers call [`TypedCodec::transform`] to
105/// produce an arbitrary output type per segment. `split` must ensure each segment
106/// is logically self-contained (e.g. cut at XML element boundaries) — the collector
107/// will NOT reassemble segments; each output is forwarded individually.
108pub trait TypedCodec: Sync {
109 /// Per-segment descriptor (e.g. a `(usize, usize)` byte range within the chunk).
110 type Seg: Send + Copy + 'static;
111
112 /// Fully-processed output produced by one worker for one segment.
113 type Output: Send + 'static;
114
115 /// Find segment boundaries in `data`. Each segment must be logically complete
116 /// (no partial elements spanning segments). `consumed` = bytes up to the last
117 /// safe boundary; remainder becomes carry for next chunk.
118 fn split(&self, data: &[u8], n_workers: usize, is_last: bool) -> Option<Split<Self::Seg>>;
119
120 /// Transform one segment of `data` into typed output. Called on N worker
121 /// threads in parallel. `data` is the full slot slice (carry + read);
122 /// use `seg` to index into it.
123 fn transform(&self, data: &[u8], seg: &Self::Seg) -> Self::Output;
124
125 /// Called once on each worker thread after it has processed all of its
126 /// segments, just before the thread exits. Lets a codec flush per-worker
127 /// accumulated state (e.g. a partial row group carried across segments via
128 /// thread-local storage) into the stream. The returned output, if any, is
129 /// forwarded to the sink after all in-order segment outputs. Default: none.
130 fn finish_worker(&self) -> Option<Self::Output> { None }
131}
132
133/// Consumer plug-in for typed mode. Receives each segment's output individually,
134/// in strict stream order. The collector thread calls this — fan out to a writer
135/// thread or batch internally as needed.
136pub trait TypedSink<T>: Send {
137 /// Process one segment's output. `is_last` is true for the final segment
138 /// of the entire stream (not just the chunk).
139 fn process(&mut self, output: T, is_last: bool) -> Result<()>;
140
141 /// Called once after all segments have been processed. Default is no-op.
142 fn finish(&mut self) -> Result<()> { Ok(()) }
143}
144
145/// How a slot's backing memory is grown/faulted when the reader fills it.
146///
147/// Both modes lazily allocate slots (a small input still only ever touches one
148/// slot); they differ only in how much of an *allocated* slot gets faulted.
149#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
150pub enum SlotFill {
151 /// Grow the slot in blocks to exactly the bytes read, never faulting the
152 /// unused tail. A 24 MB read into a 232 MB slot touches ~24 MB, not 232 MB.
153 /// Best for small / single-chunk / partial-tail inputs. This is the default.
154 #[default]
155 Incremental,
156 /// Size each slot to the full `carry_headroom + chunk_size` once, on first
157 /// use, then reuse it across reads with no re-zeroing and no reallocation.
158 /// The up-front fault is paid once and amortizes to ~zero over the many
159 /// full chunks of a large (planet-scale) input. Best when nearly every
160 /// chunk is full so the slot tail is never wasted.
161 Big,
162}
163
164/// Engine tuning. `slot size = carry_headroom + chunk_size`.
165pub struct Config {
166 /// Compressed bytes read per slot.
167 pub chunk_size: usize,
168 /// Bytes reserved at the head of each slot for prepended carry.
169 pub carry_headroom: usize,
170 /// Slot-pool depth (slots in flight).
171 pub ring_slots: usize,
172 /// Bytes to prepend to the first chunk (e.g. a stream header already read).
173 pub initial_carry: Vec<u8>,
174 /// Slot memory-growth strategy. See [`SlotFill`].
175 pub slot_fill: SlotFill,
176}
177
178// ── Internal message types ──────────────────────────────────────────────────
179
180struct WorkItem<S> {
181 chunk_id: u64,
182 seg_id: usize,
183 data_ptr: *const u8,
184 data_len: usize,
185 seg: S,
186}
187// SAFETY: `data_ptr` points into a slot held alive in `InFlight` by the collector
188// until every segment of that chunk has been decoded and the slot recycled. The
189// reader never overwrites a slot while it is in flight (it is removed from the
190// free pool). `S: Send` covers the descriptor payload.
191unsafe impl<S: Send> Send for WorkItem<S> {}
192
193struct SegResult {
194 chunk_id: u64,
195 seg_id: usize,
196 output: Vec<u8>,
197}
198
199struct InFlight {
200 chunk_id: u64,
201 slot: Vec<u8>,
202 n_segments: usize,
203 results: Vec<Option<Vec<u8>>>,
204 done: usize,
205 is_last: bool,
206}
207
208enum Msg {
209 New(InFlight),
210 Result(SegResult),
211}
212
213// ── Engine ──────────────────────────────────────────────────────────────────
214
215/// Fill `slot` with up to `chunk_size` bytes read from `src`, placed at offset
216/// `carry_headroom` (leaving the front free for a prepended carry). Returns the
217/// number of bytes read.
218///
219/// The [`SlotFill`] mode controls how much slot memory gets faulted:
220/// - [`SlotFill::Incremental`] grows the slot in blocks to exactly the bytes
221/// read, never touching the unused tail — the win on small / partial inputs.
222/// - [`SlotFill::Big`] sizes the slot to its full extent once and reuses it with
223/// no re-zeroing or reallocation — the up-front fault amortizes over the many
224/// full chunks of a planet-scale input.
225///
226/// For a full chunk (`got == chunk_size`) both modes touch the same region, so
227/// large reads never regress regardless of mode.
228///
229/// Returns `(bytes_read, is_last)`. End-of-stream is detected precisely: a short
230/// read means EOF was hit mid-chunk, and after a *full* chunk we peek one byte
231/// past it — an empty peek means the stream ends exactly here (this is the last
232/// chunk), while a peeked byte is stashed in `pending` as the next chunk's first
233/// data byte. This flags `is_last` correctly even when the input length is an
234/// exact multiple of `chunk_size` (where `got == chunk_size` on every read and no
235/// trailing short read ever occurs). The one-byte peek costs nothing meaningful
236/// and needs no extra slot, so it never deadlocks the pool.
237fn fill_slot<R: Read>(
238 src: &mut R,
239 slot: &mut Vec<u8>,
240 carry_headroom: usize,
241 chunk_size: usize,
242 fill: SlotFill,
243 pending: &mut Option<u8>,
244) -> (usize, bool) {
245 // A byte peeked past the previous full chunk belongs at the front of this
246 // chunk's data region.
247 let prefill = pending.take();
248 let got = match fill {
249 SlotFill::Big => {
250 let slot_size = carry_headroom + chunk_size;
251 // Pay the fault once: a fresh slot is sized here; a recycled slot is
252 // already full-length, so this is a no-op (no realloc, no re-zero).
253 if slot.len() != slot_size {
254 slot.resize(slot_size, 0);
255 }
256 let mut got = 0usize;
257 if let Some(b) = prefill {
258 slot[carry_headroom] = b;
259 got = 1;
260 }
261 while got < chunk_size {
262 match src.read(&mut slot[carry_headroom + got..slot_size]) {
263 Ok(0) => break,
264 Ok(k) => got += k,
265 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
266 Err(_) => break,
267 }
268 }
269 // Leave the slot full-length so the next reuse skips the resize.
270 got
271 }
272 SlotFill::Incremental => {
273 const READ_BLOCK: usize = 8 * 1024 * 1024;
274 slot.clear();
275 // Carry prefix: bytes [data_start..carry_headroom] are overwritten by
276 // the prepended carry before use; [0..data_start] is never read.
277 slot.resize(carry_headroom, 0);
278 let mut got = 0usize;
279 if let Some(b) = prefill {
280 slot.push(b);
281 got = 1;
282 }
283 while got < chunk_size {
284 let want = READ_BLOCK.min(chunk_size - got);
285 let base = slot.len();
286 slot.resize(base + want, 0);
287 match src.read(&mut slot[base..base + want]) {
288 Ok(0) => {
289 slot.truncate(base);
290 break;
291 }
292 Ok(k) => {
293 slot.truncate(base + k);
294 got += k;
295 }
296 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
297 slot.truncate(base);
298 continue;
299 }
300 Err(_) => {
301 slot.truncate(base);
302 break;
303 }
304 }
305 }
306 got
307 }
308 };
309
310 // Determine end-of-stream. A short read already means EOF; after a full
311 // chunk, probe one byte forward without losing it.
312 let is_last = if got < chunk_size {
313 true
314 } else {
315 let mut probe = [0u8; 1];
316 loop {
317 match src.read(&mut probe) {
318 Ok(0) => break true,
319 Ok(_) => {
320 *pending = Some(probe[0]);
321 break false;
322 }
323 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
324 Err(_) => break true,
325 }
326 }
327 };
328
329 (got, is_last)
330}
331
332/// Run the engine to completion over `reader`, driving `sink`.
333///
334/// `sink` is borrowed mutably for the duration; counts / outputs the consumer
335/// accumulates are read from it after `run` returns (e.g. via a `finish` method
336/// of your own). `codec` is shared across the decode workers.
337pub fn run<C: Codec, S: Sink>(
338 reader: impl Read + Send,
339 codec: C,
340 sink: &mut S,
341 n_workers: usize,
342 cfg: Config,
343) -> Result<()> {
344 let n_workers = n_workers.max(1);
345 let slot_size = cfg.carry_headroom + cfg.chunk_size;
346 assert!(
347 cfg.initial_carry.len() <= cfg.carry_headroom,
348 "initial_carry {} exceeds carry_headroom {}",
349 cfg.initial_carry.len(),
350 cfg.carry_headroom
351 );
352
353 let codec_ref = &codec;
354 let sink_ref: &mut S = sink;
355
356 std::thread::scope(|s| -> Result<()> {
357 // ── Slot pool (lazy) ──────────────────────────────────────────────────
358 // Slots are allocated on demand, capped at `ring_slots`. A small input only
359 // ever touches one slot, so it never pays the full `ring_slots × slot_size`
360 // up front (a 5 MB file used to zero 6 × 232 MB = 1.4 GB before reading a byte).
361 let (slot_return_tx, slot_return_rx) = mpsc::sync_channel::<Vec<u8>>(cfg.ring_slots);
362
363 // ── Reader thread ─────────────────────────────────────────────────────
364 let (filled_tx, filled_rx) =
365 mpsc::sync_channel::<(Vec<u8>, usize, bool)>(cfg.ring_slots);
366 let chunk_size = cfg.chunk_size;
367 let carry_headroom = cfg.carry_headroom;
368 let ring_slots = cfg.ring_slots;
369 let slot_fill = cfg.slot_fill;
370 s.spawn(move || {
371 use std::sync::mpsc::TryRecvError;
372 let mut src = reader;
373 let mut allocated = 0usize;
374 let mut pending: Option<u8> = None;
375 loop {
376 // Reuse a recycled slot if one is waiting; else allocate a fresh
377 // slot while under budget; else block for a slot to come back.
378 let mut slot = match slot_return_rx.try_recv() {
379 Ok(sl) => sl,
380 Err(TryRecvError::Empty) if allocated < ring_slots => {
381 allocated += 1;
382 Vec::with_capacity(slot_size)
383 }
384 Err(TryRecvError::Empty) => match slot_return_rx.recv() {
385 Ok(sl) => sl,
386 Err(_) => break,
387 },
388 Err(TryRecvError::Disconnected) => break,
389 };
390 let (got, is_last) =
391 fill_slot(&mut src, &mut slot, carry_headroom, chunk_size, slot_fill, &mut pending);
392 if filled_tx.send((slot, got, is_last)).is_err() {
393 break;
394 }
395 if is_last {
396 break;
397 }
398 }
399 });
400
401 // ── Decode workers (persistent, no barrier) ───────────────────────────
402 let (work_tx, work_rx) = mpsc::sync_channel::<WorkItem<C::Seg>>(n_workers * 2);
403 let (collector_tx, collector_rx) = mpsc::sync_channel::<Msg>(n_workers * 4);
404 let work_rx = Arc::new(Mutex::new(work_rx));
405
406 for _ in 0..n_workers {
407 let work_rx = Arc::clone(&work_rx);
408 let collector_tx = collector_tx.clone();
409 s.spawn(move || {
410 loop {
411 let item = {
412 let rx = work_rx.lock().expect("work_rx lock");
413 match rx.recv() {
414 Ok(item) => item,
415 Err(_) => break,
416 }
417 };
418 // SAFETY: the slot backing `data_ptr` is held in `InFlight`
419 // by the collector and not recycled until all its segments
420 // are decoded, so the pointer is valid for this read.
421 let data = unsafe {
422 std::slice::from_raw_parts(item.data_ptr, item.data_len)
423 };
424 let output = codec_ref.decode(data, &item.seg);
425 let _ = collector_tx.send(Msg::Result(SegResult {
426 chunk_id: item.chunk_id,
427 seg_id: item.seg_id,
428 output,
429 }));
430 }
431 });
432 }
433 drop(work_rx);
434
435 // Main keeps an inflight sender; workers keep their clones.
436 let inflight_tx = collector_tx.clone();
437 drop(collector_tx);
438
439 // ── Collector thread ──────────────────────────────────────────────────
440 let slot_return_for_collector = slot_return_tx.clone();
441 let collector = s.spawn(move || -> Result<()> {
442 let mut in_flight: Vec<InFlight> = Vec::new();
443 let mut next_id: u64 = 0;
444 let mut carry: Vec<u8> = Vec::new();
445
446 while let Ok(msg) = collector_rx.recv() {
447 match msg {
448 Msg::New(slot) => in_flight.push(slot),
449 Msg::Result(r) => {
450 for slot in in_flight.iter_mut() {
451 if slot.chunk_id == r.chunk_id {
452 slot.results[r.seg_id] = Some(r.output);
453 slot.done += 1;
454 break;
455 }
456 }
457 }
458 }
459
460 // Flush completed slots strictly in stream order.
461 loop {
462 let Some(idx) = in_flight.iter().position(|s| s.chunk_id == next_id) else {
463 break;
464 };
465 if in_flight[idx].done < in_flight[idx].n_segments {
466 break;
467 }
468
469 let mut done = in_flight.remove(idx);
470 let is_last = done.is_last;
471
472 let mut buf = std::mem::take(&mut carry);
473 for seg in done.results.drain(..) {
474 if let Some(bytes) = seg {
475 buf.extend_from_slice(&bytes);
476 }
477 }
478
479 let safe = sink_ref.safe_end(&buf, is_last);
480 if safe == 0 {
481 carry = buf;
482 } else {
483 carry = buf[safe..].to_vec();
484 buf.truncate(safe);
485 sink_ref.process(&buf)?;
486 }
487
488 slot_return_for_collector.send(done.slot).ok();
489 next_id += 1;
490 }
491 }
492
493 if !carry.is_empty() {
494 sink_ref.process(&carry)?;
495 }
496 Ok(())
497 });
498
499 // ── Main: carry + split + dispatch ────────────────────────────────────
500 // A zero-read chunk whose carry has not grown past the seed length means
501 // the stream is finished (the carry holds only the prepended header).
502 let carry_floor = cfg.initial_carry.len();
503 let mut carry: Vec<u8> = cfg.initial_carry;
504 let mut chunk_id: u64 = 0;
505 // Truncation/corruption guard — see the twin in `run_typed`. Bails when the
506 // codec can never split a growing carry (e.g. a truncated / zero-padded
507 // input) instead of accumulating to OOM.
508 let max_carry = cfg.chunk_size.saturating_mul(16).max(1 << 30);
509
510 for (mut slot, read_len, is_last) in filled_rx.iter() {
511 if carry.len() > max_carry {
512 return Err(anyhow::anyhow!(
513 "gatling: no split boundary found across {} MiB of accumulated carry — \
514 the input stream is truncated or corrupt (e.g. an incomplete / zero-padded download)",
515 carry.len() >> 20
516 ));
517 }
518 if read_len == 0 && carry.len() <= carry_floor {
519 slot_return_tx.send(slot).ok();
520 break;
521 }
522
523 // Prepend the carry so the worker sees a single contiguous
524 // `carry ++ read` slice. The common case fits in the slot's reserved
525 // front headroom (zero-copy). When a single unit spans enough chunks
526 // that the carry outgrows the headroom, fall back to a dedicated
527 // buffer sized to the actual carry — the headroom is a fast-path
528 // optimization, not a hard cap on unit size.
529 let carry_len = carry.len();
530 let (buf, data_start, data_end) = if carry_len <= carry_headroom {
531 let data_start = carry_headroom - carry_len;
532 slot[data_start..carry_headroom].copy_from_slice(&carry);
533 (slot, data_start, carry_headroom + read_len)
534 } else {
535 let mut combined = Vec::with_capacity(carry_len + read_len);
536 combined.extend_from_slice(&carry);
537 combined.extend_from_slice(&slot[carry_headroom..carry_headroom + read_len]);
538 // The read bytes are copied out; recycle the slot immediately.
539 slot_return_tx.send(slot).ok();
540 let end = combined.len();
541 (combined, 0, end)
542 };
543
544 let data = &buf[data_start..data_end];
545 let split = match codec_ref.split(data, n_workers, is_last) {
546 Some(sp) => sp,
547 None => {
548 carry.clear();
549 carry.extend_from_slice(data);
550 slot_return_tx.send(buf).ok();
551 if is_last {
552 break;
553 }
554 continue;
555 }
556 };
557
558 carry.clear();
559 carry.extend_from_slice(&data[split.consumed..]);
560
561 let n_segments = split.segments.len();
562 let inflight = InFlight {
563 chunk_id,
564 slot: buf,
565 n_segments,
566 results: (0..n_segments).map(|_| None).collect(),
567 done: 0,
568 is_last,
569 };
570
571 // Pointer into the slot before it is moved into the collector.
572 let data_ptr = inflight.slot[data_start..].as_ptr();
573 let data_len = data_end - data_start;
574
575 inflight_tx
576 .send(Msg::New(inflight))
577 .map_err(|_| anyhow::anyhow!("collector closed"))?;
578
579 for (seg_id, seg) in split.segments.into_iter().enumerate() {
580 work_tx
581 .send(WorkItem { chunk_id, seg_id, data_ptr, data_len, seg })
582 .map_err(|_| anyhow::anyhow!("work channel closed"))?;
583 }
584
585 chunk_id += 1;
586 if is_last {
587 break;
588 }
589 }
590
591 drop(work_tx);
592 drop(inflight_tx);
593 drop(slot_return_tx);
594
595 collector.join().expect("collector panicked")?;
596 Ok(())
597 })
598}
599
600// ── Typed mode internals ────────────────────────────────────────────────────
601
602struct TypedWorkItem<S> {
603 chunk_id: u64,
604 seg_id: usize,
605 data_ptr: *const u8,
606 data_len: usize,
607 seg: S,
608}
609unsafe impl<S: Send> Send for TypedWorkItem<S> {}
610
611struct TypedSegResult<T> {
612 chunk_id: u64,
613 seg_id: usize,
614 output: T,
615}
616
617struct TypedInFlight<T> {
618 chunk_id: u64,
619 slot: Vec<u8>,
620 n_segments: usize,
621 results: Vec<Option<T>>,
622 done: usize,
623 is_last: bool,
624}
625
626enum TypedMsg<T> {
627 New(TypedInFlight<T>),
628 Result(TypedSegResult<T>),
629 /// End-of-worker remainder, forwarded to the sink after all ordered outputs.
630 Tail(T),
631}
632
633// ── Typed engine ────────────────────────────────────────────────────────────
634
635/// Run the typed engine to completion. Workers call [`TypedCodec::transform`] to
636/// produce arbitrary output per segment; the collector forwards each output to
637/// `sink` in strict stream order (no byte assembly, no boundary logic).
638///
639/// Use this when `split` already guarantees logically-complete segments (e.g. cut
640/// at XML element boundaries) and workers produce fully-processed typed output
641/// (parsed records, encoded row groups, etc.).
642pub fn run_typed<C: TypedCodec, S: TypedSink<C::Output>>(
643 reader: impl Read + Send,
644 codec: C,
645 sink: &mut S,
646 n_workers: usize,
647 cfg: Config,
648) -> Result<()> {
649 let n_workers = n_workers.max(1);
650 let slot_size = cfg.carry_headroom + cfg.chunk_size;
651 assert!(
652 cfg.initial_carry.len() <= cfg.carry_headroom,
653 "initial_carry {} exceeds carry_headroom {}",
654 cfg.initial_carry.len(),
655 cfg.carry_headroom
656 );
657
658 let codec_ref = &codec;
659 let sink_ref: &mut S = sink;
660
661 std::thread::scope(|s| -> Result<()> {
662 // ── Slot pool (lazy) ──────────────────────────────────────────────────
663 let (slot_return_tx, slot_return_rx) = mpsc::sync_channel::<Vec<u8>>(cfg.ring_slots);
664
665 // ── Reader thread ─────────────────────────────────────────────────────
666 let (filled_tx, filled_rx) =
667 mpsc::sync_channel::<(Vec<u8>, usize, bool)>(cfg.ring_slots);
668 let chunk_size = cfg.chunk_size;
669 let carry_headroom = cfg.carry_headroom;
670 let ring_slots = cfg.ring_slots;
671 let slot_fill = cfg.slot_fill;
672 s.spawn(move || {
673 use std::sync::mpsc::TryRecvError;
674 let mut src = reader;
675 let mut allocated = 0usize;
676 let mut pending: Option<u8> = None;
677 loop {
678 let mut slot = match slot_return_rx.try_recv() {
679 Ok(sl) => sl,
680 Err(TryRecvError::Empty) if allocated < ring_slots => {
681 allocated += 1;
682 Vec::with_capacity(slot_size)
683 }
684 Err(TryRecvError::Empty) => match slot_return_rx.recv() {
685 Ok(sl) => sl,
686 Err(_) => break,
687 },
688 Err(TryRecvError::Disconnected) => break,
689 };
690 let (got, is_last) =
691 fill_slot(&mut src, &mut slot, carry_headroom, chunk_size, slot_fill, &mut pending);
692 if filled_tx.send((slot, got, is_last)).is_err() {
693 break;
694 }
695 if is_last {
696 break;
697 }
698 }
699 });
700
701 // ── Transform workers (persistent, no barrier) ────────────────────────
702 let (work_tx, work_rx) =
703 mpsc::sync_channel::<TypedWorkItem<C::Seg>>(n_workers * 2);
704 let (collector_tx, collector_rx) =
705 mpsc::sync_channel::<TypedMsg<C::Output>>(n_workers * 4);
706 let work_rx = Arc::new(Mutex::new(work_rx));
707
708 for _ in 0..n_workers {
709 let work_rx = Arc::clone(&work_rx);
710 let collector_tx = collector_tx.clone();
711 s.spawn(move || {
712 loop {
713 let item = {
714 let rx = work_rx.lock().expect("work_rx lock");
715 match rx.recv() {
716 Ok(item) => item,
717 Err(_) => break,
718 }
719 };
720 // SAFETY: slot is held alive by collector until all segments done.
721 let data = unsafe {
722 std::slice::from_raw_parts(item.data_ptr, item.data_len)
723 };
724 let output = codec_ref.transform(data, &item.seg);
725 let _ = collector_tx.send(TypedMsg::Result(TypedSegResult {
726 chunk_id: item.chunk_id,
727 seg_id: item.seg_id,
728 output,
729 }));
730 }
731 // This worker has drained all its segments; flush any per-worker
732 // remainder (e.g. a partial row group accumulated across segments).
733 if let Some(tail) = codec_ref.finish_worker() {
734 let _ = collector_tx.send(TypedMsg::Tail(tail));
735 }
736 });
737 }
738 drop(work_rx);
739
740 let inflight_tx = collector_tx.clone();
741 drop(collector_tx);
742
743 // ── Collector thread (trivial: forward in order, no assembly) ─────────
744 let slot_return_for_collector = slot_return_tx.clone();
745 let collector = s.spawn(move || -> Result<()> {
746 let mut in_flight: Vec<TypedInFlight<C::Output>> = Vec::new();
747 let mut next_id: u64 = 0;
748 let mut tails: Vec<C::Output> = Vec::new();
749
750 while let Ok(msg) = collector_rx.recv() {
751 match msg {
752 TypedMsg::New(slot) => in_flight.push(slot),
753 TypedMsg::Result(r) => {
754 for slot in in_flight.iter_mut() {
755 if slot.chunk_id == r.chunk_id {
756 slot.results[r.seg_id] = Some(r.output);
757 slot.done += 1;
758 break;
759 }
760 }
761 }
762 TypedMsg::Tail(out) => tails.push(out),
763 }
764
765 // Flush completed slots strictly in stream order.
766 loop {
767 let Some(idx) = in_flight.iter().position(|s| s.chunk_id == next_id) else {
768 break;
769 };
770 if in_flight[idx].done < in_flight[idx].n_segments {
771 break;
772 }
773
774 let mut done = in_flight.remove(idx);
775 let is_last = done.is_last;
776 let n = done.results.len();
777
778 // Forward each segment's output in order.
779 for (i, output) in done.results.drain(..).enumerate() {
780 if let Some(val) = output {
781 let last_segment = is_last && i == n - 1;
782 sink_ref.process(val, last_segment)?;
783 }
784 }
785
786 slot_return_for_collector.send(done.slot).ok();
787 next_id += 1;
788 }
789 }
790
791 // All ordered segment outputs are flushed; emit per-worker remainders.
792 for tail in tails {
793 sink_ref.process(tail, false)?;
794 }
795 sink_ref.finish()?;
796 Ok(())
797 });
798
799 // ── Main: carry + split + dispatch ────────────────────────────────────
800 let carry_floor = cfg.initial_carry.len();
801 let mut carry: Vec<u8> = cfg.initial_carry;
802 let mut chunk_id: u64 = 0;
803 // Truncation/corruption guard. A well-formed stream never accumulates an
804 // un-splittable carry larger than a few chunks (a single OSM element / bz2
805 // block is < the headroom). If `split` keeps returning None while carry
806 // grows without bound — the exact symptom of a truncated or zero-padded
807 // input (e.g. an incomplete torrent download, whose un-fetched tail is a
808 // sparse hole of zeros with no block magic) — bail loudly here instead of
809 // growing the carry buffer to many GiB and hanging forever.
810 let max_carry = chunk_size.saturating_mul(16).max(1 << 30);
811
812 for (mut slot, read_len, is_last) in filled_rx.iter() {
813 if carry.len() > max_carry {
814 return Err(anyhow::anyhow!(
815 "gatling: no split boundary found across {} MiB of accumulated carry — \
816 the input stream is truncated or corrupt (e.g. an incomplete / zero-padded download)",
817 carry.len() >> 20
818 ));
819 }
820 if read_len == 0 && carry.len() <= carry_floor {
821 slot_return_tx.send(slot).ok();
822 break;
823 }
824
825 // Prepend the carry so the worker sees a single contiguous
826 // `carry ++ read` slice. The common case fits in the slot's reserved
827 // front headroom (zero-copy). When a single element/record spans
828 // enough chunks that the carry outgrows the headroom (planet-scale
829 // OSM: a run with no internal split boundary larger than the 32 MiB
830 // headroom), fall back to a dedicated buffer sized to the actual
831 // carry. The headroom is a fast-path optimization, not a hard cap on
832 // element size.
833 let carry_len = carry.len();
834 let (buf, data_start, data_end) = if carry_len <= carry_headroom {
835 let data_start = carry_headroom - carry_len;
836 slot[data_start..carry_headroom].copy_from_slice(&carry);
837 (slot, data_start, carry_headroom + read_len)
838 } else {
839 let mut combined = Vec::with_capacity(carry_len + read_len);
840 combined.extend_from_slice(&carry);
841 combined.extend_from_slice(&slot[carry_headroom..carry_headroom + read_len]);
842 // The read bytes are copied out; recycle the slot immediately.
843 slot_return_tx.send(slot).ok();
844 let end = combined.len();
845 (combined, 0, end)
846 };
847
848 let data = &buf[data_start..data_end];
849 let split = match codec_ref.split(data, n_workers, is_last) {
850 Some(sp) => sp,
851 None => {
852 carry.clear();
853 carry.extend_from_slice(data);
854 slot_return_tx.send(buf).ok();
855 if is_last {
856 break;
857 }
858 continue;
859 }
860 };
861
862 carry.clear();
863 carry.extend_from_slice(&data[split.consumed..]);
864
865 let n_segments = split.segments.len();
866 let inflight = TypedInFlight {
867 chunk_id,
868 slot: buf,
869 n_segments,
870 results: (0..n_segments).map(|_| None).collect(),
871 done: 0,
872 is_last,
873 };
874
875 let data_ptr = inflight.slot[data_start..].as_ptr();
876 let data_len = data_end - data_start;
877
878 inflight_tx
879 .send(TypedMsg::New(inflight))
880 .map_err(|_| anyhow::anyhow!("collector closed"))?;
881
882 for (seg_id, seg) in split.segments.into_iter().enumerate() {
883 work_tx
884 .send(TypedWorkItem { chunk_id, seg_id, data_ptr, data_len, seg })
885 .map_err(|_| anyhow::anyhow!("work channel closed"))?;
886 }
887
888 chunk_id += 1;
889 if is_last {
890 break;
891 }
892 }
893
894 drop(work_tx);
895 drop(inflight_tx);
896 drop(slot_return_tx);
897
898 collector.join().expect("collector panicked")?;
899 Ok(())
900 })
901}