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