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