gatling/src/revolver/coalesce.rs
//! Coalescing / ref-counted mode for [`ChunkRevolver`] — Phase 1 of the
//! "Magazine → chunk_revolver consolidation" plan
//! (`.nornir/chunk_revolver-unsorted.md`).
//!
//! The fixed-slot mode ([`ChunkRevolver::try_get_chunk`]) is 1 write = 1 whole
//! slot: perfect for a stream cut into equal chunks, ruinous for a corpus of
//! small files, where a 4 KiB file would burn a whole slot. znippy forked the
//! pool into `znippy-common/src/slotpool.rs` ("Magazine") to add the one thing
//! that was missing:
//!
//! - **coalescing** — many variable-width payloads packed into ONE slot
//! (invariant: a slice ⊆ one payload AND ⊆ one slot, contiguous),
//! - **per-slot reference counting** — the slot goes back to the free list only
//! after the LAST worker holding a slice of it calls [`Returner::release_one`],
//! - **a blocking claim** — [`ChunkRevolver::claim_slot`] blocks, and that block
//! is the pipeline's only backpressure.
//!
//! This module adds those three to the canonical pool so the fork can be
//! deleted. It is strictly ADDITIVE: `try_get_chunk` / `return_chunk` / `Chunk`
//! are untouched, and a pool built with [`ChunkRevolver::new`] behaves exactly
//! as before.
//!
//! Slot lifecycle: `FREE → (reader claims) FILLING → (reader publishes)
//! DRAINING → (workers release_one; the last one frees it) FREE`.
//!
//! **Safety discipline.** While a slot is FILLING only the reader touches it,
//! exclusively, through [`Filler::writable`] (`&mut self`). After `publish` the
//! reader has dropped the `Filler` and only workers read the slot, shared, via
//! [`PackedSlice::as_slice`]. The slot is not handed back to a reader until its
//! outstanding counter hits zero, so the `&mut` and `&` windows never overlap.
//!
//! **Allocation.** Steady state is zero: the slot memory is the pool's, and the
//! per-slot `Vec<PackedSlice>` a `Filler` fills is recycled through the pool's
//! scratch stack by [`Filler::publish_into`] (the hot-path publish) instead of
//! being allocated per slot.
use std::mem;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use super::revolver::ChunkRevolver;
/// Raw base pointer into a slot buffer. `Send`/`Sync` by the lifecycle
/// discipline documented at the module head — same contract as
/// [`super::SendPtr`], but mutable (the reader fills through it).
#[derive(Copy, Clone)]
pub(super) struct SlotPtr(pub(super) *mut u8);
// SAFETY: a slot is written by exactly one reader while FILLING and read by
// workers only after publish; the pool never hands the same slot to a reader
// and a worker at the same time (the outstanding counter enforces it).
unsafe impl Send for SlotPtr {}
// SAFETY: as above — sharing the base address is sound; the aliasing discipline
// is on the slot's *state*, not on the address.
unsafe impl Sync for SlotPtr {}
/// One unit of work: a borrowed, variable-width view into a published slot.
///
/// Carries the packing metadata a caller needs to rebuild its own record for
/// the payload (znippy's `ChunkMeta`): which input `file_index` the bytes came
/// from, at which `fdata_offset`, which `chunk_seq` within that file, and
/// whether the sink should `skip` recompression.
///
/// The bytes are valid until the slice's slot is released — the holder MUST
/// call [`Returner::release_one`] with `slot_id` exactly once, after it is done
/// reading.
pub struct PackedSlice {
/// Which slot the bytes live in — the id to hand [`Returner::release_one`].
pub slot_id: u32,
ptr: *const u8,
/// Length of this slice in bytes.
pub len: usize,
/// Caller-defined: the sink should pass the bytes through untouched.
pub skip: bool,
/// Caller-defined: index of the input this slice came from.
pub file_index: u64,
/// Caller-defined: byte offset of this slice within that input.
pub fdata_offset: u64,
/// Caller-defined: sequence number of this slice within that input.
pub chunk_seq: u32,
}
// SAFETY: the pointer addresses an immutable region of a published slot. No
// thread writes the slot between publish and the last release, so moving the
// read-only view to a worker thread is sound.
unsafe impl Send for PackedSlice {}
impl PackedSlice {
/// Borrow the slice's bytes — zero copy, straight out of the slot.
///
/// # Safety
/// The caller must hold this `PackedSlice` only while its slot is
/// unreleased: call [`Returner::release_one`] for `slot_id` only after the
/// returned borrow is dropped (or the bytes have been written out).
pub unsafe fn as_slice<'a>(&self) -> &'a [u8] {
// SAFETY: caller upholds the "slot still unreleased" contract above;
// `ptr`/`len` were recorded inside the slot by `Filler::commit_slice`.
unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
}
}
/// The free-slot list. A `Mutex<Vec<u32>>` + `Condvar` rather than a channel:
/// it is touched once per SLOT (not per slice, and never per byte), so the lock
/// is uncontended, and the condvar gives the blocking claim that is the
/// pipeline's backpressure without pulling a channel crate into this leaf.
struct FreeList {
ids: Mutex<Vec<u32>>,
ready: Condvar,
}
impl FreeList {
fn push(&self, id: u32) {
self.ids.lock().expect("free-list mutex poisoned").push(id);
self.ready.notify_one();
}
fn try_pop(&self) -> Option<u32> {
self.ids.lock().expect("free-list mutex poisoned").pop()
}
fn pop_blocking(&self) -> u32 {
let mut g = self.ids.lock().expect("free-list mutex poisoned");
loop {
if let Some(id) = g.pop() {
return id;
}
g = self.ready.wait(g).expect("free-list mutex poisoned");
}
}
}
struct ReturnerInner {
outstanding: Vec<AtomicUsize>,
free: Arc<FreeList>,
}
/// Shared handle workers use to hand a slot back. Cloneable (an `Arc` inside) —
/// move a clone into each worker thread.
///
/// This is the piece that cannot live on `&mut ChunkRevolver`: the reader owns
/// the pool and is off filling the next slot while N workers finish with the
/// previous one.
#[derive(Clone)]
pub struct Returner {
inner: Arc<ReturnerInner>,
}
impl Returner {
/// One slice of `slot_id` has been fully consumed. When the slot's last
/// slice is released the slot returns to the free list, unblocking a reader
/// waiting in [`ChunkRevolver::claim_slot`].
///
/// Call it exactly once per [`PackedSlice`].
pub fn release_one(&self, slot_id: u32) {
let prev = self.inner.outstanding[slot_id as usize].fetch_sub(1, Ordering::AcqRel);
debug_assert!(prev >= 1, "release_one underflow on slot {slot_id}");
if prev == 1 {
self.inner.free.push(slot_id);
}
}
}
/// The coalescing state armed onto a [`ChunkRevolver`] by
/// [`ChunkRevolver::with_coalescing`].
pub(super) struct Coalescing {
slot_size: usize,
slice_size: usize,
/// Base pointer per global slot id — the offset-pointer form the reader and
/// the workers both address the slot through, so no byte is ever copied to
/// hand a payload from one to the other.
///
/// The plan sketch (`.nornir/chunk_revolver-unsorted.md`) proposed a
/// `slot_ring: Vec<u8>` so a freed slot could be pushed back to *its* ring.
/// It is not carried: a coalescing slot has no ring affinity to preserve
/// (its consumers are N workers pulling off one global queue, not one
/// worker bound to one ring), so a flat free list is both the smaller state
/// and the correct one. Ring identity is recoverable from the base pointer
/// if it is ever wanted.
slot_base: Vec<SlotPtr>,
free: Arc<FreeList>,
ret: Returner,
/// Recycled `Vec<PackedSlice>` buffers — one per in-flight slot, handed to
/// a `Filler` on claim and handed back on publish. Preallocated: the
/// steady-state publish allocates nothing.
scratch: Mutex<Vec<Vec<PackedSlice>>>,
}
impl Coalescing {
pub(super) fn new(slot_size: usize, slice_size: usize, slot_base: Vec<SlotPtr>) -> Self {
let n = slot_base.len();
let free = Arc::new(FreeList {
ids: Mutex::new((0..n as u32).rev().collect()),
ready: Condvar::new(),
});
let outstanding = (0..n).map(|_| AtomicUsize::new(0)).collect();
let ret = Returner {
inner: Arc::new(ReturnerInner {
outstanding,
free: Arc::clone(&free),
}),
};
// One scratch buffer per slot, sized for the case where every slice is a
// full `slice_size` cut. Smaller payloads make it grow ONCE and then it
// is reused for the life of the pool.
let cap = (slot_size / slice_size.max(1)).max(16);
let scratch = Mutex::new(
(0..n)
.map(|_| Vec::<PackedSlice>::with_capacity(cap))
.collect(),
);
Self {
slot_size,
slice_size,
slot_base,
free,
ret,
scratch,
}
}
fn take_scratch(&self) -> Vec<PackedSlice> {
self.scratch
.lock()
.expect("scratch mutex poisoned")
.pop()
.unwrap_or_default()
}
fn give_scratch(&self, mut v: Vec<PackedSlice>) {
if v.capacity() == 0 {
return;
}
v.clear();
self.scratch.lock().expect("scratch mutex poisoned").push(v);
}
}
impl ChunkRevolver {
/// Build a pool in **coalescing mode**: `num_slots` slots of `slot_size`
/// bytes, into which the reader packs variable-width payloads.
///
/// `slice_size` is the granularity at which the caller is expected to cut a
/// payload that is larger than one slice. The pool only records it (see
/// [`ChunkRevolver::slice_size`]) — it is the caller's cut policy, and it is
/// output-affecting for an archiver, so it is planned by the caller, never
/// derived here.
///
/// Memory layout is one `Box<[u8]>` per slot, so the reservation is
/// `num_slots * slot_size` and nothing is contiguous across slots. The
/// fixed-slot free lists are drained at construction: a coalescing pool can
/// never hand the same memory out through [`ChunkRevolver::try_get_chunk`].
///
/// # Panics
/// If `num_slots`, `slot_size` or `slice_size` is zero, or if `slice_size >
/// slot_size` (a slice must fit contiguously in one slot).
pub fn with_coalescing(num_slots: usize, slot_size: usize, slice_size: usize) -> Self {
assert!(
num_slots > 0 && slot_size > 0 && slice_size > 0,
"coalescing pool needs num_slots/slot_size/slice_size > 0"
);
assert!(
slice_size <= slot_size,
"slice_size {slice_size} exceeds slot_size {slot_size}: a slice must fit \
contiguously in one slot"
);
let mut rev = Self::new(slot_size, num_slots, num_slots);
rev.arm_coalescing(slice_size);
rev
}
/// The cut granularity this pool was planned with.
///
/// # Panics
/// If the pool is not in coalescing mode.
pub fn slice_size(&self) -> usize {
self.coalescing().slice_size
}
/// Number of coalescing slots.
///
/// # Panics
/// If the pool is not in coalescing mode.
pub fn num_slots(&self) -> usize {
self.coalescing().slot_base.len()
}
/// Handle workers use to release slots. Clone it into each worker thread.
///
/// # Panics
/// If the pool is not in coalescing mode.
pub fn returner(&self) -> Returner {
self.coalescing().ret.clone()
}
/// Release one slice of `slot_id` — the convenience form of
/// [`Returner::release_one`] for a caller that still has the pool in hand.
///
/// # Panics
/// If the pool is not in coalescing mode.
pub fn release_one(&self, slot_id: u32) {
self.coalescing().ret.release_one(slot_id);
}
/// Reader: claim a free slot to fill. **Blocks** until one is free — this
/// block is the only backpressure in the pipeline.
///
/// Returns `Some` for as long as the pool lives; the `Option` mirrors the
/// fallible shape of [`ChunkRevolver::try_claim_slot`] so a caller can move
/// between them without changing shape. A caller that never releases a slot
/// will block here forever — that is a lost [`Returner::release_one`], not a
/// pool bug.
///
/// # Panics
/// If the pool is not in coalescing mode.
pub fn claim_slot(&self) -> Option<Filler<'_>> {
let c = self.coalescing();
let slot_id = c.free.pop_blocking();
Some(self.filler_for(slot_id))
}
/// Non-blocking claim: `None` when every slot is in flight.
///
/// # Panics
/// If the pool is not in coalescing mode.
pub fn try_claim_slot(&self) -> Option<Filler<'_>> {
let c = self.coalescing();
let slot_id = c.free.try_pop()?;
Some(self.filler_for(slot_id))
}
fn filler_for(&self, slot_id: u32) -> Filler<'_> {
let slices = self.coalescing().take_scratch();
Filler {
pool: self,
slot_id,
cursor: 0,
slices,
published: false,
}
}
}
/// Reader-side handle for filling one claimed slot.
///
/// Pack payloads with [`writable`](Filler::writable) + [`commit_slice`](Filler::commit_slice),
/// then [`publish_into`](Filler::publish_into) (hot path, recycles the buffer)
/// or [`publish`](Filler::publish) to arm the slot's reference count and hand
/// the slices to the workers.
///
/// Dropping a `Filler` without publishing returns the slot to the free list —
/// a panic in the reader loses the packed bytes but never leaks the slot.
pub struct Filler<'a> {
pool: &'a ChunkRevolver,
slot_id: u32,
cursor: usize,
slices: Vec<PackedSlice>,
published: bool,
}
impl<'a> Filler<'a> {
/// The slot being filled.
pub fn slot_id(&self) -> u32 {
self.slot_id
}
/// Bytes of free space left in the slot.
pub fn remaining(&self) -> usize {
self.pool.coalescing().slot_size - self.cursor
}
/// Number of slices committed so far.
pub fn committed(&self) -> usize {
self.slices.len()
}
/// Writable view at the cursor, up to `max` bytes (clamped to
/// [`remaining`](Filler::remaining)). The reader reads payload bytes
/// **directly** into this — there is no intermediate buffer and no copy.
///
/// The borrow must be dropped before calling
/// [`commit_slice`](Filler::commit_slice).
pub fn writable(&mut self, max: usize) -> &mut [u8] {
let n = max.min(self.remaining());
let base = self.pool.coalescing().slot_base[self.slot_id as usize].0;
// SAFETY: this slot is FILLING — it is not in any free list and no
// slices of it are published, so the reader holding `&mut self` is the
// only accessor. `cursor + n <= slot_size` by the clamp above.
unsafe { std::slice::from_raw_parts_mut(base.add(self.cursor), n) }
}
/// Commit the `len` bytes just written at the cursor as one slice and
/// advance the cursor. The metadata arguments are passed through verbatim
/// onto the resulting [`PackedSlice`].
pub fn commit_slice(
&mut self,
len: usize,
skip: bool,
file_index: u64,
fdata_offset: u64,
chunk_seq: u32,
) {
debug_assert!(
len <= self.remaining(),
"commit_slice({len}) exceeds the {} bytes left in slot {}",
self.remaining(),
self.slot_id
);
let base = self.pool.coalescing().slot_base[self.slot_id as usize].0 as *const u8;
// SAFETY: the cursor stays inside the slot (asserted above); this only
// records the address, it does not dereference it.
let ptr = unsafe { base.add(self.cursor) };
self.slices.push(PackedSlice {
slot_id: self.slot_id,
ptr,
len,
skip,
file_index,
fdata_offset,
chunk_seq,
});
self.cursor += len;
}
/// Publish the slot into a caller-owned buffer: arms the slot's reference
/// count to the number of committed slices and appends them to `out`,
/// returning how many were appended. An empty slot is returned to the free
/// list immediately and appends nothing.
///
/// This is the **hot-path** publish: the `Filler`'s internal buffer goes
/// back to the pool's scratch stack, so a steady-state reader loop that
/// reuses one `out` allocates nothing at all.
pub fn publish_into(mut self, out: &mut Vec<PackedSlice>) -> usize {
let mut slices = mem::take(&mut self.slices);
let n = self.arm(slices.len());
out.extend(slices.drain(..));
self.published = true;
// Straight back to the scratch stack with its capacity intact — this is
// the allocation the fork made once per slot and this one does not.
self.pool.coalescing().give_scratch(slices);
n
}
/// Publish the slot and return the slices as a fresh `Vec`.
///
/// Convenience shape (it matches znippy's `Clip::publish`); it hands the
/// buffer to the caller rather than recycling it, so prefer
/// [`publish_into`](Filler::publish_into) on a hot path.
#[must_use = "the published slices are the work — dropping them leaks the slot"]
pub fn publish(mut self) -> Vec<PackedSlice> {
let slices = mem::take(&mut self.slices);
self.arm(slices.len());
self.published = true;
slices
}
/// Arm the outstanding counter for `n` slices; free the slot outright when
/// `n == 0`. Returns `n`.
fn arm(&self, n: usize) -> usize {
let c = self.pool.coalescing();
if n == 0 {
c.free.push(self.slot_id);
return 0;
}
c.ret.inner.outstanding[self.slot_id as usize].store(n, Ordering::Release);
n
}
}
impl<'a> Drop for Filler<'a> {
fn drop(&mut self) {
let c = self.pool.coalescing();
let slices = mem::take(&mut self.slices);
if !self.published {
// Un-published: nothing was handed to a worker, so the slot is free
// the moment the reader lets go of it. Without this a panic (or an
// early `?`) in the reader loop would retire a slot permanently and
// the pipeline would deadlock a claim at a time.
c.free.push(self.slot_id);
}
c.give_scratch(slices);
}
}