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