Skip to main content

mz_ore/
pool.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Prototype buffer pool for dataflow state. See
17//! `doc/developer/design/20260610_buffer_managed_state.md`.
18//!
19//! The pool is the cache: size-class anonymous virtual-memory regions whose
20//! slots hold resident chunks. Slots are scoped to residency — eviction
21//! returns a chunk's slot to the free list along with its physical pages —
22//! so slot demand tracks the resident set (bounded by the budget), not the
23//! potentially unbounded live backlog. Reads are copy-out
24//! ([`ChunkHandle::read_into`]): a resident slot is copied and an evicted
25//! extent decompressed straight into the caller's buffer, all under the
26//! chunk's state lock, so no reference into pool memory escapes the pool and
27//! a read leaves residency untouched. The backing is the swap-backed extent
28//! store of the design's Layer 1: a slot in a pool-owned anonymous-memory
29//! extent arena holding the chunk's lz4-compressed bytes.
30//!
31//! Memory descends a ladder of tiers, each with its own ceiling and each
32//! cheaper to vacate than the one above:
33//!
34//! * **Slots** (uncompressed, free reads) — bounded by the budget; crossing
35//!   it compresses the oldest chunks into extents and releases their slots.
36//! * **Warm free slots** (pages kept for fault-free reuse) — bounded by the
37//!   warm cap.
38//! * **Compressed-resident extents** (reads decompress, no device) — bounded
39//!   by the headroom the RSS target leaves above the first two; crossing it
40//!   pushes the oldest extents to the swap device with `MADV_PAGEOUT`.
41//! * **The swap device** — overflow; reads fault and decompress.
42//!
43//! Residency is a state, not a type. It descends through eviction and
44//! ascends through exactly one transition: an admitting read
45//! ([`ChunkHandle::read_into_admit`]) lifts an evicted chunk back to
46//! `BackedResident` when a slot is free within the budget or stealable from
47//! a clean backed victim of the same class, never by evicting or
48//! compressing anything. Plain reads ([`ChunkHandle::read_into`]) leave
49//! residency untouched. Eviction I/O runs on spill threads when enabled —
50//! `WriteInFlight` marks a chunk whose compression a spill thread owns — and
51//! inline on the evicting caller otherwise. Chunks are immutable after
52//! [`Pool::insert_with`], which is what makes a `BackedResident` slot always
53//! identical to its extent and its eviction free of I/O.
54//!
55//! Freeing an `UnbackedResident` chunk is a pure memory operation — the
56//! design's "never write dead data" win, surfaced as `writes_elided` in
57//! [`PoolStats`]. Budget pressure evicts cold chunks via second-chance
58//! FIFOs banded by the caller-supplied generational depth ([`ChunkHints`]).
59
60mod extent;
61mod region;
62
63use std::collections::VecDeque;
64use std::ops::Range;
65use std::sync::atomic::{AtomicU64, Ordering};
66use std::sync::{Arc, Mutex, MutexGuard, Weak};
67
68use crate::cast::CastFrom;
69use crate::pool::extent::{ExtentArena, Scratch, SwapExtent};
70use crate::pool::region::{Region, SIZE_CLASSES};
71
72/// Virtual reservation per size class. Purely virtual: physical memory
73/// materializes only for slots in use, and slots are scoped to residency,
74/// so this must exceed the largest plausible *resident* set per class, the
75/// budget plus in-flight slack, not the backlog. It is deliberately enormous
76/// (address space costs nothing, and touched pages are bounded by peak
77/// residency) so that no realistic budget, on any machine size, reaches the
78/// heap-fallback path.
79///
80/// NOTE: Seen OoMs with Miri since it actually allocates the capacity.
81const CLASS_CAPACITY_BYTES: usize = if cfg!(miri) { 16 << 20 } else { 1 << 40 };
82
83/// A chunk-provided transform between a chunk's body bytes and the stored
84/// bytes its extent holds. The pool owns scheduling: spill threads, the
85/// residency state machine, cancellation, and the ledger. It invokes the
86/// codec on opaque bytes at the extent boundary, `encode` when backing a
87/// chunk (on a spill thread, or inline under overload) and `decode` when
88/// reading an evicted one, under the chunk's state lock. The pool itself
89/// has no opinion on the stored form: framing, compression, and validation
90/// all belong to the codec.
91///
92/// Implementations must be pure transforms: no locking, no calls back into
93/// the pool (the state lock is held at `decode` sites), and no panic on
94/// bytes their own `encode` produced. `decode` must exactly invert
95/// `encode`, and `encode`'s output must never exceed
96/// [`max_stored_len`]`(body.len())`, the bound the extent store's size
97/// classes are provisioned to.
98pub trait ExtentCodec: std::fmt::Debug + Send + Sync {
99    /// Transforms `body` into its stored form, replacing `out`'s contents.
100    /// `out`'s capacity is reused across calls; implementations size it
101    /// themselves.
102    fn encode(&self, body: &[u8], out: &mut Vec<u8>);
103
104    /// Inverts [`ExtentCodec::encode`]: reconstructs into `body` exactly
105    /// the bytes whose encoding produced `stored`. `body` is exactly the
106    /// original body's length, and implementations must panic on a length
107    /// mismatch rather than truncate or pad.
108    fn decode(&self, stored: &[u8], body: &mut [u8]);
109}
110
111/// The identity [`ExtentCodec`]: the stored form is the body. Encode and
112/// decode are copies, and range reads copy the range directly, so a chunk
113/// stored under this codec pays no compression work in either direction
114/// while remaining fully budgeted and swap-backed like any other extent.
115#[derive(Debug)]
116pub struct IdentityCodec;
117
118/// The [`IdentityCodec`] instance to pass to [`Pool::insert_with`].
119pub static IDENTITY_CODEC: IdentityCodec = IdentityCodec;
120
121impl ExtentCodec for IdentityCodec {
122    fn encode(&self, body: &[u8], out: &mut Vec<u8>) {
123        out.clear();
124        out.extend_from_slice(body);
125    }
126
127    fn decode(&self, stored: &[u8], body: &mut [u8]) {
128        assert_eq!(stored.len(), body.len(), "identity stored form is the body");
129        body.copy_from_slice(stored);
130    }
131}
132
133/// The largest stored form [`ExtentCodec::encode`] may produce for a
134/// `body_len`-byte body: an incompressible-input expansion matching lz4's
135/// worst case plus a four-byte length prefix. The extent store's size-class
136/// ladder is provisioned to this bound, so a codec that exceeds it can
137/// strand payloads with no class to hold them (they degrade to unpageable
138/// heap fallbacks).
139pub fn max_stored_len(body_len: usize) -> usize {
140    4 + body_len + body_len / 255 + 16
141}
142
143/// Advisory placement hints for a chunk, supplied at insert and immutable
144/// thereafter (merges mint new chunks, so a chunk's generation never
145/// changes). Hints steer policy — eviction order and write-behind
146/// candidacy — never correctness: a mislabeled chunk performs worse, while
147/// the budget and residency invariants hold regardless.
148#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
149pub struct ChunkHints {
150    /// Generational depth of the chunk in its producer's merge structure,
151    /// 0 for the youngest generation (and the unannotated default). Deeper
152    /// chunks are treated as colder: preferred write-behind candidates and
153    /// preferred eviction victims, cheap to evict once backed.
154    pub depth: u8,
155}
156
157/// Number of depth bands the eviction queues are split into; depths at or
158/// beyond the last band share it.
159const DEPTH_BANDS: usize = 4;
160
161/// The eviction-queue band for a chunk of `depth`.
162fn band(depth: u8) -> usize {
163    usize::from(depth).min(DEPTH_BANDS - 1)
164}
165
166/// Residency state of a chunk.
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168enum Residency {
169    /// Lives only in the pool; no extent copy exists. Freeing it never
170    /// touches the backing store.
171    UnbackedResident,
172    /// Resident, and an identical extent copy exists; eviction releases
173    /// physical pages without I/O.
174    BackedResident,
175    /// Resident and readable, with compression into an extent scheduled on a
176    /// spill thread. Completion moves an evicting chunk to
177    /// [`Residency::Evicted`] and an eagerly backed one to
178    /// [`Residency::BackedResident`]; a free observed at dequeue cancels the
179    /// write instead.
180    WriteInFlight,
181    /// Extent copy only; the chunk holds no slot. The extent itself may
182    /// still be RAM-resident (the compressed tier) or paged out to the swap
183    /// device. Reads decompress the extent straight into the caller's
184    /// buffer and leave the chunk evicted, except that an admitting read
185    /// may lift it back to [`Residency::BackedResident`].
186    Evicted,
187    /// Larger than the largest size class; held as a plain heap allocation,
188    /// always resident. A prototype limitation, not a design state.
189    Oversize,
190}
191
192/// Snapshot of pool counters.
193#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
194pub struct PoolStats {
195    /// Chunks inserted.
196    pub inserts: u64,
197    /// Chunks freed (handle dropped).
198    pub frees: u64,
199    /// Backing writes elided: chunks dead before their compression
200    /// completed, so no extent write happened. Covers chunks freed while
201    /// `UnbackedResident` and chunks freed while queued for a spill thread
202    /// that had not yet compressed them.
203    pub writes_elided: u64,
204    /// Evictions that compressed the chunk into a new extent.
205    pub evictions_compress: u64,
206    /// Evictions of `BackedResident` chunks: pure page release, no I/O.
207    pub evictions_cheap: u64,
208    /// Compressed bytes written into extents.
209    pub extent_bytes_written: u64,
210    /// Evictions handed to spill threads.
211    pub spill_scheduled: u64,
212    /// Compressions cancelled because the chunk was freed while queued or
213    /// in flight, whatever scheduled them. Eager-backing work counts here
214    /// but never in `spill_scheduled`, so this can exceed that counter.
215    pub spill_cancelled: u64,
216    /// Entries currently queued for or being processed by spill threads.
217    pub spill_in_flight: u64,
218    /// Inserts that fell back to the heap because their size class had no
219    /// free slot (the live set outgrew the class reservation). Heap-backed
220    /// chunks behave like oversize ones: always resident, never paged.
221    pub slot_exhausted_fallbacks: u64,
222    /// Inserts whose payload exceeded the largest size class and therefore
223    /// went straight to a heap-backed oversize chunk.
224    pub oversize_payloads: u64,
225    /// Live size-classed chunks across all classes, whatever their residency.
226    /// For backlog-shaped consumers this tracks the un-drained backlog in
227    /// chunks.
228    pub live_chunks: u64,
229    /// Uncompressed bytes of currently resident chunks (including oversize).
230    pub resident_bytes: u64,
231    /// Uncompressed bytes of live oversize chunks.
232    pub oversize_bytes: u64,
233    /// Class bytes of free slots currently kept warm (pages resident for
234    /// fault-free reuse). Bounded by a fraction of the budget; RSS exceeds
235    /// `resident_bytes` by up to this amount.
236    pub warm_bytes: u64,
237    /// Slot allocations served from the warm list: reuses that faulted no
238    /// pages and skipped the kernel's page zeroing.
239    pub warm_reuses: u64,
240    /// Chunks eagerly compressed to `BackedResident` by idle spill threads
241    /// (write-behind): still readable in their slots, with eviction
242    /// pre-paid.
243    pub eager_backs: u64,
244    /// Evicted chunks re-admitted to `BackedResident` by an admitting read
245    /// out of free budget headroom.
246    pub admissions_budget: u64,
247    /// Evicted chunks re-admitted to `BackedResident` by an admitting read
248    /// stealing the slot of a clean backed victim of the same size class.
249    /// The victim becomes `Evicted` with zero I/O and its extent intact.
250    pub admissions_steal: u64,
251    /// Admitting reads of evicted chunks served as a plain decompress
252    /// instead: no budget headroom (or an exhausted size class), and no
253    /// clean victim whose growth the budget could absorb.
254    pub admissions_denied: u64,
255    /// Allocation bytes of compressed extents currently resident — the
256    /// compressed-but-resident middle tier. Bounded by the RSS target;
257    /// exceeding it pages the oldest extents out to the swap device.
258    pub extent_resident_bytes: u64,
259    /// Allocation bytes of resident extents the RSS target cannot push out:
260    /// retry-capped arena extents (the kernel declined the reclaim advice
261    /// until the retry budget ran out) and heap-fallback extents. The
262    /// compressed tier settles above its capacity by this amount.
263    pub extent_unreclaimable_bytes: u64,
264    /// Extents pushed to the swap device by RSS-target enforcement, with
265    /// the whole range observed nonresident afterwards.
266    pub extent_pageouts: u64,
267    /// Pageout passes whose observation found some of the extent's pages
268    /// still mapped: `MADV_PAGEOUT` may decline pages and still succeed, so
269    /// the page table decides. The extent keeps its full resident
270    /// accounting and is retried until its per-extent retry cap. Climbing
271    /// steadily on a loaded pool means pages cannot actually be unmapped to
272    /// the swap device (no swap, or a cgroup that cannot reclaim).
273    pub extent_pageout_incomplete: u64,
274    /// Extent writes that fell back to the heap because their extent-arena
275    /// class had no free slot. Heap-backed extents stay readable but are
276    /// never paged out, so their compressed bytes hold RAM until freed.
277    pub extent_arena_fallbacks: u64,
278}
279
280#[derive(Debug, Default)]
281struct Counters {
282    inserts: AtomicU64,
283    spill_scheduled: AtomicU64,
284    spill_cancelled: AtomicU64,
285    slot_exhausted_fallbacks: AtomicU64,
286    oversize_payloads: AtomicU64,
287    frees: AtomicU64,
288    writes_elided: AtomicU64,
289    evictions_compress: AtomicU64,
290    evictions_cheap: AtomicU64,
291    extent_bytes_written: AtomicU64,
292    resident_bytes: AtomicU64,
293    oversize_bytes: AtomicU64,
294    warm_bytes: AtomicU64,
295    warm_reuses: AtomicU64,
296    eager_backs: AtomicU64,
297    admissions_budget: AtomicU64,
298    admissions_steal: AtomicU64,
299    admissions_denied: AtomicU64,
300    extent_resident_bytes: AtomicU64,
301    extent_unreclaimable_bytes: AtomicU64,
302    extent_pageouts: AtomicU64,
303    extent_pageout_incomplete: AtomicU64,
304}
305
306/// A buffer pool over swap-backed extents. Cheap to clone; all clones share
307/// one budget and one backing store.
308#[derive(Debug, Clone)]
309pub struct Pool(Arc<PoolInner>);
310
311/// The shared state behind every [`Pool`] handle. One per process in
312/// practice; [`Pool`] clones and chunk handles share it through an `Arc`,
313/// so it lives until the last handle and spill thread release it.
314///
315/// Lock order: a chunk's `state` mutex may be held while taking any of the
316/// leaf locks — the eviction `queue`, the `extent_queue`, the spill queue,
317/// and the region slot allocators — but never the reverse. The enforcement
318/// and backing scans additionally drop the queue guard before trying a
319/// chunk's state lock (and only ever `try_lock` it), so no path holds a
320/// queue lock while waiting on chunk state. The admitting read's victim
321/// steal is the one place a chunk's state lock is held while probing
322/// another chunk's, and the victim is only ever `try_lock`ed, so two
323/// admitters stealing toward each other skip instead of deadlocking. Reads
324/// copy out under the chunk's state lock — the same lock eviction takes —
325/// so there is no reader-side count and no reader the evictor must account
326/// for.
327#[derive(Debug)]
328struct PoolInner {
329    /// Resident-bytes target, enforced against evictable bytes (resident
330    /// minus heap-backed, which no eviction can reclaim). Atomic so a
331    /// running pool can be retuned in place (operator-driven budget
332    /// changes) without orphaning live handles, which share this value
333    /// through their `Arc<PoolInner>`.
334    budget_bytes: AtomicU64,
335    /// Ceiling on the pool's *total* RSS: slots (the budget) plus warm free
336    /// slots plus compressed-resident extents. The compressed tier's
337    /// capacity derives as `max(0, rss_target - budget - warm cap)`; zero
338    /// (the default) collapses the tier, paging every extent out as soon as
339    /// it is written.
340    rss_target_bytes: AtomicU64,
341    /// One region per entry of [`SIZE_CLASSES`], same order.
342    regions: Vec<Region>,
343    /// The arena backing extents. Shared with every live [`SwapExtent`],
344    /// whose drop returns its slot.
345    extent_arena: Arc<ExtentArena>,
346    /// Second-chance FIFOs of eviction candidates, one per depth band; a
347    /// chunk joins the band of its [`ChunkHints`] depth at insert and again
348    /// on re-admission. Entries for freed chunks go stale in place and are
349    /// dropped by [`PoolInner::prune_queues`].
350    ///
351    /// Two scanners walk them with different obligations, both visiting the
352    /// deepest band first. Budget enforcement is the one that ages chunks:
353    /// it spends the touched bit (second chance) and drops entries it
354    /// evicts. Eager backing ([`PoolInner::back_one`]) rotates visited
355    /// entries to the back but never spends a touched bit, so a backing
356    /// pass shuffles FIFO order without aging any chunk toward eviction.
357    queues: [Mutex<VecDeque<Weak<ChunkMeta>>>; DEPTH_BANDS],
358    /// FIFO of chunks whose extents are resident, oldest first — the
359    /// RSS-target enforcement's victim queue. Entries go stale when an
360    /// extent pages out, is dropped, or its chunk dies; visits drop them,
361    /// and [`PoolInner::prune_extent_queue`] compacts dead-chunk entries
362    /// that under-cap operation never visits.
363    extent_queue: Mutex<VecDeque<Weak<ChunkMeta>>>,
364    /// Number of live size-classed chunks (whatever their residency), which
365    /// is the number of non-stale queue entries across all bands;
366    /// [`PoolInner::prune_queues`] compacts the queues against it.
367    live_chunks: AtomicU64,
368    /// Number of live chunks whose extent is currently resident, including
369    /// unreclaimable extents that deliberately hold no `extent_queue` entry
370    /// (heap-backed and retry-capped ones). It therefore upper-bounds the
371    /// queue's non-stale entries, and [`PoolInner::prune_extent_queue`]'s
372    /// compaction threshold is conservative by the unreclaimable count.
373    extent_residents: AtomicU64,
374    /// Single-flight claim for budget enforcement.
375    enforcing: Mutex<()>,
376    /// Set by an insert turned away from `enforcing`. The holder re-runs its
377    /// pass while it is set, so a caller turned away after the holder's final
378    /// counter read still has its bytes enforced rather than dropped.
379    enforce_pending: std::sync::atomic::AtomicBool,
380    counters: Counters,
381    spill: Spill,
382}
383
384/// Hand-off point between budget enforcement and spill threads. Eviction I/O
385/// (compression and the synchronous-reclaim `pageout`) runs on spill threads
386/// when enabled, keeping multi-millisecond work off the threads that trip the
387/// budget; with no spill threads, eviction runs inline on the caller.
388#[derive(Debug, Default)]
389struct Spill {
390    /// Chunks in `WriteInFlight`, awaiting a spill thread.
391    queue: Mutex<VecDeque<Arc<ChunkMeta>>>,
392    /// Parks idle spill threads. Notified when work lands in `queue`, when
393    /// eager backing turns on, and at shutdown; threads additionally wake
394    /// on a timeout so eager backing scans for write-behind work without a
395    /// dedicated wakeup per candidate.
396    cv: std::sync::Condvar,
397    /// Whether evictions are handed to spill threads. Set when threads are
398    /// first spawned; cleared to fall back to inline eviction.
399    enabled: std::sync::atomic::AtomicBool,
400    /// Whether idle spill threads eagerly compress unbacked chunks to
401    /// `BackedResident` (write-behind); see [`Pool::set_eager_backing`].
402    eager: std::sync::atomic::AtomicBool,
403    /// Number of spill threads spawned (spawn-once; later config changes
404    /// only toggle `enabled`).
405    threads: AtomicU64,
406    /// Queued plus currently-processing entries; `quiesce` waits on zero.
407    in_flight: AtomicU64,
408    /// Test-only lifecycle: production spill threads are immortal (the pool
409    /// is a process singleton), but Miri rejects a test binary exiting with
410    /// live threads, so tests stop and join them.
411    #[cfg(test)]
412    stop: std::sync::atomic::AtomicBool,
413    #[cfg(test)]
414    handles: Mutex<Vec<std::thread::JoinHandle<()>>>,
415}
416
417/// Beyond this many queued or in-flight spill entries, eviction degrades to
418/// inline on the caller: bounded memory overshoot under burst beats an
419/// unbounded queue of still-resident chunks.
420const SPILL_IN_FLIGHT_MAX: usize = 64;
421
422/// What a spill thread does with a chunk once compressed.
423#[derive(Clone, Copy, PartialEq, Eq)]
424enum SpillKind {
425    /// Budget-driven: release the slot, leaving the chunk `Evicted`.
426    Evict,
427    /// Eager write-behind: keep the slot, leaving the chunk
428    /// `BackedResident`.
429    Back,
430}
431
432#[derive(Debug)]
433struct ChunkMeta {
434    pool: Arc<PoolInner>,
435    /// Length in `u64` words; immutable.
436    len: usize,
437    /// Size class for slot allocations; `None` for empty chunks and payloads
438    /// beyond the largest class. Immutable: the chunk's *slot* comes and goes
439    /// with residency, but it is always drawn from this class.
440    class: Option<usize>,
441    /// The insert-time [`ChunkHints`] depth; immutable. Names the eviction
442    /// band the chunk's queue entries belong to.
443    depth: u8,
444    /// The insert-time [`ExtentCodec`]; immutable. Encodes the chunk when it
445    /// is backed and decodes its extent on reads, so it must outlive any
446    /// extent it produced, hence `'static`.
447    codec: &'static dyn ExtentCodec,
448    state: Mutex<ChunkState>,
449}
450
451#[derive(Debug)]
452struct ChunkState {
453    residency: Residency,
454    /// Second-chance bit, set on read and cleared (in lieu of eviction) when
455    /// the budget enforcer first visits the chunk.
456    touched: bool,
457    /// Set when the owning handle is dropped, so a queue entry upgraded
458    /// concurrently with the free cannot touch a recycled slot.
459    freed: bool,
460    /// The chunk's slot index within its class's region, held exactly while
461    /// the chunk occupies pool memory (the resident states and
462    /// `WriteInFlight`). Eviction returns the slot to the region free list.
463    /// Reads copy the slot out under the state lock; no pointer into the
464    /// slot outlives the lock under which it was formed.
465    slot: Option<u32>,
466    /// The backing copy; present exactly in the `BackedResident` and
467    /// `Evicted` states.
468    extent: Option<SwapExtent>,
469    /// The payload of an `Oversize` chunk.
470    oversize: Option<Vec<u64>>,
471}
472
473impl ChunkMeta {
474    /// A fresh chunk in its insert-time state.
475    fn new(
476        pool: &Arc<PoolInner>,
477        len: usize,
478        class: Option<usize>,
479        depth: u8,
480        codec: &'static dyn ExtentCodec,
481        residency: Residency,
482        slot: Option<u32>,
483        oversize: Option<Vec<u64>>,
484    ) -> ChunkMeta {
485        ChunkMeta {
486            pool: Arc::clone(pool),
487            len,
488            class,
489            depth,
490            codec,
491            state: Mutex::new(ChunkState {
492                residency,
493                touched: false,
494                freed: false,
495                slot,
496                extent: None,
497                oversize,
498            }),
499        }
500    }
501
502    fn len_bytes(&self) -> usize {
503        self.len * std::mem::size_of::<u64>()
504    }
505
506    /// Locks the chunk's state.
507    fn state(&self) -> MutexGuard<'_, ChunkState> {
508        self.state.lock().expect("chunk state poisoned")
509    }
510}
511
512/// Handle to one immutable chunk in a [`Pool`]. Dropping the handle frees the
513/// chunk: the slot (if resident) returns to the region free list with its
514/// physical pages released, and the extent (if any) is deallocated,
515/// discarding any swapped copy for free.
516#[derive(Debug)]
517pub struct ChunkHandle {
518    meta: Arc<ChunkMeta>,
519}
520
521// Test hook fired inside `PoolInner::enforce_budget`, between a pass's final
522// counter read and the release of the `enforcing` guard. A test arms it on the
523// thread whose pass it wants to freeze, to interleave a concurrent over-budget
524// insert. One-shot: the hook is taken before it runs, so a re-enforcing pass
525// does not re-arm.
526#[cfg(test)]
527thread_local! {
528    static ENFORCE_BUDGET_HOOK: std::cell::RefCell<Option<Box<dyn FnOnce()>>> =
529        const { std::cell::RefCell::new(None) };
530}
531
532#[cfg(test)]
533fn run_enforce_budget_hook() {
534    let hook = ENFORCE_BUDGET_HOOK.with(|cell| cell.borrow_mut().take());
535    if let Some(hook) = hook {
536        hook();
537    }
538}
539
540impl Pool {
541    /// Creates a pool, reserving one virtual region per size class. The
542    /// pool starts with an unlimited budget — nothing is evicted until
543    /// [`Pool::set_budget`] tunes it.
544    pub fn new() -> std::io::Result<Pool> {
545        Pool::with_class_capacity(CLASS_CAPACITY_BYTES)
546    }
547
548    /// As [`Pool::new`], with a caller-chosen virtual reservation per size
549    /// class. Small reservations let tests exercise slot exhaustion.
550    fn with_class_capacity(class_capacity_bytes: usize) -> std::io::Result<Pool> {
551        let regions = SIZE_CLASSES
552            .iter()
553            .map(|&class_size| Region::new(class_size, class_capacity_bytes))
554            .collect::<std::io::Result<Vec<_>>>()?;
555        let extent_arena = Arc::new(ExtentArena::new(class_capacity_bytes)?);
556        Ok(Pool(Arc::new(PoolInner {
557            budget_bytes: AtomicU64::new(u64::MAX),
558            rss_target_bytes: AtomicU64::new(0),
559            regions,
560            extent_arena,
561            queues: std::array::from_fn(|_| Mutex::new(VecDeque::new())),
562            extent_queue: Mutex::new(VecDeque::new()),
563            live_chunks: AtomicU64::new(0),
564            extent_residents: AtomicU64::new(0),
565            enforcing: Mutex::new(()),
566            enforce_pending: std::sync::atomic::AtomicBool::new(false),
567            counters: Counters::default(),
568            spill: Spill::default(),
569        })))
570    }
571
572    /// Allocates a chunk of `len` words and fills it in place: `fill`
573    /// receives the chunk's slot memory directly and must overwrite all of
574    /// it (the slot's prior contents are unspecified), so serialization
575    /// writes its single copy straight into pool memory. The returned handle
576    /// starts `UnbackedResident`. A zero `len` returns a length-0 handle
577    /// holding no slot; payloads beyond the largest size class fall back to
578    /// a plain heap allocation, always resident, a prototype limitation.
579    /// `hints` steer eviction and write-behind policy; callers without
580    /// placement knowledge pass the default. `codec` is the chunk's
581    /// [`ExtentCodec`], fixed for its lifetime: the pool invokes it whenever
582    /// the chunk moves across the extent boundary, and takes no interest in
583    /// the stored form it produces.
584    ///
585    /// Relies on abort-on-panic: a panic in `fill` that was caught would
586    /// leak the slot and its resident-bytes accounting. All hosting
587    /// binaries abort via `mz_ore::panic::install_enhanced_handler`, and
588    /// pool consumers are dataflow operators, never code hosted under a
589    /// `catch_unwind` boundary the way the optimizer is.
590    pub fn insert_with(
591        &self,
592        len: usize,
593        hints: ChunkHints,
594        codec: &'static dyn ExtentCodec,
595        fill: impl FnOnce(&mut [u64]),
596    ) -> ChunkHandle {
597        let inner = &self.0;
598        inner.counters.inserts.fetch_add(1, Ordering::Relaxed);
599        let len_bytes = len * std::mem::size_of::<u64>();
600        if len == 0 {
601            fill(&mut []);
602            let meta = ChunkMeta::new(
603                inner,
604                0,
605                None,
606                hints.depth,
607                codec,
608                Residency::UnbackedResident,
609                None,
610                None,
611            );
612            return ChunkHandle {
613                meta: Arc::new(meta),
614            };
615        }
616        let class = region::size_class_for(len_bytes);
617        if class.is_none() {
618            inner
619                .counters
620                .oversize_payloads
621                .fetch_add(1, Ordering::Relaxed);
622        }
623        // A class with no free slot degrades to the heap path below: an
624        // unpageable chunk beats a dead replica.
625        let slot = class.and_then(|class| inner.alloc_slot(class, len_bytes));
626        // Whichever home the payload found, it is resident.
627        inner
628            .counters
629            .resident_bytes
630            .fetch_add(u64::cast_from(len_bytes), Ordering::Relaxed);
631        let meta = match (class, slot) {
632            (Some(class), Some(slot)) => {
633                let region = &inner.regions[class];
634                // SAFETY: the freshly allocated slot is at least `len_bytes`
635                // long (the class fits the payload) and is exclusively owned
636                // by this not-yet-shared chunk, so the mutable borrow is
637                // unique; region memory is mapped and writable, and `u64` has
638                // no validity requirements beyond size, so exposing the
639                // unspecified prior contents through `&mut [u64]` is sound.
640                let dst = unsafe {
641                    std::slice::from_raw_parts_mut(region.slot_ptr(slot).cast::<u64>(), len)
642                };
643                // The fill contract (overwrite all `len` words) is
644                // discipline-only. Poison in debug builds so an
645                // under-writing fill reads back as deterministic garbage
646                // instead of a previous occupant's bytes, which the heap
647                // path's zero fill would otherwise mask in tests.
648                #[cfg(debug_assertions)]
649                dst.fill(u64::from_ne_bytes([0xDE; 8]));
650                fill(dst);
651                ChunkMeta::new(
652                    inner,
653                    len,
654                    Some(class),
655                    hints.depth,
656                    codec,
657                    Residency::UnbackedResident,
658                    Some(slot),
659                    None,
660                )
661            }
662            _ => {
663                let mut payload = vec![0u64; len];
664                fill(&mut payload);
665                inner
666                    .counters
667                    .oversize_bytes
668                    .fetch_add(u64::cast_from(len_bytes), Ordering::Relaxed);
669                ChunkMeta::new(
670                    inner,
671                    len,
672                    None,
673                    hints.depth,
674                    codec,
675                    Residency::Oversize,
676                    None,
677                    Some(payload),
678                )
679            }
680        };
681        let meta = Arc::new(meta);
682        if meta.class.is_some() {
683            inner.live_chunks.fetch_add(1, Ordering::Relaxed);
684            inner
685                .queue(band(meta.depth))
686                .push_back(Arc::downgrade(&meta));
687        }
688        inner.enforce_budget();
689        ChunkHandle { meta }
690    }
691
692    /// Snapshot of the pool's counters.
693    pub fn stats(&self) -> PoolStats {
694        let c = &self.0.counters;
695        PoolStats {
696            inserts: c.inserts.load(Ordering::Relaxed),
697            frees: c.frees.load(Ordering::Relaxed),
698            writes_elided: c.writes_elided.load(Ordering::Relaxed),
699            evictions_compress: c.evictions_compress.load(Ordering::Relaxed),
700            evictions_cheap: c.evictions_cheap.load(Ordering::Relaxed),
701            extent_bytes_written: c.extent_bytes_written.load(Ordering::Relaxed),
702            resident_bytes: c.resident_bytes.load(Ordering::Relaxed),
703            oversize_bytes: c.oversize_bytes.load(Ordering::Relaxed),
704            warm_bytes: c.warm_bytes.load(Ordering::Relaxed),
705            warm_reuses: c.warm_reuses.load(Ordering::Relaxed),
706            eager_backs: c.eager_backs.load(Ordering::Relaxed),
707            admissions_budget: c.admissions_budget.load(Ordering::Relaxed),
708            admissions_steal: c.admissions_steal.load(Ordering::Relaxed),
709            admissions_denied: c.admissions_denied.load(Ordering::Relaxed),
710            extent_resident_bytes: c.extent_resident_bytes.load(Ordering::Relaxed),
711            extent_unreclaimable_bytes: c.extent_unreclaimable_bytes.load(Ordering::Relaxed),
712            extent_pageouts: c.extent_pageouts.load(Ordering::Relaxed),
713            extent_pageout_incomplete: c.extent_pageout_incomplete.load(Ordering::Relaxed),
714            extent_arena_fallbacks: self.0.extent_arena.fallbacks(),
715            spill_scheduled: c.spill_scheduled.load(Ordering::Relaxed),
716            spill_cancelled: c.spill_cancelled.load(Ordering::Relaxed),
717            spill_in_flight: self.0.spill.in_flight.load(Ordering::Relaxed),
718            slot_exhausted_fallbacks: c.slot_exhausted_fallbacks.load(Ordering::Relaxed),
719            oversize_payloads: c.oversize_payloads.load(Ordering::Relaxed),
720            live_chunks: self.0.live_chunks.load(Ordering::Relaxed),
721        }
722    }
723
724    /// Enables or disables off-worker eviction I/O. The first call with
725    /// `threads > 0` spawns that many spill threads (spawn-once: later calls
726    /// only toggle participation); `threads == 0` falls back to inline
727    /// eviction on the caller for subsequent victims, letting any queued
728    /// work drain.
729    pub fn set_spill_threads(&self, threads: usize) {
730        if threads == 0 {
731            self.0.spill.enabled.store(false, Ordering::Relaxed);
732            return;
733        }
734        let spawned = self.0.spill.threads.load(Ordering::Relaxed);
735        if spawned == 0 {
736            let to_spawn = u64::cast_from(threads);
737            if self
738                .0
739                .spill
740                .threads
741                .compare_exchange(0, to_spawn, Ordering::Relaxed, Ordering::Relaxed)
742                .is_ok()
743            {
744                for i in 0..threads {
745                    let inner = Arc::clone(&self.0);
746                    let handle = std::thread::Builder::new()
747                        .name(format!("pool-spill-{i}"))
748                        .spawn(move || inner.spill_worker())
749                        .expect("spawn pool spill thread");
750                    #[cfg(test)]
751                    self.0
752                        .spill
753                        .handles
754                        .lock()
755                        .expect("spill handles poisoned")
756                        .push(handle);
757                    #[cfg(not(test))]
758                    drop(handle);
759                }
760            }
761        }
762        self.0.spill.enabled.store(true, Ordering::Relaxed);
763    }
764
765    /// Enables or disables eager backing: when on, idle spill threads
766    /// compress unbacked chunks to `BackedResident` ahead of pressure, so
767    /// budget-driven eviction becomes a pure page release. Costs CPU on
768    /// chunks that die before eviction would have reached them; pays at
769    /// every pressure event. Only meaningful with spill threads spawned.
770    pub fn set_eager_backing(&self, eager: bool) {
771        self.0.spill.eager.store(eager, Ordering::Relaxed);
772        if eager {
773            self.0.spill.cv.notify_all();
774        }
775    }
776
777    /// Test hook: performs one eager-backing step on the calling thread.
778    /// Returns whether progress was made.
779    #[cfg(test)]
780    fn back_step(&self) -> bool {
781        self.0.back_one()
782    }
783
784    /// Test hook: waits until the spill queue is empty and no entry is being
785    /// processed, so tests observe deterministic post-eviction states.
786    #[cfg(test)]
787    fn quiesce_spill(&self) {
788        while self.0.spill.in_flight.load(Ordering::Relaxed) > 0 {
789            std::thread::yield_now();
790        }
791    }
792
793    /// Test hook: stops and joins the spill threads, so a test binary exits
794    /// with none alive (which Miri requires). Stopped threads process no
795    /// further queued work; call [`Pool::quiesce_spill`] first when the test
796    /// depends on the queue draining.
797    #[cfg(test)]
798    fn join_spill_threads(&self) {
799        self.0.spill.stop.store(true, Ordering::Relaxed);
800        self.0.spill.cv.notify_all();
801        let handles =
802            std::mem::take(&mut *self.0.spill.handles.lock().expect("spill handles poisoned"));
803        for handle in handles {
804            handle.join().expect("spill thread panicked");
805        }
806    }
807
808    /// Test hook: enables spill scheduling without spawning threads, so tests
809    /// drive the queue deterministically via [`Pool::spill_step`].
810    #[cfg(test)]
811    fn enable_spill_without_threads(&self) {
812        self.0.spill.enabled.store(true, Ordering::Relaxed);
813    }
814
815    /// Test hook: processes one queued spill entry on the calling thread.
816    /// Returns whether an entry was processed.
817    #[cfg(test)]
818    fn spill_step(&self) -> bool {
819        let popped = self.0.spill_queue().pop_front();
820        let Some(meta) = popped else {
821            return false;
822        };
823        self.0.spill_process(&meta, SpillKind::Evict);
824        self.0.spill.in_flight.fetch_sub(1, Ordering::Relaxed);
825        true
826    }
827
828    /// Test hook: runs one compressed-cap enforcement pass on the calling
829    /// thread.
830    #[cfg(test)]
831    fn enforce_compressed(&self) {
832        self.0.enforce_compressed_cap();
833    }
834
835    /// Test hook: evicts cold chunks until resident bytes fall to the budget
836    /// or every queued chunk has been visited once. Enforcement runs
837    /// automatically on every insert and budget shrink.
838    #[cfg(test)]
839    fn enforce_budget(&self) {
840        self.0.enforce_budget();
841    }
842
843    /// Test hook: runs one compressed-cap enforcement pass inline on the
844    /// calling thread, where the fake residency observation applies.
845    #[cfg(test)]
846    fn enforce_rss_target(&self) {
847        self.0.enforce_compressed_cap();
848    }
849
850    /// Retunes the resident-bytes budget in place and enforces it. Live
851    /// handles share the new value immediately through their `Arc<PoolInner>`;
852    /// a shrink takes effect by evicting on this call, a grow simply leaves
853    /// more headroom for future inserts.
854    pub fn set_budget(&self, budget_bytes: usize) {
855        let new = u64::cast_from(budget_bytes);
856        let prev = self.0.budget_bytes.swap(new, Ordering::Relaxed);
857        // Config application calls this per worker per tick; only a change
858        // warrants an enforcement pass (a grow needs none, and inserts
859        // enforce continuously anyway).
860        if new < prev {
861            self.0.trim_warm_pool();
862            self.0.enforce_budget();
863        }
864    }
865
866    /// Retunes the ceiling on the pool's total RSS — slots plus warm slots
867    /// plus compressed-resident extents. The compressed tier's capacity is
868    /// the gap above the budget and warm cap; zero (the default) collapses
869    /// the tier, paging extents out as soon as they are written. A shrink
870    /// takes effect by paging out the oldest extents on this call.
871    pub fn set_rss_target(&self, target_bytes: usize) {
872        let new = u64::cast_from(target_bytes);
873        let prev = self.0.rss_target_bytes.swap(new, Ordering::Relaxed);
874        if new < prev {
875            self.0.enforce_compressed_cap();
876        }
877    }
878
879    /// Test-only: the number of entries across the second-chance queues,
880    /// live and stale.
881    #[cfg(test)]
882    fn queue_len(&self) -> usize {
883        (0..DEPTH_BANDS).map(|band| self.0.queue(band).len()).sum()
884    }
885
886    /// Test-only: the number of resident-extent queue entries, live and
887    /// stale.
888    #[cfg(test)]
889    fn extent_queue_len(&self) -> usize {
890        self.0.extent_queue().len()
891    }
892
893    /// Test hook: explicitly evicts one chunk. No-op if the chunk is already
894    /// evicted, in flight, empty, or oversize. With spill threads enabled the
895    /// compression is handed off and completes asynchronously (observable via
896    /// [`Residency::WriteInFlight`]); without them it runs inline.
897    #[cfg(test)]
898    fn evict(&self, handle: &ChunkHandle) {
899        let meta = &handle.meta;
900        let mut state = meta.state();
901        if !meta.pool.spill_handoff(meta, &mut state) {
902            meta.pool.evict_locked(meta, &mut state);
903        }
904        drop(state);
905        meta.pool.enforce_or_defer_compressed_cap();
906    }
907
908    /// Test hook: overwrites every free slot's bytes with `0xDE`. The free
909    /// list keeps a freed slot's old bytes on platforms where
910    /// `MADV_DONTNEED` retains contents (macOS); poisoning lets tests prove
911    /// that reads of evicted chunks decompress from the extent rather than
912    /// passing stale slot memory through.
913    #[cfg(test)]
914    fn poison_free_slots(&self) {
915        for region in &self.0.regions {
916            region.poison_free_slots();
917        }
918    }
919}
920
921impl PoolInner {
922    /// Locks the eviction queue of one depth band.
923    fn queue(&self, band: usize) -> MutexGuard<'_, VecDeque<Weak<ChunkMeta>>> {
924        self.queues[band].lock().expect("pool queue poisoned")
925    }
926
927    /// Locks the resident-extent queue.
928    fn extent_queue(&self) -> MutexGuard<'_, VecDeque<Weak<ChunkMeta>>> {
929        self.extent_queue.lock().expect("extent queue poisoned")
930    }
931
932    /// Locks the spill hand-off queue.
933    fn spill_queue(&self) -> MutexGuard<'_, VecDeque<Arc<ChunkMeta>>> {
934        self.spill.queue.lock().expect("spill queue poisoned")
935    }
936
937    /// The region behind a slotted chunk's size class.
938    fn region_of(&self, meta: &ChunkMeta) -> &Region {
939        &self.regions[meta.class.expect("slotted chunk has a class")]
940    }
941
942    /// Borrows the payload of a slotted chunk.
943    ///
944    /// # Safety
945    ///
946    /// `slot` must be `meta`'s slot, its contents must be initialized (they
947    /// are from insert onward), and nothing may write the slot while the
948    /// borrow lives.
949    unsafe fn slot_data(&self, meta: &ChunkMeta, slot: u32) -> &[u64] {
950        let ptr = self
951            .region_of(meta)
952            .slot_ptr(slot)
953            .cast_const()
954            .cast::<u64>();
955        // SAFETY: per the function contract; `meta.len` words fit the class
956        // by construction.
957        unsafe { std::slice::from_raw_parts(ptr, meta.len) }
958    }
959
960    /// Records a freshly written extent under the chunk's state lock: the
961    /// compressed-bytes counter, the compressed-tier accounting, and the
962    /// state's extent field.
963    fn commit_extent(&self, meta: &Arc<ChunkMeta>, state: &mut ChunkState, extent: SwapExtent) {
964        self.counters
965            .extent_bytes_written
966            .fetch_add(u64::cast_from(extent.comp_len()), Ordering::Relaxed);
967        // A heap-fallback extent is born permanently capped and counts as
968        // unreclaimable from the start.
969        self.note_extent_resident(meta, extent.alloc_size(), !extent.pageout_capped());
970        state.extent = Some(extent);
971    }
972
973    /// Drops queue entries whose chunk has been freed, detected by their
974    /// `Weak` no longer holding a live chunk. Each band compacts only when
975    /// its stale entries outnumber all live chunks (plus a small floor), so
976    /// the cost amortizes to a constant per insert and the total queue
977    /// length stays proportional to the number of live slotted chunks even
978    /// when the pool never comes under budget pressure.
979    fn prune_queues(&self) {
980        let live = usize::cast_from(self.live_chunks.load(Ordering::Relaxed));
981        for band in 0..DEPTH_BANDS {
982            let mut queue = self.queue(band);
983            if queue.len() > 2 * live + 16 {
984                queue.retain(|weak| weak.strong_count() > 0);
985            }
986        }
987    }
988
989    fn enforce_budget(&self) {
990        // Single-flight: enforcement runs synchronously on whichever thread
991        // trips it (every insert), and concurrent passes would
992        // convoy on the queue mutex doing redundant scans of the same
993        // candidates. One pass at a time reaches the budget just as well;
994        // skipped callers hand their bytes to the in-progress pass through
995        // `enforce_pending`. A poisoned claim means a prior pass panicked.
996        // Recover and keep enforcing rather than silently disabling the
997        // budget for the process's lifetime.
998        let guard = match self.enforcing.try_lock() {
999            Ok(guard) => guard,
1000            Err(std::sync::TryLockError::WouldBlock) => {
1001                // Release pairs with the holder's Acquire: the `resident_bytes`
1002                // bump this caller just made must be visible to the re-read.
1003                self.enforce_pending.store(true, Ordering::Release);
1004                return;
1005            }
1006            Err(std::sync::TryLockError::Poisoned(poisoned)) => poisoned.into_inner(),
1007        };
1008        loop {
1009            self.enforce_budget_inner();
1010            #[cfg(test)]
1011            run_enforce_budget_hook();
1012            // A caller turned away since this pass's counter reads may have
1013            // left bytes unenforced; re-run rather than drop them. The Acquire
1014            // pairs with the turned-away Release so the re-read sees the bump,
1015            // and this swap is the only place the flag is cleared, so no set
1016            // can be lost.
1017            //
1018            // NOTE: a caller turned away between this swap and `drop(guard)`
1019            // sets the flag but finds no re-reader. That residual window is a
1020            // few instructions wide, versus the whole pass before.
1021            if !self.enforce_pending.swap(false, Ordering::Acquire) {
1022                break;
1023            }
1024        }
1025        drop(guard);
1026        // Inline evictions above may have grown the compressed tier.
1027        self.enforce_or_defer_compressed_cap();
1028    }
1029
1030    /// Bytes budget enforcement can actually reclaim: resident bytes minus
1031    /// heap-backed (oversize and class-exhaustion) chunks, which hold no
1032    /// slot and can never be evicted. Enforcing against raw resident bytes
1033    /// would, once unevictable bytes alone exceed the budget, compress
1034    /// every slotted chunk on arrival forever.
1035    fn evictable_bytes(&self) -> u64 {
1036        self.counters
1037            .resident_bytes
1038            .load(Ordering::Relaxed)
1039            .saturating_sub(self.counters.oversize_bytes.load(Ordering::Relaxed))
1040    }
1041
1042    fn enforce_budget_inner(&self) {
1043        self.prune_queues();
1044        // Deepest band first: deep chunks are the coldest, and once eager
1045        // backing has visited them (same order) their eviction is a pure
1046        // page release. The youngest band is reached only when the deeper
1047        // bands cannot satisfy the budget, keeping young data's
1048        // die-before-write chance longest.
1049        for band in (0..DEPTH_BANDS).rev() {
1050            if self.evictable_bytes() <= self.budget_bytes.load(Ordering::Relaxed) {
1051                return;
1052            }
1053            self.enforce_budget_band(band);
1054        }
1055    }
1056
1057    fn enforce_budget_band(&self, band: usize) {
1058        // The queue holds resident chunks only (entries for evicted chunks
1059        // are dropped on visit and never re-added), so a full pass is
1060        // proportional to the resident set. Visit each queued chunk at most
1061        // twice per call: a first visit may only clear the second-chance
1062        // bit, so a second is needed before an over-budget call is
1063        // guaranteed to evict every chunk it saw. The bound keeps contended
1064        // and in-flight entries from spinning this loop forever.
1065        let mut remaining = self.queue(band).len().saturating_mul(2);
1066        while remaining > 0 && self.evictable_bytes() > self.budget_bytes.load(Ordering::Relaxed) {
1067            remaining -= 1;
1068            let popped = self.queue(band).pop_front();
1069            let Some(weak) = popped else {
1070                break;
1071            };
1072            let Some(meta) = weak.upgrade() else {
1073                continue;
1074            };
1075            let requeue = {
1076                // `try_lock`: a chunk mid-eviction or mid-read holds its
1077                // lock for milliseconds; skipping it beats convoying every
1078                // budget enforcer in the process behind one chunk's I/O.
1079                let Ok(mut state) = meta.state.try_lock() else {
1080                    self.queue(band).push_back(weak);
1081                    continue;
1082                };
1083                if state.freed {
1084                    false
1085                } else if matches!(state.residency, Residency::Evicted | Residency::Oversize) {
1086                    // Nothing to evict: drop the entry. A chunk re-enters
1087                    // the queue only when it becomes resident again (insert
1088                    // or re-admission), so the queue stays proportional to
1089                    // the resident set rather than accumulating every chunk
1090                    // ever evicted.
1091                    false
1092                } else if state.touched {
1093                    state.touched = false;
1094                    true
1095                } else if self.spill_handoff(&meta, &mut state) {
1096                    // Stays queued while in flight; once the spill commits to
1097                    // `Evicted`, the next visit drops the entry.
1098                    true
1099                } else {
1100                    self.evict_locked(&meta, &mut state);
1101                    state.residency != Residency::Evicted
1102                }
1103            };
1104            if requeue {
1105                self.queue(band).push_back(weak);
1106            }
1107        }
1108    }
1109
1110    fn evict_locked(&self, meta: &Arc<ChunkMeta>, state: &mut ChunkState) {
1111        let Some(slot) = state.slot else {
1112            return;
1113        };
1114        if state.freed {
1115            return;
1116        }
1117        match state.residency {
1118            Residency::UnbackedResident => {
1119                // SAFETY: the slot belongs to this live chunk and the state
1120                // lock is held, so nothing else touches the slot while this
1121                // borrow is live (reads copy out under the same lock).
1122                let data = unsafe { self.slot_data(meta, slot) };
1123                // Inline eviction runs on whichever thread tripped the
1124                // budget, so the compression scratch must not stay parked
1125                // on it.
1126                let extent =
1127                    SwapExtent::write(&self.extent_arena, data, meta.codec, Scratch::Shrink);
1128                self.counters
1129                    .evictions_compress
1130                    .fetch_add(1, Ordering::Relaxed);
1131                self.commit_extent(meta, state, extent);
1132            }
1133            Residency::BackedResident => {
1134                self.counters
1135                    .evictions_cheap
1136                    .fetch_add(1, Ordering::Relaxed);
1137            }
1138            Residency::WriteInFlight | Residency::Evicted | Residency::Oversize => return,
1139        }
1140        // `release_slot`'s precondition holds: the state lock is held and
1141        // `!freed` was checked above under it.
1142        self.release_slot(meta, state);
1143        state.residency = Residency::Evicted;
1144    }
1145
1146    /// Whether the next eviction should be handed to spill threads: enabled,
1147    /// and the queue is below the backpressure bound (beyond it, callers
1148    /// evict inline rather than growing an unbounded queue of still-resident
1149    /// chunks).
1150    fn spill_eligible(&self) -> bool {
1151        self.spill.enabled.load(Ordering::Relaxed)
1152            && usize::cast_from(self.spill.in_flight.load(Ordering::Relaxed)) < SPILL_IN_FLIGHT_MAX
1153    }
1154
1155    /// Hands a `WriteInFlight` chunk to the spill threads.
1156    fn spill_schedule(&self, meta: Arc<ChunkMeta>) {
1157        self.counters
1158            .spill_scheduled
1159            .fetch_add(1, Ordering::Relaxed);
1160        self.spill.in_flight.fetch_add(1, Ordering::Relaxed);
1161        self.spill_queue().push_back(meta);
1162        self.spill.cv.notify_one();
1163    }
1164
1165    /// Spill-thread main loop. The thread owns an `Arc<PoolInner>`, so the
1166    /// pool (a process-wide singleton in production) lives as long as its
1167    /// threads. Queued (budget-driven) evictions take priority; with eager
1168    /// backing enabled, idle threads compress unbacked chunks to
1169    /// `BackedResident` instead of parking, and park with a timeout once
1170    /// everything reachable is backed.
1171    fn spill_worker(self: Arc<Self>) {
1172        loop {
1173            #[cfg(test)]
1174            if self.spill.stop.load(Ordering::Relaxed) {
1175                return;
1176            }
1177            // Tier-2 pageouts ride the spill threads: every pass through the
1178            // loop (job completion, condvar wakeup, park timeout) trims the
1179            // compressed tier if needed. A single atomic load when under cap.
1180            self.enforce_compressed_cap();
1181            let popped = self.spill_queue().pop_front();
1182            if let Some(meta) = popped {
1183                self.spill_process(&meta, SpillKind::Evict);
1184                self.spill.in_flight.fetch_sub(1, Ordering::Relaxed);
1185                continue;
1186            }
1187            if self.spill.eager.load(Ordering::Relaxed) && self.back_one() {
1188                continue;
1189            }
1190            // Nothing to evict or back: park. Re-checking emptiness under
1191            // the queue lock closes the lost-wakeup window (hand-offs push
1192            // under this lock before notifying); the timeout backstops
1193            // everything else (fresh inserts, tier growth, lost notifies).
1194            let queue = self.spill_queue();
1195            if queue.is_empty() {
1196                let _ = self
1197                    .spill
1198                    .cv
1199                    .wait_timeout(queue, std::time::Duration::from_millis(100))
1200                    .expect("spill queue poisoned");
1201            }
1202        }
1203    }
1204
1205    /// Eagerly compresses one unbacked chunk from the eviction queues into
1206    /// `BackedResident`, returning whether a chunk was backed — `false`
1207    /// means nothing was actionable (queues empty, or the bounded scans
1208    /// found only already-backed, in-flight, contended, or stale entries)
1209    /// and the caller should park rather than rescan. Bands are visited
1210    /// deepest first, mirroring eviction order so the chunks evicted first
1211    /// are the ones whose backing is already pre-paid.
1212    fn back_one(&self) -> bool {
1213        for band in (0..DEPTH_BANDS).rev() {
1214            if self.back_one_from(band) {
1215                return true;
1216            }
1217        }
1218        false
1219    }
1220
1221    /// One bounded backing scan over a single band's queue. Non-actionable
1222    /// entries are requeued or dropped per the same rules budget
1223    /// enforcement uses, except that the second-chance `touched` bit is
1224    /// left alone — backing is not an eviction and must not consume a
1225    /// chunk's reprieve.
1226    fn back_one_from(&self, band: usize) -> bool {
1227        for _ in 0..16 {
1228            let popped = self.queue(band).pop_front();
1229            let Some(weak) = popped else {
1230                return false;
1231            };
1232            let Some(meta) = weak.upgrade() else {
1233                continue;
1234            };
1235            {
1236                let Ok(mut state) = meta.state.try_lock() else {
1237                    self.queue(band).push_back(weak);
1238                    continue;
1239                };
1240                if state.freed {
1241                    continue;
1242                }
1243                match state.residency {
1244                    Residency::Evicted | Residency::Oversize => {
1245                        continue;
1246                    }
1247                    Residency::UnbackedResident => {
1248                        state.residency = Residency::WriteInFlight;
1249                    }
1250                    Residency::BackedResident | Residency::WriteInFlight => {
1251                        self.queue(band).push_back(weak);
1252                        continue;
1253                    }
1254                }
1255            }
1256            self.spill.in_flight.fetch_add(1, Ordering::Relaxed);
1257            self.spill_process(&meta, SpillKind::Back);
1258            self.spill.in_flight.fetch_sub(1, Ordering::Relaxed);
1259            // The chunk remains an eviction candidate (now a cheap one).
1260            self.queue(band).push_back(weak);
1261            return true;
1262        }
1263        false
1264    }
1265
1266    /// Performs (or cancels) one scheduled compression. Lock discipline: the
1267    /// chunk lock is held only to validate and to commit — never across the
1268    /// compression or the `pageout` reclaim, which are the multi-millisecond
1269    /// costs this path exists to keep off budget-enforcing threads.
1270    fn spill_process(&self, meta: &Arc<ChunkMeta>, kind: SpillKind) {
1271        // Validate under the lock, then release it for the I/O. The slot is
1272        // captured under the lock and remains owned by this chunk for the
1273        // unlocked compression: in `WriteInFlight`, eviction skips the chunk
1274        // and `ChunkHandle::drop` defers slot release to this thread.
1275        let slot;
1276        {
1277            let mut state = meta.state();
1278            if state.freed {
1279                // Freed while queued: the deferred cleanup is ours, and the
1280                // chunk dies without ever compressing — the write-behind
1281                // cancellation window. `ChunkHandle::drop` already counted
1282                // the free and the live-chunks decrement.
1283                self.counters
1284                    .spill_cancelled
1285                    .fetch_add(1, Ordering::Relaxed);
1286                self.counters.writes_elided.fetch_add(1, Ordering::Relaxed);
1287                self.release_slot(meta, &mut state);
1288                return;
1289            }
1290            if state.residency != Residency::WriteInFlight {
1291                return;
1292            }
1293            slot = state.slot.expect("write-in-flight chunk has a slot");
1294        }
1295        // SAFETY: the chunk is live (the queue holds an `Arc`) and in
1296        // `WriteInFlight`, so the slot is not recycled (`ChunkHandle::drop`
1297        // defers slot release to this thread in that state) and its contents
1298        // are immutable; concurrent copy-out reads take the state lock and
1299        // read the slot, but nothing writes it.
1300        let data = unsafe { self.slot_data(meta, slot) };
1301        // Spill threads see a steady job stream, so they keep the grown
1302        // compression scratch for the next job.
1303        let extent = SwapExtent::write(&self.extent_arena, data, meta.codec, Scratch::Retain);
1304        let mut state = meta.state();
1305        if state.freed {
1306            // Freed during compression: the extent is garbage; cleanup is
1307            // ours as above. Compression ran, so this is not an elided free.
1308            self.counters
1309                .spill_cancelled
1310                .fetch_add(1, Ordering::Relaxed);
1311            self.release_slot(meta, &mut state);
1312            return;
1313        }
1314        self.commit_extent(meta, &mut state, extent);
1315        match kind {
1316            SpillKind::Back => {
1317                // The slot stays for write-behind: the chunk remains
1318                // readable, and the extent makes a later budget eviction a
1319                // pure page release.
1320                self.counters.eager_backs.fetch_add(1, Ordering::Relaxed);
1321                state.residency = Residency::BackedResident;
1322            }
1323            SpillKind::Evict => {
1324                // `release_slot`'s precondition holds: the state lock is
1325                // held and `!freed` was observed under it.
1326                self.counters
1327                    .evictions_compress
1328                    .fetch_add(1, Ordering::Relaxed);
1329                self.release_slot(meta, &mut state);
1330                state.residency = Residency::Evicted;
1331            }
1332        };
1333        drop(state);
1334        // Counted a fresh resident extent: the tier may need trimming. Kept
1335        // here (rather than relying on the spill loop alone) so the
1336        // threadless test hooks observe deterministic post-commit states.
1337        self.enforce_compressed_cap();
1338    }
1339
1340    /// Releases `state`'s slot — slot returned to the region free list,
1341    /// physical pages discarded unless the slot joins the bounded warm pool —
1342    /// and decrements resident bytes. Releasing pages beyond the warm pool is
1343    /// what keeps RSS aligned with the `resident_bytes` gauge the budget
1344    /// enforcer trusts; the warm pool relaxes that alignment by an explicit,
1345    /// bounded amount (`warm_bytes`, capped at a fraction of the budget) so
1346    /// slot reuse faults no pages and skips the kernel's page zeroing.
1347    ///
1348    /// Precondition: the caller holds the chunk's state lock, and no
1349    /// reference into the slot exists — copy-out reads borrow the slot only
1350    /// under that same lock, and a `WriteInFlight` chunk's unlocked
1351    /// compression read belongs to the spill thread, which is the only
1352    /// caller that releases the slot in that state. This is what makes the
1353    /// `dontneed` below sound, and what makes keeping a warm slot's stale
1354    /// contents safe: the slot's next occupant fully overwrites every byte
1355    /// it reads, satisfying the contents-undefined contract either way.
1356    fn release_slot(&self, meta: &ChunkMeta, state: &mut ChunkState) {
1357        let slot = state.slot.take().expect("slotted chunk");
1358        let region = self.region_of(meta);
1359        let warm = self.try_keep_warm(region.class_size());
1360        if !warm {
1361            // SAFETY: no reference into the slot exists (the function-level
1362            // precondition, established under the held state lock).
1363            unsafe {
1364                region::dontneed(region.slot_ptr(slot), region.class_size());
1365            }
1366        }
1367        region.free(slot, warm);
1368        self.counters
1369            .resident_bytes
1370            .fetch_sub(u64::cast_from(meta.len_bytes()), Ordering::Relaxed);
1371    }
1372
1373    /// The warm pool's byte ceiling: an eighth of the budget, clamped at an
1374    /// absolute maximum. The fraction sizes fault amortization at small
1375    /// budgets; the clamp keeps large budgets from parking gigabytes of idle
1376    /// warm slots no fault rate could justify.
1377    fn warm_cap(&self) -> u64 {
1378        (self.budget_bytes.load(Ordering::Relaxed) / 8).min(1 << 30)
1379    }
1380
1381    /// Cools warm free slots until `warm_bytes` falls to the warm cap. A
1382    /// budget shrink lowers the cap, and warm capacity is checked only when
1383    /// a slot is freed, so without this pass slots parked under the old cap
1384    /// would hold their pages until same-class reuse happened to drain them,
1385    /// exactly when the shrink wanted the memory back.
1386    fn trim_warm_pool(&self) {
1387        let mut over = self
1388            .counters
1389            .warm_bytes
1390            .load(Ordering::Relaxed)
1391            .saturating_sub(self.warm_cap());
1392        for region in &self.regions {
1393            if over == 0 {
1394                return;
1395            }
1396            let cooled = u64::cast_from(region.cool_warm_slots(usize::cast_from(over)));
1397            self.counters
1398                .warm_bytes
1399                .fetch_sub(cooled, Ordering::Relaxed);
1400            over = over.saturating_sub(cooled);
1401        }
1402    }
1403
1404    /// Claims warm-pool capacity for a slot of `class_size` bytes, returning
1405    /// whether the slot may keep its pages. The RSS overshoot the warm pool
1406    /// introduces is bounded by [`PoolInner::warm_cap`] and visible as the
1407    /// `warm_bytes` stat.
1408    fn try_keep_warm(&self, class_size: usize) -> bool {
1409        let cap = self.warm_cap();
1410        let class_bytes = u64::cast_from(class_size);
1411        self.counters
1412            .warm_bytes
1413            .try_update(Ordering::Relaxed, Ordering::Relaxed, |cur| {
1414                (cur + class_bytes <= cap).then_some(cur + class_bytes)
1415            })
1416            .is_ok()
1417    }
1418
1419    /// Allocates a slot in `class` for a payload of `len_bytes` with
1420    /// warm-pool accounting (a warm allocation is counted as a reuse and
1421    /// trimmed to the payload), or `None` when the class has no free slot.
1422    fn try_alloc_slot(&self, class: usize, len_bytes: usize) -> Option<u32> {
1423        let (index, warm) = self.regions[class].alloc()?;
1424        if warm {
1425            let class_bytes = u64::cast_from(self.regions[class].class_size());
1426            self.counters
1427                .warm_bytes
1428                .fetch_sub(class_bytes, Ordering::Relaxed);
1429            self.counters.warm_reuses.fetch_add(1, Ordering::Relaxed);
1430            // A warm slot keeps the prior occupant's resident pages, which
1431            // may extend past the new payload while the ledger credits only
1432            // `len_bytes`.
1433            self.trim_slot_tail(class, index, len_bytes);
1434        }
1435        Some(index)
1436    }
1437
1438    /// Releases a slot's pages beyond the first `len_bytes` (rounded up to
1439    /// a page), so a slot reused for a smaller payload does not keep its
1440    /// prior occupant's tail pages resident with no bytes in the ledger to
1441    /// answer for them.
1442    ///
1443    /// Precondition: the caller exclusively owns the slot (freshly
1444    /// allocated, or taken from a victim under the victim's state lock)
1445    /// with no reference into it.
1446    fn trim_slot_tail(&self, class: usize, slot: u32, len_bytes: usize) {
1447        let region = &self.regions[class];
1448        // Hugepage-class slots trim at huge-page granularity: a base-page
1449        // trim would split the slot's `MADV_HUGEPAGE` folios, and khugepaged
1450        // may later re-collapse a partially trimmed range, re-instantiating
1451        // pages the ledger counts as released. Whole-folio trims leave no
1452        // partial folio to split or resurrect.
1453        let granule = if region.class_size() >= region::HUGE_PAGE {
1454            region::HUGE_PAGE
1455        } else {
1456            region::page_size()
1457        };
1458        let keep = len_bytes.next_multiple_of(granule).min(region.class_size());
1459        let tail = region.class_size() - keep;
1460        if tail == 0 {
1461            return;
1462        }
1463        // SAFETY: the caller exclusively owns the slot per the
1464        // precondition, and `keep + tail` is exactly the class size, so the
1465        // range stays within the slot.
1466        unsafe {
1467            region::dontneed(region.slot_ptr(slot).add(keep), tail);
1468        }
1469    }
1470
1471    /// Allocates a slot in `class` for an insert: as
1472    /// [`PoolInner::try_alloc_slot`], with an exhausted class counted as a
1473    /// heap fallback for a `len_bytes` payload (warned about once). `None`
1474    /// means the caller must degrade to the heap.
1475    fn alloc_slot(&self, class: usize, len_bytes: usize) -> Option<u32> {
1476        match self.try_alloc_slot(class, len_bytes) {
1477            Some(index) => Some(index),
1478            None => {
1479                self.counters
1480                    .slot_exhausted_fallbacks
1481                    .fetch_add(1, Ordering::Relaxed);
1482                static EXHAUSTED_ONCE: std::sync::Once = std::sync::Once::new();
1483                EXHAUSTED_ONCE.call_once(|| {
1484                    tracing::warn!(
1485                        len_bytes,
1486                        "buffer pool size class exhausted; falling back to heap chunks \
1487                         (raise the pool's per-class virtual reservation)",
1488                    );
1489                });
1490                None
1491            }
1492        }
1493    }
1494
1495    /// Acquires a slot for re-admitting an evicted chunk, from free budget
1496    /// headroom or by stealing a clean backed victim's slot, never by
1497    /// evicting or compressing anything. `None` counts a denied admission.
1498    /// On success the admitted chunk's resident-bytes accounting and the
1499    /// admission counter are settled, and the caller (who holds the chunk's
1500    /// state lock) owns the slot: its contents are unspecified (fresh,
1501    /// warm, or the victim's stale bytes) and must be fully overwritten.
1502    fn admit_slot(&self, meta: &ChunkMeta) -> Option<u32> {
1503        let class = meta.class.expect("evicted chunk has a class");
1504        let len_bytes = u64::cast_from(meta.len_bytes());
1505        // Free budget first: reserve the bytes, then a slot. The
1506        // reservation never pushes resident bytes past the budget, and a
1507        // class with no free slot hands the reservation back rather than
1508        // evicting anything to make room. The headroom test uses evictable
1509        // bytes, matching budget enforcement: unevictable heap-backed bytes
1510        // must not permanently veto budget-path admissions the enforcer
1511        // would never need to undo.
1512        let reserved = self
1513            .counters
1514            .resident_bytes
1515            .try_update(Ordering::Relaxed, Ordering::Relaxed, |cur| {
1516                // Loaded inside the closure so a CAS retry sees oversize
1517                // frees that landed since the last attempt.
1518                let oversize = self.counters.oversize_bytes.load(Ordering::Relaxed);
1519                let next = cur.checked_add(len_bytes)?;
1520                (next.saturating_sub(oversize) <= self.budget_bytes.load(Ordering::Relaxed))
1521                    .then_some(next)
1522            })
1523            .is_ok();
1524        if reserved {
1525            if let Some(slot) = self.try_alloc_slot(class, meta.len_bytes()) {
1526                self.counters
1527                    .admissions_budget
1528                    .fetch_add(1, Ordering::Relaxed);
1529                return Some(slot);
1530            }
1531            self.counters
1532                .resident_bytes
1533                .fetch_sub(len_bytes, Ordering::Relaxed);
1534        }
1535        if let Some(slot) = self.steal_clean_victim(class, meta.len_bytes()) {
1536            // The slot's physical pages transfer deliberately, but only up
1537            // to the admitted payload: the victim's pages past it would
1538            // stay resident with no ledger bytes to answer for them.
1539            self.trim_slot_tail(class, slot, meta.len_bytes());
1540            self.counters
1541                .admissions_steal
1542                .fetch_add(1, Ordering::Relaxed);
1543            return Some(slot);
1544        }
1545        self.counters
1546            .admissions_denied
1547            .fetch_add(1, Ordering::Relaxed);
1548        None
1549    }
1550
1551    /// Takes the slot of a clean victim in `class` for an admitted payload
1552    /// of `admitted_len_bytes`: a `BackedResident` chunk with a clear
1553    /// touched bit, whose extent already duplicates its slot, so the victim
1554    /// transitions to `Evicted` with zero I/O, its extent intact, and its
1555    /// queue entry dropped. The returned slot keeps its physical pages (no
1556    /// `dontneed`, no free-list round trip); they hold the victim's stale
1557    /// bytes. `None` when the bounded scan finds no such victim, or none
1558    /// whose growth the budget can absorb.
1559    ///
1560    /// The caller holds its own chunk's state lock. The scan follows the
1561    /// enforcement discipline (deepest band first, queue guard dropped
1562    /// before any chunk lock, victims only ever `try_lock`ed), which is
1563    /// what keeps the chunk-lock-while-probing-chunk-lock window
1564    /// deadlock-free: two admitters stealing toward each other both fail
1565    /// the `try_lock` and skip. Unlike enforcement, the scan rotates
1566    /// unsuitable entries (touched, wrong class, unbacked) to the back
1567    /// without spending touched bits, shuffling FIFO order the way the
1568    /// backing scan does.
1569    fn steal_clean_victim(&self, class: usize, admitted_len_bytes: usize) -> Option<u32> {
1570        // Bound on entries examined, per band rather than shared across the
1571        // scan: a deep band dense with touched resident chunks would
1572        // otherwise spend the whole scan on hopeless candidates and starve
1573        // the shallower bands where eager backing stocks the clean victims.
1574        // Eight visits absorb a few lock-busy or freshly touched entries
1575        // without degrading a hopeless scan into a full queue walk.
1576        const VISITS_PER_BAND: usize = 8;
1577        for band in (0..DEPTH_BANDS).rev() {
1578            let mut visits = VISITS_PER_BAND;
1579            while visits > 0 {
1580                let popped = self.queue(band).pop_front();
1581                let Some(weak) = popped else {
1582                    // Band exhausted; the next band has its own budget.
1583                    break;
1584                };
1585                let Some(meta) = weak.upgrade() else {
1586                    // Stale entries drop for free and do not spend a visit.
1587                    continue;
1588                };
1589                visits -= 1;
1590                let Ok(mut state) = meta.state.try_lock() else {
1591                    self.queue(band).push_back(weak);
1592                    continue;
1593                };
1594                if state.freed {
1595                    continue;
1596                }
1597                match state.residency {
1598                    // Entries for non-resident chunks drop, as in
1599                    // enforcement.
1600                    Residency::Evicted | Residency::Oversize => continue,
1601                    Residency::UnbackedResident | Residency::WriteInFlight => {
1602                        self.queue(band).push_back(weak);
1603                        continue;
1604                    }
1605                    Residency::BackedResident => {}
1606                }
1607                if state.touched || meta.class != Some(class) {
1608                    self.queue(band).push_back(weak);
1609                    continue;
1610                }
1611                // Settle the ledger in one step: the victim's bytes out, the
1612                // admitted payload's in. A steal that grows resident bytes
1613                // is an admission and must fit the budget (against evictable
1614                // bytes, as everywhere); a shrinking steal always may
1615                // proceed. On failure the victim is requeued untouched.
1616                let victim_len = u64::cast_from(meta.len_bytes());
1617                let admitted_len = u64::cast_from(admitted_len_bytes);
1618                let settled = self
1619                    .counters
1620                    .resident_bytes
1621                    .try_update(Ordering::Relaxed, Ordering::Relaxed, |cur| {
1622                        let next = cur.checked_add(admitted_len)?.saturating_sub(victim_len);
1623                        let oversize = self.counters.oversize_bytes.load(Ordering::Relaxed);
1624                        (next <= cur
1625                            || next.saturating_sub(oversize)
1626                                <= self.budget_bytes.load(Ordering::Relaxed))
1627                        .then_some(next)
1628                    })
1629                    .is_ok();
1630                if !settled {
1631                    self.queue(band).push_back(weak);
1632                    continue;
1633                }
1634                let slot = state.slot.take().expect("backed chunk has a slot");
1635                state.residency = Residency::Evicted;
1636                return Some(slot);
1637            }
1638        }
1639        None
1640    }
1641
1642    /// Capacity of the compressed-resident tier: the RSS target's headroom
1643    /// above the slot budget and the warm cap. With no target set the tier
1644    /// has zero capacity, so extents page out as soon as they are written.
1645    fn compressed_cap(&self) -> u64 {
1646        let target = self.rss_target_bytes.load(Ordering::Relaxed);
1647        let floor = self
1648            .budget_bytes
1649            .load(Ordering::Relaxed)
1650            .saturating_add(self.warm_cap());
1651        target.saturating_sub(floor)
1652    }
1653
1654    /// Counts a newly resident extent (written, or revived by a read)
1655    /// against the compressed tier. A reclaimable extent additionally
1656    /// enqueues its chunk for RSS-target enforcement; an unreclaimable one
1657    /// (a heap-fallback extent, which is never advised out) counts against
1658    /// the unreclaimable gauge instead and stays out of the queue, so
1659    /// enforcement never walks entries it cannot act on. Callers hold the
1660    /// chunk's state lock with the extent present and resident, and follow
1661    /// up with [`PoolInner::enforce_compressed_cap`] once the lock is
1662    /// released.
1663    ///
1664    /// Invariant: `extent_resident_bytes` equals the sum of `alloc_size`
1665    /// over live chunks' extents whose `is_resident()` is true, and
1666    /// `extent_residents` counts those extents; `extent_unreclaimable_bytes`
1667    /// is the subset whose `pageout_capped()` is true. This method,
1668    /// [`PoolInner::note_extent_reclaimable`],
1669    /// [`PoolInner::note_extent_released`], and the pageout arms in
1670    /// [`PoolInner::enforce_compressed_cap`] are the only adjusters; every
1671    /// flag flip pairs with one of them under the chunk's state lock.
1672    fn note_extent_resident(&self, meta: &Arc<ChunkMeta>, extent_alloc: usize, reclaimable: bool) {
1673        self.counters
1674            .extent_resident_bytes
1675            .fetch_add(u64::cast_from(extent_alloc), Ordering::Relaxed);
1676        self.extent_residents.fetch_add(1, Ordering::Relaxed);
1677        if reclaimable {
1678            self.prune_extent_queue();
1679            self.extent_queue().push_back(Arc::downgrade(meta));
1680        } else {
1681            self.counters
1682                .extent_unreclaimable_bytes
1683                .fetch_add(u64::cast_from(extent_alloc), Ordering::Relaxed);
1684        }
1685    }
1686
1687    /// Returns a retry-capped resident extent to the reclaimable set after a
1688    /// read restored its pageout budget: uncounts it from the unreclaimable
1689    /// gauge and re-enqueues its chunk for RSS-target enforcement. The
1690    /// caller holds the chunk's state lock with the extent present, resident,
1691    /// and no longer `pageout_capped()`.
1692    fn note_extent_reclaimable(&self, meta: &Arc<ChunkMeta>, extent_alloc: usize) {
1693        self.counters
1694            .extent_unreclaimable_bytes
1695            .fetch_sub(u64::cast_from(extent_alloc), Ordering::Relaxed);
1696        self.prune_extent_queue();
1697        self.extent_queue().push_back(Arc::downgrade(meta));
1698    }
1699
1700    /// Uncounts a resident extent that is being dropped (chunk freed or
1701    /// degraded). Its queue entry goes stale and is dropped on visit or by
1702    /// [`PoolInner::prune_extent_queue`].
1703    fn note_extent_released(&self, extent: &SwapExtent) {
1704        if extent.is_resident() {
1705            self.counters
1706                .extent_resident_bytes
1707                .fetch_sub(u64::cast_from(extent.alloc_size()), Ordering::Relaxed);
1708            self.extent_residents.fetch_sub(1, Ordering::Relaxed);
1709            if extent.pageout_capped() {
1710                self.counters
1711                    .extent_unreclaimable_bytes
1712                    .fetch_sub(u64::cast_from(extent.alloc_size()), Ordering::Relaxed);
1713            }
1714        }
1715    }
1716
1717    /// Drops extent-queue entries whose chunk has been freed, mirroring
1718    /// [`PoolInner::prune_queues`]: compact only when the queue outgrows
1719    /// all live resident extents (plus a small floor), so the cost
1720    /// amortizes to a constant per push. Enforcement drops stale entries
1721    /// too, but only while the tier is over capacity. A pool that stays
1722    /// under its compressed cap would otherwise accumulate an entry (and a
1723    /// pin on the dead chunk's allocation) per freed extent forever.
1724    fn prune_extent_queue(&self) {
1725        let live = usize::cast_from(self.extent_residents.load(Ordering::Relaxed));
1726        let mut queue = self.extent_queue();
1727        if queue.len() > 2 * live + 16 {
1728            queue.retain(|weak| weak.strong_count() > 0);
1729        }
1730    }
1731
1732    /// Routes compressed-cap enforcement off latency-sensitive threads: with
1733    /// spill threads spawned, wakes one to perform the pageouts
1734    /// (`MADV_PAGEOUT` is synchronous reclaim, bounded per extent but not
1735    /// free at chunk rates); without them, enforces inline. The test is for
1736    /// thread existence, not `spill.enabled`: spawned threads trim the tier
1737    /// in their loop even with eviction hand-off disabled.
1738    ///
1739    /// Deferral makes the target eventually-enforced with bounded lag, and
1740    /// the backstop below turns the lag into a bound by construction: a
1741    /// caller finding the reclaimable tier at double its capacity enforces
1742    /// inline regardless, so sustained creation can never outrun trimming
1743    /// by more than one capacity's worth.
1744    fn enforce_or_defer_compressed_cap(&self) {
1745        if self.spill.threads.load(Ordering::Relaxed) > 0 {
1746            // The inline backstop keys on the bytes enforcement can actually
1747            // reclaim. Unreclaimable extents (retry-capped, heap-backed)
1748            // would otherwise hold the backstop permanently over threshold
1749            // and put a full enforcement pass on every caller.
1750            let resident = self.counters.extent_resident_bytes.load(Ordering::Relaxed);
1751            let unreclaimable = self
1752                .counters
1753                .extent_unreclaimable_bytes
1754                .load(Ordering::Relaxed);
1755            if resident.saturating_sub(unreclaimable) > self.compressed_cap().saturating_mul(2) {
1756                self.enforce_compressed_cap();
1757            } else {
1758                self.spill.cv.notify_one();
1759            }
1760        } else {
1761            self.enforce_compressed_cap();
1762        }
1763    }
1764
1765    /// Pages out the oldest resident extents until the compressed tier falls
1766    /// to its capacity. The compression is already paid and the device write
1767    /// is the kernel's async writeback, so each pageout is one bounded
1768    /// madvise plus a page-table observation; spill threads run this between
1769    /// jobs, and other threads only when no spill threads exist (see
1770    /// [`PoolInner::enforce_or_defer_compressed_cap`]). Not single-flighted:
1771    /// concurrent passes pop disjoint victims. Visits are bounded by the
1772    /// queue's length at entry; stale entries (extent paged out, dropped, or
1773    /// chunk dead) are dropped. Incomplete extents are requeued with their
1774    /// accounting intact until their retry budget runs out, at which point
1775    /// they leave the queue with their bytes on the unreclaimable gauge, so
1776    /// the tier may settle above its capacity by the bytes the kernel
1777    /// declined to reclaim without enforcement re-walking them.
1778    fn enforce_compressed_cap(&self) {
1779        let cap = self.compressed_cap();
1780        let resident = |c: &Counters| c.extent_resident_bytes.load(Ordering::Relaxed);
1781        // Under-cap is the common case: answer it with one atomic load and
1782        // no queue lock, so frequent callers (the spill loop) stay cheap.
1783        if resident(&self.counters) <= cap {
1784            return;
1785        }
1786        let mut remaining = self.extent_queue().len();
1787        while remaining > 0 && resident(&self.counters) > cap {
1788            remaining -= 1;
1789            let popped = self.extent_queue().pop_front();
1790            let Some(weak) = popped else {
1791                break;
1792            };
1793            let Some(meta) = weak.upgrade() else {
1794                continue;
1795            };
1796            // `try_lock`: a chunk mid-read or mid-compression holds its lock
1797            // for milliseconds; requeue rather than convoy behind it.
1798            let Ok(mut state) = meta.state.try_lock() else {
1799                self.extent_queue().push_back(weak);
1800                continue;
1801            };
1802            match &mut state.extent {
1803                Some(extent) if extent.is_resident() => {
1804                    if extent.pageout_capped() {
1805                        // A leftover entry for an already-capped extent (its
1806                        // capping transition below accounted it and dropped
1807                        // its entry): drop this one too. The read that
1808                        // restores the retry budget re-enqueues the chunk.
1809                    } else if extent.pageout() {
1810                        self.counters
1811                            .extent_resident_bytes
1812                            .fetch_sub(u64::cast_from(extent.alloc_size()), Ordering::Relaxed);
1813                        self.extent_residents.fetch_sub(1, Ordering::Relaxed);
1814                        self.counters
1815                            .extent_pageouts
1816                            .fetch_add(1, Ordering::Relaxed);
1817                    } else {
1818                        // The advice left pages resident. The extent keeps
1819                        // its full accounting (the ledger may over-count
1820                        // RSS, the safe direction).
1821                        self.counters
1822                            .extent_pageout_incomplete
1823                            .fetch_add(1, Ordering::Relaxed);
1824                        if extent.pageout_capped() {
1825                            // The retry budget just ran out: the extent
1826                            // leaves the queue and its bytes move to the
1827                            // unreclaimable gauge, so enforcement and the
1828                            // inline backstop stop chasing memory the kernel
1829                            // will not give back. A read that restores the
1830                            // budget re-counts and re-enqueues it.
1831                            self.counters
1832                                .extent_unreclaimable_bytes
1833                                .fetch_add(u64::cast_from(extent.alloc_size()), Ordering::Relaxed);
1834                        } else {
1835                            // Budget remains: keep the queue slot so later
1836                            // passes retry it up to the cap.
1837                            self.extent_queue().push_back(weak);
1838                        }
1839                    }
1840                }
1841                // Paged out already or dropped: the entry is stale. A later
1842                // resident event re-enqueues.
1843                _ => {}
1844            }
1845        }
1846    }
1847
1848    /// If the chunk is a live `UnbackedResident` holding a slot and the
1849    /// spill threads have capacity, transitions it to `WriteInFlight` and
1850    /// hands it to them, returning `true`. The hand-off happens under the
1851    /// held state lock; the spill thread blocks on that lock only after this
1852    /// call returns and the caller releases it.
1853    fn spill_handoff(&self, meta: &Arc<ChunkMeta>, state: &mut ChunkState) -> bool {
1854        // The slot check excludes empty chunks, which are `UnbackedResident`
1855        // without a slot: handing one off would panic the spill thread on
1856        // the missing slot.
1857        if state.residency != Residency::UnbackedResident
1858            || state.freed
1859            || state.slot.is_none()
1860            || !self.spill_eligible()
1861        {
1862            return false;
1863        }
1864        state.residency = Residency::WriteInFlight;
1865        self.spill_schedule(Arc::clone(meta));
1866        true
1867    }
1868}
1869
1870impl ChunkHandle {
1871    /// Test hook: the chunk's current residency state.
1872    #[cfg(test)]
1873    fn residency(&self) -> Residency {
1874        self.meta.state().residency
1875    }
1876
1877    /// Copies the whole contents into `dst` (cleared first), leaving the
1878    /// chunk's residency untouched: a resident slot is copied out directly,
1879    /// and an evicted extent decompresses straight into `dst` without
1880    /// allocating a slot. A read therefore never raises resident bytes,
1881    /// never converts the chunk's state, and hands out no reference into
1882    /// pool memory.
1883    ///
1884    /// The copy runs under the chunk's state lock, which is what makes the
1885    /// no-reference contract cheap: eviction takes the same lock, so there
1886    /// is no reader it could race. The admitting variant is
1887    /// [`ChunkHandle::read_into_admit`].
1888    pub fn read_into(&self, dst: &mut Vec<u64>) {
1889        self.read_impl(0..self.meta.len, dst, false);
1890    }
1891
1892    /// As [`ChunkHandle::read_into`], restricted to the word range `range`
1893    /// of the chunk's contents, which must lie within them. `dst` receives
1894    /// exactly the range.
1895    ///
1896    /// The range narrows only the copy into `dst`: the swap backend's
1897    /// stored form is a whole compressed block, so a cold read still
1898    /// faults and decompresses the entire extent, and accounting is that
1899    /// of a whole-chunk read.
1900    pub fn read_range_into(&self, range: Range<usize>, dst: &mut Vec<u64>) {
1901        self.read_impl(range, dst, false);
1902    }
1903
1904    /// As [`ChunkHandle::read_into`], except that an evicted chunk is
1905    /// re-admitted to `BackedResident` (its extent kept, its touched bit
1906    /// set) when a slot is available from free budget headroom or by
1907    /// stealing from a clean backed victim of the same size class, never by
1908    /// evicting or compressing anything. When neither source yields a slot
1909    /// the read is served as a plain decompress and the chunk stays
1910    /// evicted.
1911    ///
1912    /// For demand reads on probe paths, where the same chunk is likely to
1913    /// be read again. Merge, drain, and other consume-once paths should use
1914    /// [`ChunkHandle::read_into`] or [`ChunkHandle::take`]: admitting there
1915    /// churns the clean-victim stock that eager backing exists to build,
1916    /// evicting probe targets to house data about to die.
1917    pub fn read_into_admit(&self, dst: &mut Vec<u64>) {
1918        self.read_impl(0..self.meta.len, dst, true);
1919    }
1920
1921    /// As [`ChunkHandle::read_into_admit`], restricted to the word range
1922    /// `range` per [`ChunkHandle::read_range_into`]. Admission is
1923    /// whole-chunk regardless of the range: the acquired slot holds the
1924    /// entire body.
1925    pub fn read_range_into_admit(&self, range: Range<usize>, dst: &mut Vec<u64>) {
1926        self.read_impl(range, dst, true);
1927    }
1928
1929    /// Shared body of the copy-out reads: fills `dst` with the word range
1930    /// `range` of the chunk's contents under the chunk's state lock,
1931    /// re-admitting an evicted chunk when `admit` is set and a slot is
1932    /// available. An empty range returns without locking or touching the
1933    /// chunk, like the whole-chunk read of an empty chunk always has.
1934    fn read_impl(&self, range: Range<usize>, dst: &mut Vec<u64>, admit: bool) {
1935        dst.clear();
1936        let meta = &*self.meta;
1937        assert!(
1938            range.start <= range.end && range.end <= meta.len,
1939            "range {range:?} exceeds the chunk's {} words",
1940            meta.len,
1941        );
1942        if range.is_empty() {
1943            return;
1944        }
1945        let mut state = meta.state();
1946        state.touched = true;
1947        let mut extent_revived = false;
1948        match state.residency {
1949            Residency::Oversize => {
1950                let payload = state.oversize.as_ref().expect("oversize chunk has payload");
1951                dst.extend_from_slice(&payload[range]);
1952            }
1953            Residency::Evicted => {
1954                let slot = if admit {
1955                    meta.pool.admit_slot(meta)
1956                } else {
1957                    None
1958                };
1959                let extent = state.extent.as_mut().expect("evicted chunk has an extent");
1960                // Reading faults the extent's pages back in either way, so
1961                // it is re-counted against the compressed tier below.
1962                let was_resident = extent.is_resident();
1963                let was_capped = extent.pageout_capped();
1964                let extent_alloc = extent.alloc_size();
1965                match slot {
1966                    Some(slot) => {
1967                        // Admission: the extent decompresses straight into
1968                        // the acquired slot, fully overwriting its
1969                        // unspecified prior contents, and the caller's
1970                        // buffer is filled from the slot.
1971                        let region = meta.pool.region_of(meta);
1972                        // SAFETY: the slot was acquired for this chunk
1973                        // under its held state lock (freshly allocated, or
1974                        // transferred from the victim under the victim's
1975                        // lock), so it is exclusively owned with no other
1976                        // reference into it, and `len_bytes` fits the
1977                        // class.
1978                        let slot_bytes = unsafe {
1979                            std::slice::from_raw_parts_mut(region.slot_ptr(slot), meta.len_bytes())
1980                        };
1981                        extent.read_into(meta.codec, slot_bytes);
1982                        state.slot = Some(slot);
1983                        state.residency = Residency::BackedResident;
1984                        // SAFETY: the slot belongs to this chunk while the
1985                        // state lock is held (eviction and free both take
1986                        // it).
1987                        let src = unsafe { meta.pool.slot_data(meta, slot) };
1988                        dst.extend_from_slice(&src[range.start..range.end]);
1989                        // Resident again: rejoin the eviction candidates.
1990                        // A leftover entry from before the chunk's eviction
1991                        // stays sound (each entry is validated against the
1992                        // chunk's state on visit) but costs policy: two live
1993                        // entries give the enforcer two chances to spend
1994                        // this chunk's single touched bit, halving its
1995                        // second chance until one entry drains.
1996                        meta.pool
1997                            .queue(band(meta.depth))
1998                            .push_back(Arc::downgrade(&self.meta));
1999                    }
2000                    None => {
2001                        // The zero-fill ahead of the decompress is deliberate
2002                        // waste (~a tenth of the decompress cost): the extent
2003                        // read takes an initialized `&mut [u8]`, so skipping
2004                        // the fill would mean exposing uninitialized memory
2005                        // through a safe reference.
2006                        dst.resize(range.end - range.start, 0);
2007                        let bytes: &mut [u8] = bytemuck::cast_slice_mut(dst.as_mut_slice());
2008                        extent.read_range_into(
2009                            meta.codec,
2010                            meta.len_bytes(),
2011                            range.start * 8,
2012                            bytes,
2013                        );
2014                    }
2015                }
2016                // TODO: a sub-range read of a rangeable stored form (file
2017                // extents, a sub-block-framed codec) revives only part of
2018                // the extent; the whole-extent accounting below would then
2019                // overcount and needs a partial-revival variant.
2020                if !was_resident {
2021                    // Revived from the device: the decompress reset any
2022                    // retry budget, so the extent re-enters reclaimable.
2023                    meta.pool
2024                        .note_extent_resident(&self.meta, extent_alloc, true);
2025                    extent_revived = true;
2026                } else if was_capped {
2027                    // The decompress faulted every page and reset the
2028                    // pageout retry budget, so a retry-capped extent is
2029                    // reclaimable again. Heap-backed extents stay
2030                    // structurally capped and stay out of the queue.
2031                    let capped = state
2032                        .extent
2033                        .as_ref()
2034                        .expect("evicted chunk has an extent")
2035                        .pageout_capped();
2036                    if !capped {
2037                        meta.pool.note_extent_reclaimable(&self.meta, extent_alloc);
2038                        extent_revived = true;
2039                    }
2040                }
2041            }
2042            Residency::UnbackedResident | Residency::BackedResident | Residency::WriteInFlight => {
2043                let slot = state.slot.expect("resident non-empty chunk has a slot");
2044                // SAFETY: the slot belongs to this chunk while the state lock
2045                // is held (eviction and free both take it).
2046                let src = unsafe { meta.pool.slot_data(meta, slot) };
2047                dst.extend_from_slice(&src[range.start..range.end]);
2048            }
2049        }
2050        drop(state);
2051        // The read revived the extent's compressed pages; the tier may need
2052        // trimming. Enforcement locks chunk states itself, so it must run
2053        // after the unlock.
2054        if extent_revived {
2055            meta.pool.enforce_or_defer_compressed_cap();
2056        }
2057    }
2058
2059    /// Copies the whole contents into `dst` (per [`ChunkHandle::read_into`],
2060    /// never admitting) and frees the chunk, cancelling any in-flight
2061    /// backing write.
2062    pub fn take(self, dst: &mut Vec<u64>) {
2063        self.read_into(dst);
2064    }
2065
2066    /// Advisory a consumer may issue before a bulk read: hints the kernel to
2067    /// swap an evicted chunk's extent back in, and is a no-op in every other
2068    /// state. Never blocks on I/O (`MADV_WILLNEED` is asynchronous).
2069    pub fn prefetch(&self) {
2070        let state = self.meta.state();
2071        if state.residency == Residency::Evicted {
2072            let extent = state.extent.as_ref().expect("evicted chunk has an extent");
2073            extent.prefetch();
2074        }
2075    }
2076
2077    /// As [`ChunkHandle::prefetch`], scoped to the word range `range` of the
2078    /// chunk's contents. The range is advisory: a backend hints at whatever
2079    /// granularity its stored form permits, and the swap backend's stored
2080    /// form is a whole compressed block, so it hints the entire extent.
2081    pub fn prefetch_range(&self, range: Range<usize>) {
2082        let _ = range;
2083        self.prefetch();
2084    }
2085
2086    /// Test hook: the byte size of the chunk's size class, or `None` for
2087    /// empty and oversize chunks.
2088    #[cfg(test)]
2089    fn size_class_bytes(&self) -> Option<usize> {
2090        self.meta.class.map(|class| SIZE_CLASSES[class])
2091    }
2092}
2093
2094impl Drop for ChunkHandle {
2095    fn drop(&mut self) {
2096        let pool = &self.meta.pool;
2097        let mut state = self.meta.state();
2098        pool.counters.frees.fetch_add(1, Ordering::Relaxed);
2099        state.freed = true;
2100        if self.meta.class.is_some() {
2101            pool.live_chunks.fetch_sub(1, Ordering::Relaxed);
2102        }
2103        let len_bytes = u64::cast_from(self.meta.len_bytes());
2104        // `release_slot`'s precondition holds in every arm below: the handle
2105        // is being dropped, so no copy-out read (which borrows the handle)
2106        // is in progress, and `freed` was set under the state lock held
2107        // here, so concurrent queue visitors skip the chunk.
2108        match state.residency {
2109            Residency::UnbackedResident => {
2110                if state.slot.is_some() {
2111                    pool.counters.writes_elided.fetch_add(1, Ordering::Relaxed);
2112                    pool.release_slot(&self.meta, &mut state);
2113                }
2114            }
2115            Residency::BackedResident => {
2116                pool.release_slot(&self.meta, &mut state);
2117                if let Some(extent) = &state.extent {
2118                    pool.note_extent_released(extent);
2119                }
2120                state.extent = None;
2121            }
2122            Residency::Evicted => {
2123                crate::soft_assert_no_log!(state.slot.is_none(), "evicted chunk holds no slot");
2124                if let Some(extent) = &state.extent {
2125                    pool.note_extent_released(extent);
2126                }
2127                state.extent = None;
2128            }
2129            Residency::WriteInFlight => {
2130                // A spill thread may be reading the slot to compress it.
2131                // `freed` (set above) tells it the chunk died; it owns the
2132                // slot release, the `resident_bytes` decrement, and the
2133                // cancellation accounting from here.
2134            }
2135            Residency::Oversize => {
2136                pool.counters
2137                    .resident_bytes
2138                    .fetch_sub(len_bytes, Ordering::Relaxed);
2139                pool.counters
2140                    .oversize_bytes
2141                    .fetch_sub(len_bytes, Ordering::Relaxed);
2142                state.oversize = None;
2143            }
2144        }
2145    }
2146}
2147
2148#[cfg(test)]
2149mod tests {
2150    use super::*;
2151    use crate::pool::extent::TEST_CODEC;
2152
2153    /// Keep test pools small: 64 MiB of virtual reservation per class.
2154    /// Under Miri the backing is real interpreter heap rather than lazy
2155    /// virtual memory, so shrink further. Classes above the capacity yield
2156    /// empty regions whose inserts degrade to the heap fallback, which is
2157    /// fine: slotted-chunk tests exercise only the smallest classes.
2158    fn test_pool(budget_bytes: usize) -> Pool {
2159        let capacity = if cfg!(miri) { 1 << 20 } else { 64 << 20 };
2160        let pool = Pool::with_class_capacity(capacity).expect("pool creation");
2161        pool.set_budget(budget_bytes);
2162        pool
2163    }
2164
2165    /// Scales an iteration count down under Miri, where one interpreted
2166    /// compression costs what thousands do natively.
2167    fn rounds(native: u64, miri: u64) -> u64 {
2168        if cfg!(miri) { miri } else { native }
2169    }
2170
2171    fn payload(words: usize, seed: u64) -> Vec<u64> {
2172        (0..u64::cast_from(words))
2173            .map(|i| seed.wrapping_mul(0x9E3779B97F4A7C15).wrapping_add(i))
2174            .collect()
2175    }
2176
2177    /// Copies `data` into the pool and clears it.
2178    fn insert(pool: &Pool, data: &mut Vec<u64>) -> ChunkHandle {
2179        insert_at_depth(pool, 0, data)
2180    }
2181
2182    /// Copies `data` into the pool at a hinted depth and clears it.
2183    fn insert_at_depth(pool: &Pool, depth: u8, data: &mut Vec<u64>) -> ChunkHandle {
2184        let hints = ChunkHints { depth };
2185        let handle = pool.insert_with(data.len(), hints, &TEST_CODEC, |dst| {
2186            dst.copy_from_slice(data.as_slice())
2187        });
2188        data.clear();
2189        handle
2190    }
2191
2192    /// Copies a chunk's contents out into a fresh buffer.
2193    fn read(handle: &ChunkHandle) -> Vec<u64> {
2194        let mut out = Vec::new();
2195        handle.read_into(&mut out);
2196        out
2197    }
2198
2199    /// Copies a chunk's contents out into a fresh buffer via the admitting
2200    /// read.
2201    fn read_admit(handle: &ChunkHandle) -> Vec<u64> {
2202        let mut out = Vec::new();
2203        handle.read_into_admit(&mut out);
2204        out
2205    }
2206
2207    /// Words that fill a 64 KiB class exactly.
2208    const SMALL: usize = (64 << 10) / 8;
2209
2210    #[allow(dead_code)]
2211    fn assert_handle_send_sync() {
2212        fn check<T: Send + Sync>() {}
2213        check::<Pool>();
2214        check::<ChunkHandle>();
2215    }
2216
2217    /// With an RSS target set, evicted chunks keep their extents resident
2218    /// (the compressed tier); shrinking the target pages the oldest extents
2219    /// out; reads revive them and re-count them.
2220    #[mz_ore::test]
2221    fn compressed_tier_round_trip() {
2222        let pool = test_pool(256 << 20);
2223        pool.set_rss_target(1 << 30);
2224        let orig = payload(SMALL, 21);
2225        let handle = insert(&pool, &mut orig.clone());
2226        pool.evict(&handle);
2227        assert_eq!(handle.residency(), Residency::Evicted);
2228        let stats = pool.stats();
2229        assert!(
2230            stats.extent_resident_bytes > 0,
2231            "under the target, the extent stays resident",
2232        );
2233        assert_eq!(stats.extent_pageouts, 0);
2234
2235        // Shrinking the target to zero pages the extent out.
2236        pool.set_rss_target(0);
2237        let stats = pool.stats();
2238        assert_eq!(stats.extent_resident_bytes, 0, "tier collapsed");
2239        assert_eq!(stats.extent_pageouts, 1);
2240
2241        // Reading revives the extent: contents round-trip, the chunk stays
2242        // evicted, and with the target restored the revived extent is
2243        // counted again.
2244        pool.set_rss_target(1 << 30);
2245        assert_eq!(read(&handle), orig);
2246        assert_eq!(handle.residency(), Residency::Evicted);
2247        assert!(
2248            pool.stats().extent_resident_bytes > 0,
2249            "revived and counted"
2250        );
2251
2252        // Dropping the handle uncounts the resident extent.
2253        drop(handle);
2254        assert_eq!(pool.stats().extent_resident_bytes, 0);
2255    }
2256
2257    /// A ranged read returns exactly the corresponding slice of a
2258    /// whole-chunk read in every residency state, and changes residency
2259    /// exactly as the equivalent whole-chunk read would.
2260    #[mz_ore::test]
2261    fn ranged_reads_match_full_read_slice() {
2262        let pool = test_pool(256 << 20);
2263        pool.set_rss_target(1 << 30);
2264        let orig = payload(SMALL, 33);
2265        let handle = insert(&pool, &mut orig.clone());
2266        let ranges = [
2267            (0usize, 7usize),
2268            (13, 100),
2269            (SMALL - 9, 9),
2270            (0, SMALL),
2271            (5, 0),
2272        ];
2273        let check = |label: &str| {
2274            for (start, len) in ranges {
2275                let mut out = Vec::new();
2276                handle.read_range_into(start..start + len, &mut out);
2277                assert_eq!(
2278                    out,
2279                    &orig[start..start + len],
2280                    "{label} range ({start}, {len})"
2281                );
2282            }
2283        };
2284        assert_eq!(handle.residency(), Residency::UnbackedResident);
2285        check("resident");
2286        pool.evict(&handle);
2287        assert_eq!(handle.residency(), Residency::Evicted);
2288        check("evicted");
2289        assert_eq!(
2290            handle.residency(),
2291            Residency::Evicted,
2292            "plain ranged reads do not admit"
2293        );
2294        // An admitting ranged read returns the range and admits the whole
2295        // chunk.
2296        let mut out = Vec::new();
2297        handle.read_range_into_admit(3..19, &mut out);
2298        assert_eq!(out, &orig[3..19]);
2299        assert_eq!(handle.residency(), Residency::BackedResident);
2300        check("backed");
2301    }
2302
2303    #[mz_ore::test]
2304    #[should_panic(expected = "exceeds the chunk's")]
2305    fn ranged_read_out_of_bounds_panics() {
2306        let pool = test_pool(256 << 20);
2307        let handle = insert(&pool, &mut payload(SMALL, 34));
2308        let mut out = Vec::new();
2309        handle.read_range_into(SMALL - 1..SMALL + 1, &mut out);
2310    }
2311
2312    #[mz_ore::test]
2313    fn default_target_pages_extents_immediately() {
2314        let pool = test_pool(256 << 20);
2315        let handle = insert(&pool, &mut payload(SMALL, 22));
2316        pool.evict(&handle);
2317        let stats = pool.stats();
2318        assert_eq!(stats.extent_resident_bytes, 0);
2319        assert_eq!(stats.extent_pageouts, 1);
2320    }
2321
2322    #[mz_ore::test]
2323    fn full_pageout_uncounts_exactly_the_extent() {
2324        let pool = test_pool(256 << 20);
2325        pool.set_rss_target(1 << 30);
2326        let handle = insert(&pool, &mut payload(SMALL, 50));
2327        pool.evict(&handle);
2328        let counted = pool.stats().extent_resident_bytes;
2329        assert!(counted > 0, "under the target, the extent stays counted");
2330        pool.set_rss_target(0);
2331        let stats = pool.stats();
2332        assert_eq!(stats.extent_resident_bytes, 0, "exactly `counted` left");
2333        assert_eq!(stats.extent_pageouts, 1);
2334        assert_eq!(stats.extent_pageout_incomplete, 0);
2335    }
2336
2337    #[mz_ore::test]
2338    fn incomplete_pageout_keeps_accounting_and_queue_position() {
2339        let pool = test_pool(256 << 20);
2340        pool.set_rss_target(1 << 30);
2341        let handle = insert(&pool, &mut payload(SMALL, 51));
2342        pool.evict(&handle);
2343        let counted = pool.stats().extent_resident_bytes;
2344        assert!(counted > 0);
2345        region::fake_residency::decline_next(1);
2346        pool.set_rss_target(0);
2347        let stats = pool.stats();
2348        assert_eq!(
2349            stats.extent_resident_bytes, counted,
2350            "full accounting stays"
2351        );
2352        assert_eq!(stats.extent_pageouts, 0);
2353        assert_eq!(stats.extent_pageout_incomplete, 1);
2354        assert_eq!(handle.residency(), Residency::Evicted);
2355        // The requeued entry is retried by the next enforcement pass.
2356        pool.enforce_rss_target();
2357        let stats = pool.stats();
2358        assert_eq!(stats.extent_resident_bytes, 0);
2359        assert_eq!(stats.extent_pageouts, 1);
2360        assert_eq!(stats.extent_pageout_incomplete, 1);
2361    }
2362
2363    /// A never-reclaimable extent stops being advised after the retry cap:
2364    /// the incomplete counter stops climbing, the bytes stay counted
2365    /// resident, and the tier keeps paging other extents out around it.
2366    #[mz_ore::test]
2367    fn pageout_retry_cap_stops_advising() {
2368        let pool = test_pool(256 << 20);
2369        let handle = insert(&pool, &mut payload(SMALL, 52));
2370        region::fake_residency::decline_next(u64::MAX);
2371        // RSS target zero: the eviction's enforcement pass advises at once.
2372        pool.evict(&handle);
2373        for _ in 0..5 {
2374            pool.enforce_rss_target();
2375        }
2376        let stats = pool.stats();
2377        assert_eq!(
2378            stats.extent_pageout_incomplete,
2379            u64::from(extent::PAGEOUT_RETRY_CAP),
2380            "advised exactly retry-cap times",
2381        );
2382        assert_eq!(stats.extent_pageouts, 0);
2383        let counted = stats.extent_resident_bytes;
2384        assert!(counted > 0, "capped extent stays counted resident");
2385        // The tier functions around the capped extent: a fresh extent still
2386        // pages out.
2387        region::fake_residency::decline_next(0);
2388        let other = insert(&pool, &mut payload(SMALL, 53));
2389        pool.evict(&other);
2390        let stats = pool.stats();
2391        assert_eq!(stats.extent_pageouts, 1);
2392        assert_eq!(
2393            stats.extent_resident_bytes, counted,
2394            "only the capped extent remains counted",
2395        );
2396        assert_eq!(read(&handle).len(), SMALL, "capped extent stays readable");
2397    }
2398
2399    #[mz_ore::test]
2400    fn read_resets_pageout_retry_budget() {
2401        let pool = test_pool(256 << 20);
2402        let orig = payload(SMALL, 54);
2403        let handle = insert(&pool, &mut orig.clone());
2404        region::fake_residency::decline_next(u64::MAX);
2405        pool.evict(&handle);
2406        for _ in 0..4 {
2407            pool.enforce_rss_target();
2408        }
2409        assert_eq!(
2410            pool.stats().extent_pageout_incomplete,
2411            u64::from(extent::PAGEOUT_RETRY_CAP),
2412            "capped",
2413        );
2414        assert!(pool.stats().extent_resident_bytes > 0);
2415        region::fake_residency::decline_next(0);
2416        assert_eq!(read(&handle), orig);
2417        pool.enforce_rss_target();
2418        let stats = pool.stats();
2419        assert_eq!(stats.extent_pageouts, 1, "the budget reset re-advised it");
2420        assert_eq!(stats.extent_resident_bytes, 0);
2421        // The paged-out extent still round-trips.
2422        assert_eq!(read(&handle), orig);
2423    }
2424
2425    /// Eager backing compresses a chunk to `BackedResident` while it stays
2426    /// readable in its slot; the later budget-driven eviction is a pure page
2427    /// release, and the contents round-trip through the extent.
2428    #[mz_ore::test]
2429    fn eager_backing_round_trip() {
2430        let pool = test_pool(256 << 20);
2431        let orig = payload(SMALL, 11);
2432        let handle = insert(&pool, &mut orig.clone());
2433        assert_eq!(handle.residency(), Residency::UnbackedResident);
2434
2435        assert!(pool.back_step(), "one chunk is backable");
2436        assert_eq!(handle.residency(), Residency::BackedResident);
2437        let stats = pool.stats();
2438        assert_eq!(stats.eager_backs, 1);
2439        assert_eq!(stats.evictions_compress, 0, "backing is not an eviction");
2440        assert!(stats.extent_bytes_written > 0);
2441
2442        // Still readable straight from the slot: the chunk is resident.
2443        assert_eq!(read(&handle), orig);
2444
2445        // The pre-paid eviction is cheap, and the extent round-trips.
2446        pool.evict(&handle);
2447        assert_eq!(handle.residency(), Residency::Evicted);
2448        assert_eq!(pool.stats().evictions_cheap, 1);
2449        pool.poison_free_slots();
2450        assert_eq!(read(&handle), orig);
2451    }
2452
2453    #[mz_ore::test]
2454    fn backing_reports_no_progress_when_all_backed() {
2455        let pool = test_pool(256 << 20);
2456        let _handle = insert(&pool, &mut payload(SMALL, 31));
2457        assert!(pool.back_step(), "one unbacked chunk is actionable");
2458        assert!(!pool.back_step(), "fully backed: no progress");
2459        assert_eq!(pool.stats().eager_backs, 1);
2460    }
2461
2462    /// Freeing under the warm cap parks the slot warm; the next insert of the
2463    /// same class reuses it fault-free and the accounting balances.
2464    #[mz_ore::test]
2465    fn warm_slot_reuse() {
2466        // Budget 8 MiB: warm cap = 1 MiB, so a 64 KiB slot fits warm.
2467        let pool = test_pool(8 << 20);
2468        let orig = payload(SMALL, 7);
2469        let handle = insert(&pool, &mut orig.clone());
2470        drop(handle);
2471        let after_free = pool.stats();
2472        assert_eq!(after_free.warm_bytes, 64 << 10, "freed slot parks warm");
2473        assert_eq!(after_free.warm_reuses, 0);
2474
2475        let handle = insert(&pool, &mut orig.clone());
2476        let after_reuse = pool.stats();
2477        assert_eq!(after_reuse.warm_reuses, 1, "second insert reuses warm slot");
2478        assert_eq!(after_reuse.warm_bytes, 0, "reuse drains the warm pool");
2479        // Contents are correct despite the skipped page release.
2480        assert_eq!(read(&handle), orig);
2481    }
2482
2483    #[mz_ore::test]
2484    fn warm_pool_respects_cap() {
2485        // Budget 1 MiB: warm cap = 128 KiB = two 64 KiB slots.
2486        let pool = test_pool(1 << 20);
2487        let handles: Vec<_> = (0..4)
2488            .map(|seed| insert(&pool, &mut payload(SMALL, seed)))
2489            .collect();
2490        drop(handles);
2491        let stats = pool.stats();
2492        assert_eq!(
2493            stats.warm_bytes,
2494            128 << 10,
2495            "warm pool stops at the budget/8 cap",
2496        );
2497    }
2498
2499    /// A kernel that keeps declining the reclaim advice caps the extent's
2500    /// retry budget: the extent leaves the enforcement queue and moves to
2501    /// the unreclaimable gauge, so enforcement stops walking it, and a read
2502    /// that restores the budget makes it reclaimable and pageable again.
2503    #[mz_ore::test]
2504    fn capped_extents_leave_the_enforcement_queue() {
2505        let pool = test_pool(256 << 20);
2506        let orig = payload(SMALL, 960);
2507        let handle = insert(&pool, &mut orig.clone());
2508        // Decline every observation: eviction's own enforcement pass plus
2509        // the passes below spend the whole retry budget.
2510        region::fake_residency::decline_next(u64::from(extent::PAGEOUT_RETRY_CAP));
2511        pool.evict(&handle);
2512        for _ in 0..extent::PAGEOUT_RETRY_CAP {
2513            pool.enforce_compressed();
2514        }
2515        let stats = pool.stats();
2516        assert_eq!(
2517            stats.extent_pageout_incomplete,
2518            u64::from(extent::PAGEOUT_RETRY_CAP),
2519        );
2520        assert!(stats.extent_unreclaimable_bytes > 0, "capped bytes counted");
2521        assert_eq!(pool.extent_queue_len(), 0, "capped extents leave the queue",);
2522        // Further enforcement is a no-op: nothing queued, no advice spent.
2523        pool.enforce_compressed();
2524        assert_eq!(
2525            pool.stats().extent_pageout_incomplete,
2526            u64::from(extent::PAGEOUT_RETRY_CAP),
2527        );
2528
2529        // A read faults everything back in and restores the retry budget:
2530        // the extent re-enters the reclaimable set and, with the kernel now
2531        // cooperating, the read's own enforcement pass pages it out.
2532        assert_eq!(read(&handle), orig);
2533        let stats = pool.stats();
2534        assert_eq!(stats.extent_unreclaimable_bytes, 0, "budget restored");
2535        assert_eq!(stats.extent_pageouts, 1, "re-enqueued extent pages out");
2536        assert_eq!(stats.extent_resident_bytes, 0);
2537        assert_eq!(pool.extent_queue_len(), 0);
2538    }
2539
2540    /// Shrinking the budget cools warm slots parked under the old, larger
2541    /// cap: their pages are released and `warm_bytes` falls to the new cap
2542    /// on the shrink itself, not on eventual same-class reuse.
2543    #[mz_ore::test]
2544    fn budget_shrink_trims_warm_pool() {
2545        // Budget 8 MiB: warm cap 1 MiB, so four 64 KiB frees all park warm.
2546        let pool = test_pool(8 << 20);
2547        let handles: Vec<_> = (0..4)
2548            .map(|seed| insert(&pool, &mut payload(SMALL, 950 + seed)))
2549            .collect();
2550        drop(handles);
2551        assert_eq!(pool.stats().warm_bytes, 4 * (64 << 10));
2552
2553        // Budget 1 MiB: warm cap 128 KiB, so two of the four slots cool.
2554        pool.set_budget(1 << 20);
2555        assert_eq!(pool.stats().warm_bytes, 128 << 10);
2556    }
2557
2558    #[mz_ore::test]
2559    fn round_trip_resident() {
2560        let pool = test_pool(256 << 20);
2561        let orig = payload(1000, 1);
2562        let mut data = orig.clone();
2563        let capacity = data.capacity();
2564        let handle = insert(&pool, &mut data);
2565        assert!(data.is_empty());
2566        assert_eq!(data.capacity(), capacity, "insert preserves capacity");
2567        assert_eq!(handle.residency(), Residency::UnbackedResident);
2568        assert_eq!(read(&handle), orig);
2569        drop(handle);
2570        let stats = pool.stats();
2571        assert_eq!(stats.inserts, 1);
2572        assert_eq!(stats.frees, 1);
2573        assert_eq!(stats.resident_bytes, 0);
2574    }
2575
2576    #[mz_ore::test]
2577    fn take_reads_and_frees() {
2578        let pool = test_pool(256 << 20);
2579        let orig = payload(SMALL, 40);
2580        let handle = insert(&pool, &mut orig.clone());
2581        let mut out = Vec::new();
2582        handle.take(&mut out);
2583        assert_eq!(out, orig);
2584        let stats = pool.stats();
2585        assert_eq!(stats.frees, 1);
2586        assert_eq!(stats.writes_elided, 1, "a resident take never writes");
2587        assert_eq!(stats.resident_bytes, 0);
2588        assert_eq!(stats.live_chunks, 0);
2589    }
2590
2591    /// `prefetch` is safe wherever it lands: on a resident chunk (a no-op),
2592    /// on an evicted chunk (whose read then round-trips), and issued with no
2593    /// read following it. It never changes residency or resident bytes.
2594    #[mz_ore::test]
2595    fn prefetch_is_safe_in_every_state() {
2596        let pool = test_pool(256 << 20);
2597        let orig = payload(SMALL, 41);
2598        let handle = insert(&pool, &mut orig.clone());
2599        handle.prefetch();
2600        assert_eq!(handle.residency(), Residency::UnbackedResident);
2601        assert_eq!(read(&handle), orig);
2602        pool.evict(&handle);
2603        handle.prefetch();
2604        assert_eq!(handle.residency(), Residency::Evicted);
2605        assert_eq!(read(&handle), orig);
2606        // An advisory with no read behind it leaves nothing to clean up.
2607        let idle = insert(&pool, &mut payload(SMALL, 42));
2608        idle.prefetch();
2609        drop(idle);
2610        drop(handle);
2611        assert_eq!(pool.stats().resident_bytes, 0);
2612    }
2613
2614    /// Reading an evicted chunk decompresses its extent straight into the
2615    /// caller's buffer and leaves the chunk evicted. Free slots are poisoned
2616    /// first, so a read passing stale slot memory through (the macOS
2617    /// `MADV_DONTNEED` hazard) would fail the content check.
2618    #[mz_ore::test]
2619    fn evict_then_read_preserves_contents() {
2620        let pool = test_pool(256 << 20);
2621        let orig = payload(SMALL, 2);
2622        let handle = insert(&pool, &mut orig.clone());
2623        pool.evict(&handle);
2624        assert_eq!(handle.residency(), Residency::Evicted);
2625        let stats = pool.stats();
2626        assert_eq!(stats.evictions_compress, 1);
2627        assert_eq!(stats.resident_bytes, 0);
2628        assert!(stats.extent_bytes_written > 0);
2629        pool.poison_free_slots();
2630        assert_eq!(read(&handle), orig);
2631        assert_eq!(handle.residency(), Residency::Evicted);
2632        assert_eq!(pool.stats().resident_bytes, 0, "reads copy out");
2633    }
2634
2635    /// An admitting read of an evicted chunk with budget headroom re-admits
2636    /// it: contents round-trip, the chunk lands `BackedResident` with its
2637    /// extent kept, and later reads serve from the slot without touching
2638    /// the extent.
2639    #[mz_ore::test]
2640    fn admit_from_free_budget_backs_the_chunk() {
2641        let pool = test_pool(256 << 20);
2642        let orig = payload(SMALL, 70);
2643        let handle = insert(&pool, &mut orig.clone());
2644        pool.evict(&handle);
2645        assert_eq!(handle.residency(), Residency::Evicted);
2646        assert_eq!(pool.stats().resident_bytes, 0);
2647
2648        pool.poison_free_slots();
2649        assert_eq!(read_admit(&handle), orig);
2650        assert_eq!(handle.residency(), Residency::BackedResident);
2651        let stats = pool.stats();
2652        assert_eq!(stats.admissions_budget, 1);
2653        assert_eq!(stats.admissions_steal, 0);
2654        assert_eq!(stats.admissions_denied, 0);
2655        assert_eq!(stats.resident_bytes, 64 << 10);
2656
2657        // Later reads serve from the slot and never touch the extent: a
2658        // decompress would revive its pages and move the revival and
2659        // pageout counters.
2660        let pageouts = stats.extent_pageouts;
2661        let extent_resident = stats.extent_resident_bytes;
2662        assert_eq!(read(&handle), orig);
2663        assert_eq!(handle.residency(), Residency::BackedResident);
2664        let stats = pool.stats();
2665        assert_eq!(stats.extent_pageouts, pageouts);
2666        assert_eq!(stats.extent_resident_bytes, extent_resident);
2667
2668        // The kept extent pre-pays the next eviction, and round-trips.
2669        pool.evict(&handle);
2670        assert_eq!(handle.residency(), Residency::Evicted);
2671        let stats = pool.stats();
2672        assert_eq!(stats.evictions_cheap, 1);
2673        assert_eq!(stats.evictions_compress, 1, "admission wrote no extent");
2674        pool.poison_free_slots();
2675        assert_eq!(read(&handle), orig);
2676        drop(handle);
2677        assert_eq!(pool.stats().resident_bytes, 0);
2678    }
2679
2680    /// With the budget pinned full and a clean backed victim of the same
2681    /// class, an admitting read steals the victim's slot: the victim is
2682    /// evicted with zero I/O and its extent intact, the admitted chunk
2683    /// lands `BackedResident`, and resident bytes, warm bytes, and the
2684    /// compression and pageout counters are all unchanged.
2685    #[mz_ore::test]
2686    fn admit_steals_clean_victim_slot() {
2687        let pool = test_pool(256 << 20);
2688        pool.set_rss_target(1 << 30);
2689        let victim_orig = payload(SMALL, 71);
2690        let target_orig = payload(SMALL, 72);
2691        let victim = insert(&pool, &mut victim_orig.clone());
2692        let target = insert(&pool, &mut target_orig.clone());
2693        pool.evict(&target);
2694        assert!(pool.back_step(), "victim is backable");
2695        assert_eq!(victim.residency(), Residency::BackedResident);
2696        // The budget now holds exactly the victim: no admission headroom.
2697        pool.set_budget(64 << 10);
2698        assert_eq!(victim.residency(), Residency::BackedResident);
2699        let before = pool.stats();
2700
2701        assert_eq!(read_admit(&target), target_orig);
2702        assert_eq!(target.residency(), Residency::BackedResident);
2703        assert_eq!(victim.residency(), Residency::Evicted);
2704        let after = pool.stats();
2705        assert_eq!(after.admissions_steal, 1);
2706        assert_eq!(after.admissions_budget, 0);
2707        assert_eq!(after.admissions_denied, 0);
2708        assert_eq!(
2709            after.resident_bytes, before.resident_bytes,
2710            "same class, same bytes",
2711        );
2712        assert_eq!(
2713            after.evictions_compress, before.evictions_compress,
2714            "no compression",
2715        );
2716        assert_eq!(
2717            after.evictions_cheap, before.evictions_cheap,
2718            "a steal is not an enforcement eviction",
2719        );
2720        assert_eq!(after.extent_bytes_written, before.extent_bytes_written);
2721        assert_eq!(after.extent_pageouts, 0, "no pageout");
2722        assert_eq!(
2723            after.warm_bytes, before.warm_bytes,
2724            "the stolen slot skipped the free list",
2725        );
2726        assert_eq!(after.warm_reuses, before.warm_reuses);
2727
2728        // The victim's extent is intact: its old slot now holds the
2729        // admitted chunk's bytes, so a correct read must come from the
2730        // extent.
2731        assert_eq!(read(&victim), victim_orig);
2732        assert_eq!(victim.residency(), Residency::Evicted);
2733
2734        drop(victim);
2735        drop(target);
2736        let stats = pool.stats();
2737        assert_eq!(stats.resident_bytes, 0);
2738        assert_eq!(stats.extent_resident_bytes, 0);
2739    }
2740
2741    /// With the budget full and every candidate touched, the admitting read
2742    /// still returns correct data, the chunk stays evicted, and the denial
2743    /// counter increments.
2744    #[mz_ore::test]
2745    fn admit_denied_when_victims_touched() {
2746        let pool = test_pool(256 << 20);
2747        let victim_orig = payload(SMALL, 73);
2748        let target_orig = payload(SMALL, 74);
2749        let victim = insert(&pool, &mut victim_orig.clone());
2750        let target = insert(&pool, &mut target_orig.clone());
2751        pool.evict(&target);
2752        assert!(pool.back_step());
2753        // Reading the victim sets its second-chance bit, disqualifying it.
2754        assert_eq!(read(&victim), victim_orig);
2755        pool.set_budget(64 << 10);
2756        let resident = pool.stats().resident_bytes;
2757
2758        assert_eq!(read_admit(&target), target_orig);
2759        assert_eq!(target.residency(), Residency::Evicted);
2760        assert_eq!(victim.residency(), Residency::BackedResident);
2761        let stats = pool.stats();
2762        assert_eq!(stats.admissions_denied, 1);
2763        assert_eq!(stats.admissions_budget, 0);
2764        assert_eq!(stats.admissions_steal, 0);
2765        assert_eq!(stats.resident_bytes, resident);
2766    }
2767
2768    /// An unbacked resident candidate is never stolen from: evicting it
2769    /// would require the compression that admission forbids.
2770    #[mz_ore::test]
2771    fn admit_denied_when_victims_unbacked() {
2772        let pool = test_pool(256 << 20);
2773        let victim = insert(&pool, &mut payload(SMALL, 75));
2774        let target_orig = payload(SMALL, 76);
2775        let target = insert(&pool, &mut target_orig.clone());
2776        pool.evict(&target);
2777        pool.set_budget(64 << 10);
2778        assert_eq!(read_admit(&target), target_orig);
2779        assert_eq!(target.residency(), Residency::Evicted);
2780        assert_eq!(victim.residency(), Residency::UnbackedResident);
2781        assert_eq!(pool.stats().admissions_denied, 1);
2782    }
2783
2784    /// A clean backed victim of a different size class is never stolen
2785    /// from: slot reuse in place requires the classes to match.
2786    #[mz_ore::test]
2787    fn admit_denied_when_victims_wrong_class() {
2788        let pool = test_pool(256 << 20);
2789        // The victim fills the 128 KiB class; the target lives in the
2790        // 64 KiB one.
2791        let victim = insert(&pool, &mut payload(2 * SMALL, 77));
2792        let target_orig = payload(SMALL, 78);
2793        let target = insert(&pool, &mut target_orig.clone());
2794        pool.evict(&target);
2795        assert!(pool.back_step());
2796        assert_eq!(victim.residency(), Residency::BackedResident);
2797        // The budget holds exactly the victim: no headroom for the target.
2798        pool.set_budget(128 << 10);
2799        assert_eq!(read_admit(&target), target_orig);
2800        assert_eq!(target.residency(), Residency::Evicted);
2801        assert_eq!(victim.residency(), Residency::BackedResident);
2802        assert_eq!(pool.stats().admissions_denied, 1);
2803    }
2804
2805    #[mz_ore::test]
2806    fn plain_read_and_take_never_admit() {
2807        let pool = test_pool(256 << 20);
2808        let orig = payload(SMALL, 79);
2809        let handle = insert(&pool, &mut orig.clone());
2810        pool.evict(&handle);
2811        assert_eq!(read(&handle), orig);
2812        assert_eq!(handle.residency(), Residency::Evicted);
2813        assert_eq!(pool.stats().resident_bytes, 0);
2814        let mut out = Vec::new();
2815        handle.take(&mut out);
2816        assert_eq!(out, orig);
2817        let stats = pool.stats();
2818        assert_eq!(stats.admissions_budget, 0);
2819        assert_eq!(stats.admissions_steal, 0);
2820        assert_eq!(stats.admissions_denied, 0);
2821        assert_eq!(stats.resident_bytes, 0);
2822        assert_eq!(stats.frees, 1);
2823    }
2824
2825    /// A steal settles the ledger with the payload difference: a shrinking
2826    /// steal always proceeds, while a steal that would grow resident bytes
2827    /// past the budget is denied and leaves the victim untouched.
2828    #[mz_ore::test]
2829    fn steal_admission_charges_the_budget() {
2830        // Shrinking steal: the victim is larger than the admitted payload,
2831        // so the steal lowers resident bytes and always may proceed.
2832        let pool = test_pool(256 << 20);
2833        let victim_orig = payload(SMALL, 84);
2834        let victim = insert(&pool, &mut victim_orig.clone());
2835        assert!(pool.back_step(), "victim backs");
2836        let small_orig = payload(SMALL / 2, 85);
2837        let handle = insert(&pool, &mut small_orig.clone());
2838        pool.evict(&handle);
2839        pool.set_budget(64 << 10);
2840        assert_eq!(read_admit(&handle), small_orig);
2841        let stats = pool.stats();
2842        assert_eq!(stats.admissions_steal, 1, "no headroom, so the read steals");
2843        assert_eq!(stats.resident_bytes, u64::cast_from(SMALL / 2 * 8));
2844        assert_eq!(victim.residency(), Residency::Evicted);
2845        assert_eq!(read(&victim), victim_orig, "victim serves from its extent");
2846
2847        // Growing steal: the admitted payload is larger than the only
2848        // victim, and the growth does not fit the budget, so the admission
2849        // is denied and the victim is left untouched.
2850        let pool = test_pool(256 << 20);
2851        let big_orig = payload(SMALL, 86);
2852        let big = insert(&pool, &mut big_orig.clone());
2853        pool.evict(&big);
2854        let small_victim = insert(&pool, &mut payload(SMALL / 2, 87));
2855        assert!(pool.back_step(), "victim backs");
2856        pool.set_budget(32 << 10);
2857        assert_eq!(read_admit(&big), big_orig);
2858        let stats = pool.stats();
2859        assert_eq!(stats.admissions_denied, 1, "growth exceeds the budget");
2860        assert_eq!(stats.admissions_steal, 0);
2861        assert_eq!(big.residency(), Residency::Evicted);
2862        assert_eq!(small_victim.residency(), Residency::BackedResident);
2863    }
2864
2865    /// Admission of a chunk whose extent was pushed to the device: the read
2866    /// revives the extent into the acquired slot and re-counts it.
2867    #[mz_ore::test]
2868    fn admission_revives_paged_out_extent() {
2869        let pool = test_pool(256 << 20);
2870        let orig = payload(SMALL, 88);
2871        let handle = insert(&pool, &mut orig.clone());
2872        pool.evict(&handle);
2873        assert_eq!(
2874            pool.stats().extent_resident_bytes,
2875            0,
2876            "zero RSS target pages the extent out on eviction",
2877        );
2878        // Raise the target so the read's own tier enforcement does not
2879        // page the revived extent straight back out.
2880        pool.set_rss_target(1 << 30);
2881        assert_eq!(read_admit(&handle), orig);
2882        assert_eq!(handle.residency(), Residency::BackedResident);
2883        let stats = pool.stats();
2884        assert_eq!(stats.admissions_budget, 1);
2885        assert!(stats.extent_resident_bytes > 0, "revived and re-counted");
2886    }
2887
2888    /// An admitting read of a chunk that is not evicted is a plain read:
2889    /// no admission counter moves and no state changes.
2890    #[mz_ore::test]
2891    fn admit_is_plain_read_on_non_evicted_chunks() {
2892        let pool = test_pool(256 << 20);
2893        let orig = payload(SMALL, 89);
2894        let resident = insert(&pool, &mut orig.clone());
2895        assert_eq!(read_admit(&resident), orig);
2896        assert_eq!(resident.residency(), Residency::UnbackedResident);
2897        let words = SIZE_CLASSES[SIZE_CLASSES.len() - 1] / 8 + 1;
2898        let oversize_orig = payload(words, 90);
2899        let oversize = insert(&pool, &mut oversize_orig.clone());
2900        assert_eq!(read_admit(&oversize), oversize_orig);
2901        let empty = insert(&pool, &mut Vec::new());
2902        assert!(read_admit(&empty).is_empty());
2903        let stats = pool.stats();
2904        assert_eq!(stats.admissions_budget, 0);
2905        assert_eq!(stats.admissions_steal, 0);
2906        assert_eq!(stats.admissions_denied, 0);
2907    }
2908
2909    /// Concurrent admitting reads with no budget headroom: every admission
2910    /// must go through the steal path, racing steals against each other on
2911    /// the same victims (the pool's only two-chunk lock edge). Contents are
2912    /// asserted on every read.
2913    #[mz_ore::test]
2914    #[cfg_attr(miri, ignore)] // too slow
2915    fn concurrent_admits_exercise_the_steal_path() {
2916        const CHUNKS: u64 = 8;
2917        let pool = test_pool(usize::MAX);
2918        let origs: Vec<_> = (0..CHUNKS).map(|seed| payload(SMALL, 900 + seed)).collect();
2919        let evicted: Arc<Vec<(Vec<u64>, ChunkHandle)>> = Arc::new(
2920            origs
2921                .iter()
2922                .map(|orig| {
2923                    let handle = insert(&pool, &mut orig.clone());
2924                    pool.evict(&handle);
2925                    (orig.clone(), handle)
2926                })
2927                .collect(),
2928        );
2929        let mut victims = Vec::new();
2930        for seed in 0..CHUNKS {
2931            victims.push(insert(&pool, &mut payload(SMALL, 950 + seed)));
2932            assert!(pool.back_step(), "victim backs");
2933        }
2934        // Exactly the victims' bytes: no free headroom, so every admission
2935        // steals or is denied.
2936        pool.set_budget(usize::cast_from(CHUNKS) * (64 << 10));
2937        let threads: Vec<_> = (0..2u64)
2938            .map(|t| {
2939                let evicted = Arc::clone(&evicted);
2940                std::thread::spawn(move || {
2941                    for round in 0..CHUNKS {
2942                        let (orig, handle) = &evicted[usize::cast_from((t + round) % CHUNKS)];
2943                        let mut out = Vec::new();
2944                        handle.read_into_admit(&mut out);
2945                        assert_eq!(&out, orig);
2946                    }
2947                })
2948            })
2949            .collect();
2950        for thread in threads {
2951            thread.join().expect("admitting thread panicked");
2952        }
2953        let stats = pool.stats();
2954        assert!(stats.admissions_steal > 0, "no headroom forces steals");
2955        assert!(
2956            stats.resident_bytes <= u64::cast_from(usize::cast_from(CHUNKS) * (64 << 10)),
2957            "steals never grow resident bytes past the budget",
2958        );
2959        // The victims were held live as steal targets; a steal leaves its
2960        // victim evicted, so at least one is evicted here.
2961        let stolen = victims
2962            .iter()
2963            .filter(|v| v.residency() == Residency::Evicted)
2964            .count();
2965        assert!(stolen > 0, "a steal evicts its victim");
2966    }
2967
2968    /// A re-admitted chunk keeps its insert-time depth: under budget
2969    /// pressure it is evicted from its own deeper band before a younger
2970    /// band-0 chunk, which a re-admission into band 0 would have inverted.
2971    #[mz_ore::test]
2972    fn admitted_chunk_keeps_its_depth() {
2973        let pool = test_pool(256 << 20);
2974        let deep_orig = payload(SMALL, 80);
2975        let deep = insert_at_depth(&pool, 2, &mut deep_orig.clone());
2976        let young = insert(&pool, &mut payload(SMALL, 81));
2977        pool.evict(&deep);
2978        assert_eq!(read_admit(&deep), deep_orig);
2979        assert_eq!(deep.residency(), Residency::BackedResident);
2980        assert_eq!(pool.stats().admissions_budget, 1);
2981
2982        // Budget of one chunk: enforcement visits the deep band first.
2983        pool.set_budget(64 << 10);
2984        assert_eq!(deep.residency(), Residency::Evicted);
2985        assert_eq!(young.residency(), Residency::UnbackedResident);
2986        assert_eq!(
2987            pool.stats().evictions_cheap,
2988            1,
2989            "the extent kept through admission pre-paid the eviction",
2990        );
2991    }
2992
2993    /// Admitting reads racing enforcement, opposing steals, and frees:
2994    /// contents stay correct, contended steals degrade to skips, and the
2995    /// accounting identity settles to zero.
2996    #[mz_ore::test]
2997    #[cfg_attr(miri, ignore)] // too slow
2998    fn concurrent_admits_race_cleanly() {
2999        let pool = test_pool(64 << 10);
3000        let per_thread = rounds(50, 3);
3001        let threads: Vec<_> = (0..4u64)
3002            .map(|t| {
3003                let pool = pool.clone();
3004                std::thread::spawn(move || {
3005                    let mut out = Vec::new();
3006                    for round in 0..per_thread {
3007                        let orig = payload(SMALL, t * 1000 + round);
3008                        let handle = insert(&pool, &mut orig.clone());
3009                        pool.evict(&handle);
3010                        handle.read_into_admit(&mut out);
3011                        assert_eq!(out, orig);
3012                        handle.read_into_admit(&mut out);
3013                        assert_eq!(out, orig);
3014                        assert_eq!(read(&handle), orig);
3015                    }
3016                })
3017            })
3018            .collect();
3019        for thread in threads {
3020            thread.join().expect("worker thread panicked");
3021        }
3022        let stats = pool.stats();
3023        assert_eq!(stats.inserts, 4 * per_thread);
3024        assert_eq!(stats.frees, 4 * per_thread);
3025        assert_eq!(stats.resident_bytes, 0);
3026        assert_eq!(stats.extent_resident_bytes, 0);
3027    }
3028
3029    /// Slots are scoped to residency: eviction releases the slot, so a
3030    /// capacity holding exactly one chunk can serve any number of chunks one
3031    /// at a time, and reads of evicted chunks need no slot at all.
3032    #[mz_ore::test]
3033    fn eviction_releases_the_slot() {
3034        // One 64 KiB slot per class.
3035        let pool = Pool::with_class_capacity(64 << 10).expect("pool creation");
3036        let a = insert(&pool, &mut payload(SMALL, 6));
3037        pool.evict(&a);
3038        // The class's only slot is free again: a second chunk fits without
3039        // falling back to the heap.
3040        let b = insert(&pool, &mut payload(SMALL, 7));
3041        assert_eq!(b.residency(), Residency::UnbackedResident);
3042        assert_eq!(pool.stats().slot_exhausted_fallbacks, 0);
3043        // Reading `a` decompresses straight from its extent while `b` holds
3044        // the class's only slot: copy-out allocates nothing.
3045        assert_eq!(read(&a), payload(SMALL, 6));
3046        assert_eq!(a.residency(), Residency::Evicted);
3047        assert_eq!(read(&b), payload(SMALL, 7));
3048    }
3049
3050    /// The eviction queue holds resident chunks only: an enforcement pass
3051    /// drops entries for evicted chunks, and reads never re-add them, so the
3052    /// scan each insert pays stays proportional to the resident set rather
3053    /// than every chunk ever evicted.
3054    #[mz_ore::test]
3055    fn queue_holds_resident_chunks_only() {
3056        let pool = test_pool(128 << 10);
3057        let mut handles = Vec::new();
3058        for seed in 0..8 {
3059            handles.push(insert(&pool, &mut payload(SMALL, 800 + seed)));
3060        }
3061        // Budget pressure evicted ~6 of 8; one more pass visits the evicted
3062        // entries and drops them (their first visit performed the eviction
3063        // and dropped them already, but second-chance survivors may linger).
3064        pool.enforce_budget();
3065        let resident = handles
3066            .iter()
3067            .filter(|h| h.residency() != Residency::Evicted)
3068            .count();
3069        assert!(
3070            pool.queue_len() <= resident + 1,
3071            "queue ({}) tracks the resident set ({resident}), not all 8 live chunks",
3072            pool.queue_len(),
3073        );
3074        // Reading an evicted chunk copies out of its extent and does not
3075        // re-enqueue it: the queue keeps tracking the resident set.
3076        let evicted = handles
3077            .iter()
3078            .find(|h| h.residency() == Residency::Evicted)
3079            .expect("something was evicted");
3080        let before = pool.queue_len();
3081        assert_eq!(read(evicted).len(), SMALL);
3082        assert_eq!(evicted.residency(), Residency::Evicted);
3083        assert_eq!(pool.queue_len(), before, "reads leave the queue alone");
3084    }
3085
3086    #[mz_ore::test]
3087    fn dead_data_is_never_written() {
3088        let pool = test_pool(256 << 20);
3089        let handle = insert(&pool, &mut payload(SMALL, 7));
3090        drop(handle);
3091        let stats = pool.stats();
3092        assert_eq!(stats.frees, 1);
3093        assert_eq!(stats.writes_elided, 1);
3094        assert_eq!(stats.extent_bytes_written, 0);
3095        assert_eq!(stats.resident_bytes, 0);
3096    }
3097
3098    #[mz_ore::test]
3099    fn budget_is_enforced_on_insert() {
3100        let budget = 128 << 10;
3101        let pool = test_pool(budget);
3102        let mut handles = Vec::new();
3103        for seed in 0..8 {
3104            handles.push(insert(&pool, &mut payload(SMALL, 100 + seed)));
3105        }
3106        let stats = pool.stats();
3107        assert!(
3108            stats.resident_bytes <= u64::cast_from(budget),
3109            "resident {} exceeds budget {}",
3110            stats.resident_bytes,
3111            budget,
3112        );
3113        assert!(stats.evictions_compress >= 6);
3114        let resident = handles
3115            .iter()
3116            .filter(|h| {
3117                matches!(
3118                    h.residency(),
3119                    Residency::UnbackedResident | Residency::BackedResident
3120                )
3121            })
3122            .count();
3123        assert_eq!(resident, 2, "budget holds exactly two small chunks");
3124    }
3125
3126    /// Budget enforcement is single-flight: an insert that trips it while a
3127    /// pass holds the `enforcing` guard bails on `WouldBlock`, trusting that
3128    /// pass. If the holder is already past its final `resident_bytes` read, the
3129    /// bailed insert's bytes are neither read by the holder nor enforced by the
3130    /// bailer, and no later insert re-trips enforcement, so the pool stays over
3131    /// budget. The fix re-runs the pass while any caller was turned away.
3132    ///
3133    /// The test hook freezes the holder's pass in that window to make the race
3134    /// deterministic: the holder parks having found the budget satisfied, the
3135    /// main thread inserts over budget and is turned away, then the holder
3136    /// resumes. The `gate` is used for both rendezvous.
3137    #[mz_ore::test]
3138    fn racing_insert_is_not_dropped_by_budget_single_flight() {
3139        // Budget for exactly one small chunk.
3140        let budget = 64 << 10;
3141        let pool = test_pool(budget);
3142        let gate = std::sync::Arc::new(std::sync::Barrier::new(2));
3143
3144        let holder = {
3145            let pool = pool.clone();
3146            let gate = std::sync::Arc::clone(&gate);
3147            std::thread::spawn(move || -> ChunkHandle {
3148                ENFORCE_BUDGET_HOOK.with(|cell| {
3149                    *cell.borrow_mut() = Some(Box::new(move || {
3150                        gate.wait(); // parked, holding the guard
3151                        gate.wait(); // resume once the race is done
3152                    }));
3153                });
3154                // At budget: the pass finds it satisfied and parks at the hook.
3155                insert(&pool, &mut payload(SMALL, 1))
3156            })
3157        };
3158
3159        gate.wait(); // holder is parked in enforcement, holding the guard
3160        // Push over budget; this insert is turned away by the held guard.
3161        let _over = insert(&pool, &mut payload(SMALL, 2));
3162        gate.wait(); // let the holder resume and release the guard
3163        // Kept alive past the assert: freeing it would drop its bytes and mask
3164        // the overshoot.
3165        let _held = holder.join().expect("holder panicked");
3166
3167        // Nothing re-trips enforcement, so the pool must not be left over budget.
3168        let resident = pool.stats().resident_bytes;
3169        assert!(
3170            resident <= u64::cast_from(budget),
3171            "resident {resident} exceeds budget {budget}: racing insert escaped enforcement",
3172        );
3173    }
3174
3175    #[mz_ore::test]
3176    fn set_budget_retunes_in_place() {
3177        let pool = test_pool(usize::MAX);
3178        let mut handles = Vec::new();
3179        for seed in 0..8 {
3180            handles.push(insert(&pool, &mut payload(SMALL, 200 + seed)));
3181        }
3182        assert_eq!(pool.stats().evictions_compress, 0);
3183
3184        // Shrinking the budget evicts immediately.
3185        pool.set_budget(128 << 10);
3186        let stats = pool.stats();
3187        assert!(stats.resident_bytes <= 128 << 10);
3188        assert!(stats.evictions_compress >= 6);
3189
3190        // Growing it leaves headroom: a fresh insert stays resident.
3191        pool.set_budget(usize::MAX);
3192        let h = insert(&pool, &mut payload(SMALL, 300));
3193        assert_eq!(h.residency(), Residency::UnbackedResident);
3194        for h in &handles {
3195            assert_eq!(read(h).len(), SMALL);
3196        }
3197    }
3198
3199    #[mz_ore::test]
3200    fn second_chance_prefers_untouched_victims() {
3201        // Budget holds one and a half small chunks.
3202        let pool = test_pool((64 << 10) + (32 << 10));
3203        let orig_a = payload(SMALL, 8);
3204        let handle_a = insert(&pool, &mut orig_a.clone());
3205        assert_eq!(read(&handle_a), orig_a);
3206        // Inserting B overflows the budget; A is older but touched, so the
3207        // enforcer gives it a second chance and evicts untouched B instead.
3208        let handle_b = insert(&pool, &mut payload(SMALL, 9));
3209        assert_eq!(handle_a.residency(), Residency::UnbackedResident);
3210        assert_eq!(handle_b.residency(), Residency::Evicted);
3211    }
3212
3213    /// Depth-hinted chunks are evicted before younger ones: the deep chunk
3214    /// loses even though the young chunk is older and both are untouched
3215    /// (plain FIFO would have evicted the older, young one). Also exercises
3216    /// band clamping: depths beyond the last band share it.
3217    #[mz_ore::test]
3218    fn eviction_prefers_deeper_chunks() {
3219        // Budget of one small chunk.
3220        let pool = test_pool(64 << 10);
3221        let young = insert(&pool, &mut payload(SMALL, 900));
3222        let deep = insert_at_depth(&pool, 255, &mut payload(SMALL, 901));
3223        assert_eq!(young.residency(), Residency::UnbackedResident);
3224        assert_eq!(deep.residency(), Residency::Evicted);
3225    }
3226
3227    /// Eager backing visits deeper chunks first, mirroring eviction order,
3228    /// so the chunks evicted first are the ones already backed.
3229    #[mz_ore::test]
3230    fn backing_prefers_deeper_chunks() {
3231        let pool = test_pool(256 << 20);
3232        let young = insert(&pool, &mut payload(SMALL, 902));
3233        let deep = insert_at_depth(&pool, 2, &mut payload(SMALL, 903));
3234        assert!(pool.back_step());
3235        assert_eq!(deep.residency(), Residency::BackedResident);
3236        assert_eq!(young.residency(), Residency::UnbackedResident);
3237        assert!(pool.back_step());
3238        assert_eq!(young.residency(), Residency::BackedResident);
3239    }
3240
3241    #[mz_ore::test]
3242    fn empty_insert_consumes_no_slot() {
3243        let pool = test_pool(256 << 20);
3244        let mut data = Vec::new();
3245        let handle = insert(&pool, &mut data);
3246        assert_eq!(handle.size_class_bytes(), None);
3247        assert!(read(&handle).is_empty());
3248        // Reads clear the destination even for empty chunks.
3249        let mut out = vec![1u64, 2, 3];
3250        handle.read_into(&mut out);
3251        assert!(out.is_empty());
3252        drop(handle);
3253        let stats = pool.stats();
3254        assert_eq!(stats.resident_bytes, 0);
3255        assert_eq!(stats.writes_elided, 0);
3256    }
3257
3258    #[mz_ore::test]
3259    fn oversize_round_trips() {
3260        let pool = test_pool(256 << 20);
3261        let words = SIZE_CLASSES[SIZE_CLASSES.len() - 1] / 8 + 1;
3262        let orig = payload(words, 10);
3263        let handle = insert(&pool, &mut orig.clone());
3264        assert_eq!(handle.residency(), Residency::Oversize);
3265        assert_eq!(handle.size_class_bytes(), None);
3266        let stats = pool.stats();
3267        assert_eq!(stats.oversize_bytes, u64::cast_from(words * 8));
3268        // The payload outgrew the largest class, and no class was exhausted.
3269        assert_eq!(stats.oversize_payloads, 1);
3270        assert_eq!(stats.slot_exhausted_fallbacks, 0);
3271        // Explicit eviction and budget enforcement leave oversize chunks
3272        // resident.
3273        pool.evict(&handle);
3274        pool.enforce_budget();
3275        assert_eq!(handle.residency(), Residency::Oversize);
3276        assert_eq!(read(&handle), orig);
3277        drop(handle);
3278        let stats = pool.stats();
3279        assert_eq!(stats.oversize_bytes, 0);
3280        assert_eq!(stats.resident_bytes, 0);
3281    }
3282
3283    #[mz_ore::test]
3284    fn payload_lands_in_smallest_fitting_class() {
3285        let pool = test_pool(256 << 20);
3286        let handle = insert(&pool, &mut payload((100 << 10) / 8, 11));
3287        assert_eq!(handle.size_class_bytes(), Some(128 << 10));
3288        let exact = insert(&pool, &mut payload(SMALL, 12));
3289        assert_eq!(exact.size_class_bytes(), Some(64 << 10));
3290    }
3291
3292    #[mz_ore::test]
3293    #[cfg_attr(miri, ignore)] // too slow
3294    fn multithreaded_smoke() {
3295        // Budget of one small chunk: four inserting threads keep the pool
3296        // over budget, so every insert's enforcement pass selects victims
3297        // owned by other threads, racing cross-thread eviction against
3298        // copy-out reads and frees.
3299        let pool = test_pool(64 << 10);
3300        let per_thread = rounds(50, 3);
3301        let threads: Vec<_> = (0..4u64)
3302            .map(|t| {
3303                let pool = pool.clone();
3304                std::thread::spawn(move || {
3305                    for round in 0..per_thread {
3306                        let seed = t * 1000 + round;
3307                        let orig = payload(SMALL, seed);
3308                        let handle = insert(&pool, &mut orig.clone());
3309                        pool.evict(&handle);
3310                        assert_eq!(read(&handle), orig);
3311                        // Enforcement racing reads must never corrupt them.
3312                        pool.enforce_budget();
3313                        assert_eq!(read(&handle), orig);
3314                        drop(handle);
3315                    }
3316                })
3317            })
3318            .collect();
3319        for thread in threads {
3320            thread.join().expect("worker thread panicked");
3321        }
3322        let stats = pool.stats();
3323        assert_eq!(stats.inserts, 4 * per_thread);
3324        assert_eq!(stats.frees, 4 * per_thread);
3325        assert_eq!(stats.resident_bytes, 0);
3326    }
3327
3328    #[mz_ore::test]
3329    #[cfg_attr(miri, ignore)] // too slow
3330    fn concurrent_read_enforce_churn() {
3331        // Races the three actors that can touch one chunk's slot: readers
3332        // copying shared chunks out and verifying them, an enforcer evicting
3333        // them (the zero budget makes every chunk a victim), and a churner
3334        // whose insert/free traffic turns the queue over. Contents are
3335        // asserted on every read, so an eviction or slot recycle racing a
3336        // copy-out shows up as corruption.
3337        let pool = test_pool(0);
3338        let shared: Arc<Vec<(Vec<u64>, ChunkHandle)>> = Arc::new(
3339            (0..4u64)
3340                .map(|seed| {
3341                    let orig = payload(SMALL, 600 + seed);
3342                    let handle = insert(&pool, &mut orig.clone());
3343                    (orig, handle)
3344                })
3345                .collect(),
3346        );
3347        let churn = rounds(300, 6);
3348        let mut threads = Vec::new();
3349        for t in 0..2u64 {
3350            let shared = Arc::clone(&shared);
3351            threads.push(std::thread::spawn(move || {
3352                for round in 0..churn {
3353                    let (orig, handle) = &shared[usize::cast_from((t + round) % 4)];
3354                    assert_eq!(&read(handle), orig);
3355                }
3356            }));
3357        }
3358        {
3359            let pool = pool.clone();
3360            threads.push(std::thread::spawn(move || {
3361                for _ in 0..2 * churn {
3362                    pool.enforce_budget();
3363                }
3364            }));
3365        }
3366        {
3367            let pool = pool.clone();
3368            threads.push(std::thread::spawn(move || {
3369                for round in 0..churn {
3370                    let orig = payload(SMALL, 700 + round);
3371                    let handle = insert(&pool, &mut orig.clone());
3372                    assert_eq!(read(&handle), orig);
3373                }
3374            }));
3375        }
3376        for thread in threads {
3377            thread.join().expect("worker thread panicked");
3378        }
3379        drop(shared);
3380        assert_eq!(pool.stats().resident_bytes, 0);
3381    }
3382
3383    /// Read-only traffic never raises resident bytes: every chunk starts
3384    /// evicted and is then read once, with no inserts in between. Reads copy
3385    /// out of the extents and leave every chunk evicted, so a seek-heavy
3386    /// phase costs no pool memory at all.
3387    #[mz_ore::test]
3388    fn reads_never_raise_resident_bytes() {
3389        let pool = test_pool(128 << 10);
3390        let origs: Vec<_> = (0..8u64).map(|seed| payload(SMALL, 300 + seed)).collect();
3391        let handles: Vec<_> = origs
3392            .iter()
3393            .map(|o| insert(&pool, &mut o.clone()))
3394            .collect();
3395        for handle in &handles {
3396            pool.evict(handle);
3397        }
3398        assert_eq!(pool.stats().resident_bytes, 0);
3399        for (index, handle) in handles.iter().enumerate() {
3400            assert_eq!(read(handle), origs[index]);
3401            assert_eq!(handle.residency(), Residency::Evicted);
3402            assert_eq!(pool.stats().resident_bytes, 0);
3403        }
3404    }
3405
3406    /// Evict-then-free churn under a generous RSS target: the compressed
3407    /// tier never crosses its cap, so enforcement never visits (and never
3408    /// drops) extent-queue entries, and pruning alone must keep the queue
3409    /// proportional to the live resident extents.
3410    #[mz_ore::test]
3411    #[cfg_attr(miri, ignore)] // too slow
3412    fn extent_queue_stays_bounded_under_cap() {
3413        let pool = test_pool(256 << 20);
3414        pool.set_rss_target(1 << 40);
3415        for seed in 0..rounds(1000, 48) {
3416            let handle = insert(&pool, &mut payload(SMALL, seed));
3417            pool.evict(&handle);
3418            drop(handle);
3419        }
3420        assert_eq!(pool.stats().extent_resident_bytes, 0);
3421        let len = pool.extent_queue_len();
3422        assert!(
3423            len <= 32,
3424            "extent queue holds {len} entries for zero resident extents",
3425        );
3426    }
3427
3428    /// A warm slot reused for a smaller payload round-trips: the tail
3429    /// release past the new payload must not disturb the payload itself,
3430    /// and the ledger credits exactly the payload.
3431    #[mz_ore::test]
3432    fn warm_reuse_with_smaller_payload_round_trips() {
3433        // Budget 8 MiB: warm cap = 1 MiB, so a 64 KiB slot parks warm.
3434        let pool = test_pool(8 << 20);
3435        let full = insert(&pool, &mut payload(SMALL, 60));
3436        drop(full);
3437        assert_eq!(pool.stats().warm_bytes, 64 << 10, "freed slot parks warm");
3438        // A payload of just over a page reuses the warm slot; the slot's
3439        // pages past it are released.
3440        let words = 4096 / 8 + 1;
3441        let orig = payload(words, 61);
3442        let handle = insert(&pool, &mut orig.clone());
3443        let stats = pool.stats();
3444        assert_eq!(stats.warm_reuses, 1, "reused the warm slot");
3445        assert_eq!(stats.resident_bytes, u64::cast_from(words * 8));
3446        assert_eq!(read(&handle), orig);
3447        // Round-trips through the extent as well.
3448        pool.evict(&handle);
3449        pool.poison_free_slots();
3450        assert_eq!(read(&handle), orig);
3451        drop(handle);
3452        assert_eq!(pool.stats().resident_bytes, 0);
3453    }
3454
3455    /// Heap-backed chunks count as resident but can never be evicted, so
3456    /// the budget must not force slotted chunks out on their account: with
3457    /// unevictable bytes alone exceeding the budget, a slotted chunk that
3458    /// fits the budget stays resident.
3459    #[mz_ore::test]
3460    fn unevictable_bytes_do_not_force_eviction() {
3461        // One 64 KiB slot per class: the second and third inserts fall
3462        // back to the heap.
3463        let pool = Pool::with_class_capacity(64 << 10).expect("pool creation");
3464        pool.set_budget(64 << 10);
3465        let slotted = insert(&pool, &mut payload(SMALL, 91));
3466        let heap_a = insert(&pool, &mut payload(SMALL, 92));
3467        let heap_b = insert(&pool, &mut payload(SMALL, 93));
3468        assert_eq!(heap_a.residency(), Residency::Oversize);
3469        assert_eq!(heap_b.residency(), Residency::Oversize);
3470        let stats = pool.stats();
3471        assert!(stats.oversize_bytes > 64 << 10, "unevictable exceed budget");
3472        assert_eq!(slotted.residency(), Residency::UnbackedResident);
3473        assert_eq!(stats.evictions_compress, 0);
3474        assert_eq!(read(&slotted), payload(SMALL, 91));
3475        assert_eq!(read(&heap_a), payload(SMALL, 92));
3476    }
3477
3478    /// A slotless empty chunk survives an explicit evict with spill
3479    /// scheduling enabled: nothing is handed to the spill threads and the
3480    /// chunk stays readable.
3481    #[mz_ore::test]
3482    fn evict_of_empty_chunk_is_a_no_op() {
3483        let pool = test_pool(usize::MAX);
3484        pool.enable_spill_without_threads();
3485        let empty = insert(&pool, &mut Vec::new());
3486        pool.evict(&empty);
3487        assert_eq!(empty.residency(), Residency::UnbackedResident);
3488        assert!(!pool.spill_step(), "nothing was scheduled");
3489        assert_eq!(pool.stats().spill_scheduled, 0);
3490        assert!(read(&empty).is_empty());
3491    }
3492
3493    #[mz_ore::test]
3494    fn queue_stays_bounded_under_budget() {
3495        // Chunk churn that never exceeds the budget: the enforcer's eviction
3496        // loop never runs, so stale queue entries must be reclaimed by
3497        // pruning alone.
3498        let pool = test_pool(256 << 20);
3499        for seed in 0..rounds(1000, 48) {
3500            let handle = insert(&pool, &mut payload(SMALL, seed));
3501            drop(handle);
3502        }
3503        let len = pool.queue_len();
3504        assert!(len <= 32, "queue holds {len} entries for zero live chunks");
3505    }
3506
3507    #[mz_ore::test]
3508    fn spill_async_evict_round_trip() {
3509        let pool = test_pool(usize::MAX);
3510        pool.enable_spill_without_threads();
3511        let h = insert(&pool, &mut payload(SMALL, 400));
3512        pool.evict(&h);
3513        assert_eq!(h.residency(), Residency::WriteInFlight);
3514        // Readable while in flight: the slot is still populated, and the
3515        // copy-out coexists with the spill thread's compression read.
3516        assert_eq!(read(&h), payload(SMALL, 400));
3517        // Reads leave no trace, so the eviction commits.
3518        assert!(pool.spill_step());
3519        assert_eq!(h.residency(), Residency::Evicted);
3520        let stats = pool.stats();
3521        assert_eq!(stats.spill_scheduled, 1);
3522        assert_eq!(stats.evictions_compress, 1);
3523        pool.poison_free_slots();
3524        assert_eq!(read(&h), payload(SMALL, 400));
3525    }
3526
3527    #[mz_ore::test]
3528    fn spill_freed_while_queued_is_elided() {
3529        let pool = test_pool(usize::MAX);
3530        pool.enable_spill_without_threads();
3531        let h = insert(&pool, &mut payload(SMALL, 401));
3532        pool.evict(&h);
3533        assert_eq!(h.residency(), Residency::WriteInFlight);
3534        drop(h);
3535        assert!(pool.spill_step());
3536        let stats = pool.stats();
3537        assert_eq!(stats.spill_cancelled, 1);
3538        assert_eq!(stats.writes_elided, 1, "freed before compression: elided");
3539        assert_eq!(stats.extent_bytes_written, 0, "no extent was written");
3540        assert_eq!(stats.resident_bytes, 0, "slot accounting settled");
3541    }
3542
3543    #[mz_ore::test]
3544    fn spill_take_in_flight_cancels_write() {
3545        let pool = test_pool(usize::MAX);
3546        pool.enable_spill_without_threads();
3547        let orig = payload(SMALL, 402);
3548        let h = insert(&pool, &mut orig.clone());
3549        pool.evict(&h);
3550        assert_eq!(h.residency(), Residency::WriteInFlight);
3551        let mut out = Vec::new();
3552        h.take(&mut out);
3553        assert_eq!(out, orig);
3554        assert!(pool.spill_step());
3555        let stats = pool.stats();
3556        assert_eq!(stats.frees, 1);
3557        assert_eq!(stats.spill_cancelled, 1);
3558        assert_eq!(stats.writes_elided, 1, "taken before compression: elided");
3559        assert_eq!(stats.extent_bytes_written, 0, "no extent was written");
3560        assert_eq!(stats.resident_bytes, 0, "slot accounting settled");
3561        assert_eq!(stats.live_chunks, 0);
3562    }
3563
3564    #[mz_ore::test]
3565    #[cfg_attr(miri, ignore)] // too slow
3566    fn spill_threads_end_to_end() {
3567        let pool = test_pool(128 << 10);
3568        pool.set_spill_threads(2);
3569        let mut handles = Vec::new();
3570        for seed in 0..rounds(16, 6) {
3571            handles.push(insert(&pool, &mut payload(SMALL, 500 + seed)));
3572        }
3573        pool.quiesce_spill();
3574        let stats = pool.stats();
3575        assert!(
3576            stats.spill_scheduled > 0,
3577            "budget pressure should have scheduled spills",
3578        );
3579        for (i, h) in handles.iter().enumerate() {
3580            assert_eq!(read(h), payload(SMALL, 500 + u64::cast_from(i)));
3581        }
3582        pool.join_spill_threads();
3583    }
3584
3585    /// Races the `WriteInFlight` protocol in its true concurrent form:
3586    /// spill threads compress slots without the state lock while owner
3587    /// threads copy the same chunks out under it and drop chunks mid-flight
3588    /// (both cancellation windows). Contents are asserted on every read, so
3589    /// a compression or slot release racing a copy-out shows up as
3590    /// corruption; under Miri the aliasing itself is checked.
3591    #[mz_ore::test]
3592    fn spill_threads_race_reads_and_drops() {
3593        let pool = test_pool(usize::MAX);
3594        pool.set_spill_threads(2);
3595        let iters = rounds(50, 6);
3596        let mut threads = Vec::new();
3597        for t in 0..2u64 {
3598            let pool = pool.clone();
3599            threads.push(std::thread::spawn(move || {
3600                for round in 0..iters {
3601                    let orig = payload(SMALL, t * 10_000 + round);
3602                    let handle = insert(&pool, &mut orig.clone());
3603                    // Hands the chunk to the spill threads (`WriteInFlight`).
3604                    pool.evict(&handle);
3605                    // Copy-out read racing the unlocked compression read.
3606                    assert_eq!(read(&handle), orig);
3607                    if round % 2 == 0 {
3608                        // Free while queued or mid-compression: the
3609                        // cancellation windows own the deferred cleanup.
3610                        drop(handle);
3611                    } else {
3612                        assert_eq!(read(&handle), orig);
3613                    }
3614                }
3615            }));
3616        }
3617        for thread in threads {
3618            thread.join().expect("worker thread panicked");
3619        }
3620        pool.quiesce_spill();
3621        pool.join_spill_threads();
3622        assert_eq!(pool.stats().resident_bytes, 0);
3623    }
3624
3625    /// The identity codec stores the body verbatim: eviction and reads,
3626    /// whole and by range, reconstruct it unchanged.
3627    #[mz_ore::test]
3628    fn identity_codec_round_trips() {
3629        let pool = test_pool(usize::MAX);
3630        let want = payload(SMALL, 601);
3631        let h = pool.insert_with(SMALL, ChunkHints::default(), &IDENTITY_CODEC, |dst| {
3632            dst.copy_from_slice(&want);
3633        });
3634        assert_eq!(read(&h), want);
3635        pool.evict(&h);
3636        assert_eq!(read(&h), want, "round-trips through the extent");
3637        pool.evict(&h);
3638        let mut range = Vec::new();
3639        h.read_range_into(8..24, &mut range);
3640        assert_eq!(range, want[8..24], "range reads copy the range directly");
3641    }
3642
3643    #[mz_ore::test]
3644    fn insert_with_fills_in_place() {
3645        let pool = test_pool(usize::MAX);
3646        let want = payload(SMALL, 600);
3647        let h = pool.insert_with(SMALL, ChunkHints::default(), &TEST_CODEC, |dst| {
3648            assert_eq!(dst.len(), SMALL, "fill sees exactly the chunk length");
3649            dst.copy_from_slice(&want);
3650        });
3651        assert_eq!(h.residency(), Residency::UnbackedResident);
3652        assert_eq!(read(&h), want);
3653        pool.evict(&h);
3654        assert_eq!(read(&h), want, "round-trips through the extent");
3655
3656        // Empty and oversize take their fallback paths.
3657        let empty = pool.insert_with(0, ChunkHints::default(), &TEST_CODEC, |dst| {
3658            assert!(dst.is_empty())
3659        });
3660        assert!(read(&empty).is_empty());
3661        let big_len = (SIZE_CLASSES[SIZE_CLASSES.len() - 1] / 8) + 1;
3662        let big = pool.insert_with(big_len, ChunkHints::default(), &TEST_CODEC, |dst| {
3663            dst.fill(7)
3664        });
3665        assert_eq!(big.residency(), Residency::Oversize);
3666        assert_eq!(read(&big).len(), big_len);
3667    }
3668
3669    #[mz_ore::test]
3670    fn slot_exhaustion_degrades_to_heap() {
3671        // Two 64 KiB slots per class at this capacity; the third insert finds
3672        // no slot and must fall back to the heap rather than panic.
3673        let pool = Pool::with_class_capacity(128 << 10).expect("pool creation");
3674        let a = insert(&pool, &mut payload(SMALL, 700));
3675        let b = insert(&pool, &mut payload(SMALL, 701));
3676        let c = insert(&pool, &mut payload(SMALL, 702));
3677        assert_eq!(a.residency(), Residency::UnbackedResident);
3678        assert_eq!(b.residency(), Residency::UnbackedResident);
3679        assert_eq!(
3680            c.residency(),
3681            Residency::Oversize,
3682            "fallback is heap-backed"
3683        );
3684        assert_eq!(pool.stats().slot_exhausted_fallbacks, 1);
3685        assert_eq!(read(&c), payload(SMALL, 702));
3686        // Freeing a slotted chunk lets the next insert use the region again.
3687        drop(a);
3688        let d = insert(&pool, &mut payload(SMALL, 703));
3689        assert_eq!(d.residency(), Residency::UnbackedResident);
3690        assert_eq!(read(&d), payload(SMALL, 703));
3691    }
3692}