mz_ore/pool/region.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//! Size-class virtual-memory regions for the buffer pool.
17//!
18//! One [`Region`] per size class, each a single anonymous `mmap` reservation.
19//! The reservation is virtual; physical memory materializes on first write to
20//! a slot. Slots are scoped to residency: eviction releases a slot's physical
21//! pages with [`dontneed`] and returns the slot index to the free list, so a
22//! chunk holds a slot only from insert until its eviction, and the pool reads
23//! slots strictly copy-out under the owning chunk's state lock.
24//!
25//! Two fault-amortization mechanisms soften the cost of cycling slots:
26//!
27//! * Regions whose class is at least one huge page are aligned to the huge
28//! page and advised `MADV_HUGEPAGE`, so populating a large slot costs one
29//! fault instead of one per 4 KiB.
30//! * The free list is split into a *warm* side (pages kept resident; reuse
31//! faults nothing and skips the kernel's page zeroing) and a *cold* side
32//! (pages released). The pool decides which side a freed slot joins,
33//! bounding total warm bytes as a fraction of its budget.
34//!
35//! TODO: consider `mlock`ing slot regions so kernel swap can never write
36//! out pages the engine would discard or write better itself. Needs
37//! `RLIMIT_MEMLOCK` tuning before it can be on by default.
38//!
39//! All platform access goes through the [`sys`] seam: mapping, unmapping,
40//! paging advice, and the page size. Under Miri the seam swaps to a
41//! Rust-heap backing with advice as a contents-preserving no-op, so the
42//! pool's tests (including the unsafe slot borrows they exercise) run under
43//! the interpreter.
44
45use std::io;
46use std::sync::Mutex;
47
48use crate::cast::CastFrom;
49
50/// Chunk size classes in bytes, smallest first. The pool places each chunk in
51/// the smallest class that fits its payload.
52///
53/// The top classes deliberately overshoot the batchers' nominal ~2 MiB chunk
54/// target: the ship heuristic re-targets the next 2 MiB boundary whenever a
55/// single push crosses one, so real chunk sizes are multimodal with bands
56/// just under each boundary. A class that fits only the nominal target sends
57/// the higher bands to the unpageable heap fallback. Slot internal
58/// fragmentation is virtual-only — slots populate lazily, so a chunk costs
59/// physical memory for its payload, not its class size.
60pub(crate) const SIZE_CLASSES: [usize; 8] = [
61 64 << 10,
62 128 << 10,
63 256 << 10,
64 512 << 10,
65 1 << 20,
66 2 << 20,
67 4 << 20,
68 8 << 20,
69];
70
71/// The smallest size class that fits a payload of `len_bytes`, or `None`
72/// when even the largest class is too small. A selected class always fits:
73/// `SIZE_CLASSES[class] >= len_bytes`, the bound the pool turns into slice
74/// lengths over slot memory (proved by the Kani harnesses).
75pub(crate) fn size_class_for(len_bytes: usize) -> Option<usize> {
76 SIZE_CLASSES.iter().position(|&c| c >= len_bytes)
77}
78
79/// One anonymous virtual-memory reservation serving fixed-size slots of a
80/// single size class.
81#[derive(Debug)]
82pub(crate) struct Region {
83 base: *mut u8,
84 capacity: usize,
85 /// Alignment of `base`, needed to return the mapping to [`sys::unmap`].
86 align: usize,
87 class_size: usize,
88 slots: Mutex<SlotAllocator>,
89}
90
91/// Free-list-plus-bump slot allocator. A slot index returns to a free list
92/// whenever its chunk stops being resident — eviction and free alike. Warm
93/// slots keep their physical pages (reuse is fault-free); cold slots had
94/// theirs released. Never-allocated slots beyond the high-water mark are
95/// untouched virtual space and fault on first write like cold ones.
96#[derive(Debug)]
97struct SlotAllocator {
98 free_warm: Vec<u32>,
99 free_cold: Vec<u32>,
100 high_water: u32,
101 max_slots: u32,
102}
103
104impl SlotAllocator {
105 fn new(max_slots: u32) -> SlotAllocator {
106 SlotAllocator {
107 free_warm: Vec::new(),
108 free_cold: Vec::new(),
109 high_water: 0,
110 max_slots,
111 }
112 }
113
114 /// Allocates a slot index, or `None` when every slot is in use; the flag
115 /// reports whether the slot came from the warm list. Warm slots are
116 /// preferred, then cold, then never-touched bump slots.
117 fn alloc(&mut self) -> Option<(u32, bool)> {
118 if let Some(slot) = self.free_warm.pop() {
119 return Some((slot, true));
120 }
121 if let Some(slot) = self.free_cold.pop() {
122 return Some((slot, false));
123 }
124 if self.high_water == self.max_slots {
125 return None;
126 }
127 let slot = self.high_water;
128 self.high_water += 1;
129 Some((slot, false))
130 }
131
132 /// Returns a previously allocated slot to the warm or cold free list.
133 fn free(&mut self, slot: u32, warm: bool) {
134 debug_assert!(slot < self.high_water);
135 if warm {
136 self.free_warm.push(slot);
137 } else {
138 self.free_cold.push(slot);
139 }
140 }
141}
142
143// SAFETY: `base` points at an anonymous mapping owned exclusively by this
144// `Region` for its whole lifetime. Slot allocation is serialized by the
145// `slots` mutex, and access to a slot's contents is serialized by the
146// owning chunk's state mutex; the raw pointer itself carries no thread
147// affinity.
148unsafe impl Send for Region {}
149// SAFETY: see the `Send` justification; all interior mutability is behind
150// the `slots` mutex, and disjoint slots are written only by their owning
151// chunks.
152unsafe impl Sync for Region {}
153
154impl Region {
155 /// Reserves a region of `capacity_bytes` (rounded down to a whole number
156 /// of slots) for slots of `class_size` bytes.
157 ///
158 /// On Linux, regions whose class is at least [`HUGE_PAGE`] are aligned to
159 /// the huge page and advised `MADV_HUGEPAGE`: their slots tile huge-page
160 /// boundaries exactly, so populating a slot is one huge-page fault rather
161 /// than one fault per 4 KiB, and a whole-slot [`dontneed`] frees whole
162 /// huge pages without splitting any. Regions whose class is below the
163 /// huge page are advised `MADV_NOHUGEPAGE` instead, so reclaim over a
164 /// single slot stays page-exact.
165 pub(crate) fn new(class_size: usize, capacity_bytes: usize) -> io::Result<Region> {
166 Region::with_huge_pages(class_size, capacity_bytes, class_size >= HUGE_PAGE)
167 }
168
169 /// As [`Region::new`], but the region opts out of transparent huge pages
170 /// regardless of class size. For regions whose slots are reclaimed with
171 /// `MADV_PAGEOUT` under observed residency: reclaiming a range a large
172 /// folio covers only partially needs a folio split, and a failed split
173 /// silently leaves pages resident.
174 pub(crate) fn new_nohuge(class_size: usize, capacity_bytes: usize) -> io::Result<Region> {
175 Region::with_huge_pages(class_size, capacity_bytes, false)
176 }
177
178 fn with_huge_pages(class_size: usize, capacity_bytes: usize, huge: bool) -> io::Result<Region> {
179 let page = sys::page_size();
180 assert!(class_size > 0 && class_size % page == 0);
181 let capacity = capacity_bytes - capacity_bytes % class_size;
182 if capacity == 0 {
183 // A capacity below one slot yields an empty region: `alloc`
184 // always answers `None` and the caller's exhaustion fallback
185 // carries the class. No mapping exists; drop has nothing to do.
186 return Ok(Region {
187 base: std::ptr::NonNull::<u8>::dangling().as_ptr(),
188 capacity: 0,
189 align: page,
190 class_size,
191 slots: Mutex::new(SlotAllocator::new(0)),
192 });
193 }
194 let max_slots = u32::try_from(capacity / class_size).expect("slot count fits u32");
195 let align = if cfg!(target_os = "linux") && huge {
196 HUGE_PAGE
197 } else {
198 page
199 };
200 let base = sys::map(capacity, align)?;
201 if !huge {
202 // Opt the whole region out of transparent huge pages before any
203 // slot is touched. Production hosts run
204 // `transparent_hugepage=madvise`, where an unadvised region gets
205 // no fault-time folios anyway, so this is defense in depth for
206 // hosts running `always`: there, fault-time folios would straddle
207 // sub-huge-page slot boundaries, and khugepaged (which tolerates
208 // up to `max_ptes_none` empty entries) re-collapses ranges a
209 // slot-granular [`dontneed`] just reclaimed, resurrecting the
210 // released memory.
211 nohugepage(base, capacity);
212 }
213 // Pool contents are recreatable spill data: keep the (potentially
214 // huge) anonymous reservation out of core dumps.
215 dontdump(base, capacity);
216 Ok(Region {
217 base,
218 capacity,
219 align,
220 class_size,
221 slots: Mutex::new(SlotAllocator::new(max_slots)),
222 })
223 }
224
225 /// Size in bytes of every slot in this region.
226 pub(crate) fn class_size(&self) -> usize {
227 self.class_size
228 }
229
230 /// Allocates a slot index, or `None` if every slot of the class is in
231 /// use; the flag reports whether the slot came from the warm list (its
232 /// pages are resident; writing it faults nothing).
233 ///
234 /// Slots are scoped to residency (eviction frees them), so demand scales
235 /// with the *resident* set — bounded by the pool budget plus in-flight
236 /// slack — and exhaustion means residency outgrew the class reservation;
237 /// callers degrade rather than fail.
238 pub(crate) fn alloc(&self) -> Option<(u32, bool)> {
239 self.slots
240 .lock()
241 .expect("region allocator poisoned")
242 .alloc()
243 }
244
245 /// Returns a slot to the warm or cold free list. The caller must be
246 /// freeing the chunk that owned the slot, must have released the slot's
247 /// physical pages iff `warm` is false, and owns the warm-bytes accounting
248 /// that bounds the warm side.
249 pub(crate) fn free(&self, slot: u32, warm: bool) {
250 self.slots
251 .lock()
252 .expect("region allocator poisoned")
253 .free(slot, warm);
254 }
255
256 /// Moves warm free slots to the cold list, releasing their physical
257 /// pages, until at least `want_bytes` have been cooled or no warm slot
258 /// remains. Returns the bytes cooled. The caller owns the warm-bytes
259 /// accounting that bounds the warm side.
260 pub(crate) fn cool_warm_slots(&self, want_bytes: usize) -> usize {
261 let mut slots = self.slots.lock().expect("region allocator poisoned");
262 let mut cooled = 0;
263 while cooled < want_bytes {
264 let Some(slot) = slots.free_warm.pop() else {
265 break;
266 };
267 let offset = usize::cast_from(slot) * self.class_size;
268 // SAFETY: the slot is on a free list and the allocator mutex is
269 // held, so no chunk owns it and no allocation can race; the
270 // range stays within the region's mapping.
271 unsafe {
272 dontneed(self.base.add(offset), self.class_size);
273 }
274 slots.free_cold.push(slot);
275 cooled += self.class_size;
276 }
277 cooled
278 }
279
280 /// Test hook: overwrites every free slot with `0xDE` so stale contents
281 /// cannot masquerade as correct data when a slot is reused.
282 #[cfg(test)]
283 pub(crate) fn poison_free_slots(&self) {
284 let slots = self.slots.lock().expect("region allocator poisoned");
285 for &slot in slots.free_warm.iter().chain(slots.free_cold.iter()) {
286 let offset = usize::cast_from(slot) * self.class_size;
287 // SAFETY: the slot is on a free list and the allocator mutex is
288 // held, so no chunk owns it and no allocation can race; the write
289 // stays within the region's mapping.
290 unsafe {
291 std::ptr::write_bytes(self.base.add(offset), 0xDE, self.class_size);
292 }
293 }
294 }
295
296 /// The base address of a slot, fixed while its owning chunk is resident.
297 pub(crate) fn slot_ptr(&self, slot: u32) -> *mut u8 {
298 let offset = usize::cast_from(slot) * self.class_size;
299 debug_assert!(offset + self.class_size <= self.capacity);
300 // SAFETY: `slot` was handed out by `alloc`, so `offset + class_size`
301 // lies within the single `capacity`-byte mapping that `base` points
302 // to; the add stays in bounds of one allocated object.
303 unsafe { self.base.add(offset) }
304 }
305}
306
307impl Drop for Region {
308 fn drop(&mut self) {
309 // Empty regions never created a mapping.
310 if self.capacity == 0 {
311 return;
312 }
313 // SAFETY: `base`/`capacity`/`align` describe exactly the mapping
314 // created in `new`, and dropping the region means no chunk (and
315 // hence no outstanding borrow) refers into it any longer.
316 unsafe {
317 sys::unmap(self.base, self.capacity, self.align);
318 }
319 }
320}
321
322/// The transparent-huge-page size assumed for region alignment. Linux x86-64
323/// and aarch64 (4 KiB base pages) both use 2 MiB; if a platform differs, the
324/// alignment is merely unhelpful, never wrong.
325pub(crate) const HUGE_PAGE: usize = 2 << 20;
326
327/// Releases the physical pages of the page-aligned subrange of
328/// `[ptr, ptr + len)`, keeping the virtual range mapped.
329///
330/// # Safety
331///
332/// The range must lie within a live mapping exclusively owned by the caller,
333/// with no outstanding references into it. After the call the range's contents
334/// are undefined: Linux replaces them with zero pages, but other systems
335/// (macOS in particular) may keep the old bytes resident, so callers must
336/// fully overwrite the range before reading it again.
337pub(crate) unsafe fn dontneed(ptr: *mut u8, len: usize) {
338 // SAFETY: forwarding this function's contract.
339 unsafe { sys::advise(ptr, len, sys::Advice::DontNeed) }
340}
341
342/// Hints the kernel to reclaim the page-aligned subrange of `[ptr, ptr + len)`
343/// immediately, writing it to the swap device. Contents are preserved; this is
344/// a non-destructive hint. No-op outside Linux.
345pub(crate) fn pageout(ptr: *mut u8, len: usize) {
346 // SAFETY: the advice is a contents-preserving hint, so the only
347 // obligation is that the range lies in a live mapping, which callers
348 // guarantee by passing a live slot or extent allocation.
349 unsafe { sys::advise(ptr, len, sys::Advice::PageOut) }
350}
351
352/// Hints the kernel to fault the page-aligned subrange of `[ptr, ptr + len)`
353/// back in ahead of need: asynchronous swap-in, the swap-backed extent store's
354/// readahead mechanism. Contents are preserved. No-op outside Linux.
355pub(crate) fn willneed(ptr: *mut u8, len: usize) {
356 // SAFETY: as in `pageout`, a contents-preserving hint over a live
357 // mapping.
358 unsafe { sys::advise(ptr, len, sys::Advice::WillNeed) }
359}
360
361/// Opts the page-aligned subrange of `[ptr, ptr + len)` out of transparent
362/// huge pages, so reclaim advice over the range operates on base pages and
363/// never needs a folio split. Contents are preserved. No-op outside Linux.
364pub(crate) fn nohugepage(ptr: *mut u8, len: usize) {
365 // SAFETY: as in `pageout`, a contents-preserving hint over a live
366 // mapping.
367 unsafe { sys::advise(ptr, len, sys::Advice::NoHugePage) }
368}
369
370/// Excludes the page-aligned subrange of `[ptr, ptr + len)` from core dumps.
371/// Contents are preserved. No-op outside Linux.
372pub(crate) fn dontdump(ptr: *mut u8, len: usize) {
373 // SAFETY: as in `pageout`, a contents-preserving hint over a live
374 // mapping.
375 unsafe { sys::advise(ptr, len, sys::Advice::DontDump) }
376}
377
378/// The system page size.
379pub(crate) fn page_size() -> usize {
380 sys::page_size()
381}
382
383/// Whether every page of the page-aligned subrange of `[ptr, ptr + len)`
384/// has been unmapped from this process, per the pagemap present bits: the
385/// observation the pageout ledger trusts instead of the reclaim advice's
386/// return value. A page unmapped to a swap entry counts as reclaimed even
387/// while its clean copy lingers in the kernel's swap cache. Errs toward
388/// `false` (resident) when the observation is unavailable. In test builds
389/// the answer comes from the `fake_residency` seam instead of the platform.
390pub(crate) fn nonresident(ptr: *mut u8, len: usize) -> bool {
391 #[cfg(test)]
392 {
393 let _ = (ptr, len);
394 fake_residency::observe()
395 }
396 #[cfg(not(test))]
397 sys::nonresident(ptr, len)
398}
399
400/// Test seam over the pageout residency observation. The decline queue is
401/// thread-local because observation runs on whichever thread enforces the
402/// compressed cap, so tests drive enforcement inline on their own thread.
403#[cfg(test)]
404pub(crate) mod fake_residency {
405 use std::cell::Cell;
406
407 thread_local! {
408 static DECLINES: Cell<u64> = const { Cell::new(0) };
409 }
410
411 /// Makes the next `n` observations on this thread report pages still
412 /// resident, modeling a kernel that declined the reclaim advice.
413 /// Replaces any previously queued declines.
414 pub(crate) fn decline_next(n: u64) {
415 DECLINES.with(|d| d.set(n));
416 }
417
418 /// One observation: consumes a queued decline (reporting the range
419 /// still resident), or reports it fully nonresident.
420 pub(super) fn observe() -> bool {
421 DECLINES.with(|d| {
422 let n = d.get();
423 if n > 0 {
424 d.set(n - 1);
425 false
426 } else {
427 true
428 }
429 })
430 }
431}
432
433/// The largest `page`-aligned subrange of `[addr, addr + len)`, as a
434/// `(byte offset from addr, subrange length)` pair, or `None` when the range
435/// covers no whole page (including on address-space overflow). `page` must
436/// be a power of two.
437///
438/// Guarantees, relied on by [`sys::advise`] for pointer arithmetic and
439/// proved by the Kani harnesses: `offset <= len`, `offset + sub_len <= len`,
440/// and both `addr + offset` and `sub_len` are `page`-aligned.
441#[cfg_attr(miri, allow(dead_code))]
442fn aligned_subrange(addr: usize, len: usize, page: usize) -> Option<(usize, usize)> {
443 debug_assert!(page.is_power_of_two());
444 let start = addr.checked_add(page - 1)? & !(page - 1);
445 let end = addr.checked_add(len)? & !(page - 1);
446 (start < end).then(|| (start - addr, end - start))
447}
448
449/// Splits an over-mapped range of `map_len` bytes at `addr` into
450/// `(head, tail)` trim amounts such that discarding `head` bytes from the
451/// front and `tail` from the back leaves an `align`-aligned range of exactly
452/// `len` bytes. `None` on address-space overflow or when the range cannot
453/// fit an aligned `len` bytes.
454///
455/// When `addr` and `len` are page-aligned and `align` is a page-multiple
456/// power of two, `head` and `tail` are page-aligned (so both trims are
457/// unmappable) — proved by the Kani harnesses.
458#[cfg_attr(miri, allow(dead_code))]
459fn align_trim(addr: usize, map_len: usize, len: usize, align: usize) -> Option<(usize, usize)> {
460 debug_assert!(align.is_power_of_two());
461 let aligned = addr.checked_next_multiple_of(align)?;
462 let head = aligned - addr;
463 let tail = map_len.checked_sub(head.checked_add(len)?)?;
464 Some((head, tail))
465}
466
467/// The platform seam: mapping, unmapping, paging advice, and the page size.
468///
469/// The `mmap` variant is production; the Miri variant backs regions with the
470/// Rust heap and treats every advice as a contents-preserving no-op — the
471/// weakest behavior the advice contracts allow — so the pool's tests run
472/// under the interpreter with full provenance and data-race checking:
473///
474/// ```text
475/// MIRIFLAGS=-Zmiri-disable-isolation cargo +nightly miri test -p mz-ore --features pool pool::
476/// ```
477///
478/// (The isolation flag is for the test harness's wall-clock log timestamps,
479/// not for anything the pool does.)
480#[cfg(not(miri))]
481mod sys {
482 use std::io;
483
484 use super::{align_trim, aligned_subrange};
485
486 /// Paging advice, in the vocabulary the pool needs.
487 pub(super) enum Advice {
488 /// Release physical pages; contents become undefined.
489 DontNeed,
490 /// Reclaim to the swap device now; contents preserved.
491 PageOut,
492 /// Fault back in ahead of need; contents preserved.
493 WillNeed,
494 /// Exclude from transparent huge pages; contents preserved.
495 NoHugePage,
496 /// Exclude from core dumps; contents preserved.
497 DontDump,
498 }
499
500 /// Maps `len` bytes of anonymous memory with the base aligned to
501 /// `align`. `len` must be a whole number of pages and `align` a
502 /// page-multiple power of two. Alignments beyond one page over-map and
503 /// trim; huge-page alignments additionally advise `MADV_HUGEPAGE`
504 /// (best-effort; the kernel falls back to base pages under
505 /// fragmentation).
506 pub(super) fn map(len: usize, align: usize) -> io::Result<*mut u8> {
507 let page = page_size();
508 debug_assert!(len % page == 0 && align % page == 0);
509 #[cfg(target_os = "linux")]
510 let flags = libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_NORESERVE;
511 #[cfg(not(target_os = "linux"))]
512 let flags = libc::MAP_PRIVATE | libc::MAP_ANONYMOUS;
513
514 let map_len = if align > page { len + align } else { len };
515 // SAFETY: anonymous mapping with a null hint; the kernel picks a
516 // fresh range that aliases no existing Rust object. `map_len` is
517 // positive and page-aligned by construction.
518 let raw = unsafe {
519 libc::mmap(
520 std::ptr::null_mut(),
521 map_len,
522 libc::PROT_READ | libc::PROT_WRITE,
523 flags,
524 -1,
525 0,
526 )
527 };
528 if raw == libc::MAP_FAILED {
529 return Err(io::Error::last_os_error());
530 }
531 let raw = raw.cast::<u8>();
532 if align <= page {
533 return Ok(raw);
534 }
535
536 // Trim the over-mapped head and tail so the base is aligned and the
537 // region owns exactly `len` bytes; `unmap` releases that range.
538 let Some((head, tail)) = align_trim(raw.addr(), map_len, len, align) else {
539 // Address-space arithmetic overflowed; treat the reservation as
540 // failed rather than keep an unaligned mapping.
541 // SAFETY: unmapping the mapping created above, in full.
542 unsafe { libc::munmap(raw.cast::<libc::c_void>(), map_len) };
543 return Err(io::Error::from(io::ErrorKind::OutOfMemory));
544 };
545 // SAFETY: `head` and `tail` are page-aligned subranges of the
546 // mapping just created (`align_trim`'s contract with page-aligned
547 // inputs), disjoint from the `len` bytes the region keeps; nothing
548 // references them.
549 unsafe {
550 if head > 0 {
551 libc::munmap(raw.cast::<libc::c_void>(), head);
552 }
553 if tail > 0 {
554 libc::munmap(raw.add(head + len).cast::<libc::c_void>(), tail);
555 }
556 }
557 // SAFETY: `head` stays within the original mapping.
558 let base = unsafe { raw.add(head) };
559
560 #[cfg(target_os = "linux")]
561 if align >= super::HUGE_PAGE {
562 // SAFETY: `base`/`len` describe the live aligned mapping; the
563 // advice is a non-destructive hint and failure is ignorable.
564 unsafe {
565 libc::madvise(base.cast::<libc::c_void>(), len, libc::MADV_HUGEPAGE);
566 }
567 }
568 Ok(base)
569 }
570
571 /// Releases a mapping returned by [`map`].
572 ///
573 /// # Safety
574 ///
575 /// `ptr`, `len`, and `align` must describe exactly one prior [`map`]
576 /// result, with no outstanding references into the range.
577 pub(super) unsafe fn unmap(ptr: *mut u8, len: usize, _align: usize) {
578 // SAFETY: per the function contract.
579 unsafe {
580 libc::munmap(ptr.cast::<libc::c_void>(), len);
581 }
582 }
583
584 /// Applies `advice` to the largest page-aligned subrange of
585 /// `[ptr, ptr + len)`, rounding the start up and the end down so the
586 /// advice never spills onto pages the range only partially covers.
587 ///
588 /// # Safety
589 ///
590 /// The range must lie within a live mapping. For [`Advice::DontNeed`]
591 /// the caller must additionally uphold the exclusivity contract
592 /// documented on [`super::dontneed`]; the remaining advice values are
593 /// non-mutating hints.
594 pub(super) unsafe fn advise(ptr: *mut u8, len: usize, advice: Advice) {
595 let advice = match advice {
596 Advice::DontNeed => libc::MADV_DONTNEED,
597 #[cfg(target_os = "linux")]
598 Advice::PageOut => libc::MADV_PAGEOUT,
599 #[cfg(target_os = "linux")]
600 Advice::WillNeed => libc::MADV_WILLNEED,
601 #[cfg(target_os = "linux")]
602 Advice::NoHugePage => libc::MADV_NOHUGEPAGE,
603 #[cfg(target_os = "linux")]
604 Advice::DontDump => libc::MADV_DONTDUMP,
605 // Reclaim, prefetch, THP, and dump hints have no portable
606 // equivalent.
607 #[cfg(not(target_os = "linux"))]
608 Advice::PageOut | Advice::WillNeed | Advice::NoHugePage | Advice::DontDump => return,
609 };
610 let Some((offset, sub_len)) = aligned_subrange(ptr.addr(), len, page_size()) else {
611 return;
612 };
613 // SAFETY: `offset <= len` (`aligned_subrange`'s contract), so the
614 // add stays within the caller's range and preserves provenance.
615 let aligned = unsafe { ptr.byte_add(offset) }.cast::<libc::c_void>();
616 // SAFETY: pointer and length describe a fully page-aligned subrange
617 // of the caller's live mapping; destructive advice is covered by the
618 // function contract.
619 unsafe {
620 libc::madvise(aligned, sub_len, advice);
621 }
622 }
623
624 /// Whether every page of the page-aligned subrange of `[ptr, ptr + len)`
625 /// has been unmapped from this process, per the present bits in
626 /// `/proc/self/pagemap`. A failed observation reports `false`: the
627 /// ledger keeps counting pages it cannot prove gone.
628 ///
629 /// The present bit is the signal, deliberately not `mincore(2)`:
630 /// successful asynchronous reclaim unmaps the PTE to a swap entry while
631 /// the page's clean copy lingers in the kernel's swap cache, and mincore
632 /// reports swap-cache pages as in core. Observing in-core-ness would
633 /// therefore misclassify essentially every successful pageout on an
634 /// unpressured host (the retry advice cannot change the outcome either,
635 /// since the advice skips already-unmapped PTEs). An unmapped page is
636 /// reclaimed for this ledger's purposes: a clean swap-cache copy is
637 /// memory the kernel drops for free.
638 ///
639 /// Only Linux answers from the kernel. Elsewhere the reclaim advice is
640 /// compiled out, so there is no reclaim to observe and the answer is
641 /// `true`, keeping the compressed tier cycling on development platforms.
642 #[cfg_attr(test, allow(dead_code))]
643 pub(super) fn nonresident(ptr: *mut u8, len: usize) -> bool {
644 #[cfg(target_os = "linux")]
645 {
646 use std::cell::RefCell;
647 use std::os::unix::fs::FileExt;
648 use std::sync::OnceLock;
649
650 // Reading a process's own pagemap needs no privilege since
651 // Linux 4.2: unprivileged readers see zeroed frame numbers,
652 // which this probe never looks at.
653 static PAGEMAP: OnceLock<Option<std::fs::File>> = OnceLock::new();
654 let Some(pagemap) = PAGEMAP
655 .get_or_init(|| std::fs::File::open("/proc/self/pagemap").ok())
656 .as_ref()
657 else {
658 return false;
659 };
660 let page = page_size();
661 let Some((offset, sub_len)) = aligned_subrange(ptr.addr(), len, page) else {
662 // No whole page to observe: vacuously nonresident.
663 return true;
664 };
665 let first_page = (ptr.addr() + offset) / page;
666 let pages = sub_len / page;
667 thread_local! {
668 static SCRATCH: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
669 }
670 SCRATCH.with(|scratch| {
671 let mut buf = scratch.borrow_mut();
672 buf.clear();
673 buf.resize(pages, 0);
674 let file_offset = u64::try_from(first_page).expect("page index fits u64") * 8;
675 if pagemap
676 .read_exact_at(bytemuck::cast_slice_mut(buf.as_mut_slice()), file_offset)
677 .is_err()
678 {
679 return false;
680 }
681 // One native-endian u64 per page, per the kernel ABI
682 // (Documentation/admin-guide/mm/pagemap.rst); bit 63 is
683 // "present in RAM". A swap-entry PTE (bit 62) and a
684 // never-faulted zero entry are both out of memory.
685 buf.iter().all(|entry| entry & (1 << 63) == 0)
686 })
687 }
688 #[cfg(not(target_os = "linux"))]
689 {
690 let _ = (ptr, len);
691 true
692 }
693 }
694
695 pub(super) fn page_size() -> usize {
696 // SAFETY: `sysconf` with a valid argument is safe.
697 let raw = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
698 let page = usize::try_from(raw).expect("page size is positive and fits usize");
699 // The alignment arithmetic (`aligned_subrange`, `align_trim`) masks
700 // with `page - 1`; a non-power-of-two page would mis-round silently.
701 assert!(page.is_power_of_two(), "page size is a power of two");
702 page
703 }
704}
705
706#[cfg(miri)]
707mod sys {
708 use std::alloc::Layout;
709 use std::io;
710
711 pub(super) enum Advice {
712 DontNeed,
713 PageOut,
714 WillNeed,
715 NoHugePage,
716 DontDump,
717 }
718
719 pub(super) fn map(len: usize, align: usize) -> io::Result<*mut u8> {
720 let layout = Layout::from_size_align(len, align).expect("valid region layout");
721 // SAFETY: `len` is positive (empty regions never map). The memory is
722 // deliberately left uninitialized, matching the slot contract that
723 // contents are unspecified until fully overwritten; Miri enforces
724 // that no path reads a byte before writing it.
725 let ptr = unsafe { std::alloc::alloc(layout) };
726 if ptr.is_null() {
727 std::alloc::handle_alloc_error(layout);
728 }
729 Ok(ptr)
730 }
731
732 pub(super) unsafe fn unmap(ptr: *mut u8, len: usize, align: usize) {
733 let layout = Layout::from_size_align(len, align).expect("valid region layout");
734 // SAFETY: `ptr` was returned by `map` with exactly this layout.
735 unsafe { std::alloc::dealloc(ptr, layout) }
736 }
737
738 /// Advice is a contents-preserving no-op: the weakest behavior the
739 /// contracts allow (`DontNeed` leaves contents undefined, and "old bytes
740 /// kept" is one of the permitted outcomes, as on macOS).
741 pub(super) unsafe fn advise(_ptr: *mut u8, _len: usize, _advice: Advice) {}
742
743 /// The heap backing has no residency to observe, so reclaim advice is
744 /// treated as fully effective, mirroring the non-Linux answer.
745 #[cfg_attr(test, allow(dead_code))]
746 pub(super) fn nonresident(_ptr: *mut u8, _len: usize) -> bool {
747 true
748 }
749
750 pub(super) fn page_size() -> usize {
751 4096
752 }
753}
754
755/// Kani proof harnesses over the pure arithmetic and the slot allocator:
756/// the sequential facts the pool's unsafe blocks rest on. Run
757/// `cargo kani --features pool` from `src/ore`; each harness is exhaustive
758/// over its symbolic inputs. Kani ships its own toolchain, so it needs a
759/// release whose rustc is at or above the workspace `rust-version` (or the
760/// manifest check bypassed for the run). Concurrency-justified claims (the
761/// lock-protocol aliasing arguments in `pool.rs`) are outside Kani's
762/// sequential model and are exercised under Miri instead.
763#[cfg(kani)]
764mod proofs {
765 use super::*;
766
767 /// `aligned_subrange`'s contract: the subrange lies within the input
768 /// range, is nonempty, and is page-aligned at both ends.
769 #[kani::proof]
770 fn aligned_subrange_stays_in_bounds() {
771 let addr: usize = kani::any();
772 let len: usize = kani::any();
773 let shift: u32 = kani::any();
774 kani::assume(shift < usize::BITS);
775 let page = 1usize << shift;
776 if let Some((offset, sub_len)) = aligned_subrange(addr, len, page) {
777 assert!(offset <= len);
778 assert!(sub_len <= len - offset);
779 assert!(sub_len > 0);
780 assert!((addr + offset) % page == 0);
781 assert!(sub_len % page == 0);
782 }
783 }
784
785 /// `align_trim`'s contract: head and tail partition the over-map
786 /// exactly around an aligned range of the requested length, and with
787 /// page-aligned inputs both trims are page-aligned, so each can be
788 /// unmapped independently.
789 #[kani::proof]
790 fn align_trim_partitions_the_overmap() {
791 let addr: usize = kani::any();
792 let len: usize = kani::any();
793 let page_shift: u32 = kani::any();
794 let align_shift: u32 = kani::any();
795 kani::assume(page_shift <= align_shift && align_shift < usize::BITS - 1);
796 let page = 1usize << page_shift;
797 let align = 1usize << align_shift;
798 kani::assume(addr % page == 0);
799 kani::assume(len % page == 0);
800 let Some(map_len) = len.checked_add(align) else {
801 return;
802 };
803 if let Some((head, tail)) = align_trim(addr, map_len, len, align) {
804 assert_eq!(head + len + tail, map_len);
805 assert!(head < align);
806 assert!(head % page == 0);
807 assert!(tail % page == 0);
808 assert!((addr + head) % align == 0);
809 }
810 }
811
812 /// The slot allocator, over every alloc/free sequence of a bounded
813 /// length: an allocated slot is always in range and never aliases a
814 /// live one. This is the disjointness fact `slot_ptr`'s callers turn
815 /// into non-aliasing slices.
816 #[kani::proof]
817 #[kani::unwind(8)]
818 fn slot_allocator_hands_out_disjoint_slots() {
819 const MAX: u32 = 3;
820 let mut slots = SlotAllocator::new(MAX);
821 let mut live = [false; MAX as usize];
822 for _ in 0..5 {
823 if kani::any() {
824 if let Some((slot, _warm)) = slots.alloc() {
825 let slot = usize::try_from(slot).expect("fits");
826 assert!(slot < live.len(), "slot out of range");
827 assert!(!live[slot], "live slot handed out twice");
828 live[slot] = true;
829 }
830 } else {
831 let slot: usize = kani::any();
832 kani::assume(slot < live.len());
833 if live[slot] {
834 live[slot] = false;
835 slots.free(u32::try_from(slot).expect("fits"), kani::any());
836 }
837 }
838 }
839 }
840
841 /// The class-selection contract the pool turns into slice bounds: a
842 /// selected class always fits the payload.
843 #[kani::proof]
844 #[kani::unwind(10)]
845 fn selected_class_fits_payload() {
846 let len_bytes: usize = kani::any();
847 if let Some(class) = size_class_for(len_bytes) {
848 assert!(class < SIZE_CLASSES.len());
849 assert!(SIZE_CLASSES[class] >= len_bytes);
850 } else {
851 assert!(len_bytes > SIZE_CLASSES[SIZE_CLASSES.len() - 1]);
852 }
853 }
854
855 /// The offset arithmetic behind `slot_ptr`: for every size class and
856 /// every slot index the allocator can hand out, the slot lies within
857 /// the region's capacity.
858 #[kani::proof]
859 #[kani::unwind(9)]
860 fn slot_offsets_lie_within_capacity() {
861 for &class_size in &SIZE_CLASSES {
862 let capacity_bytes: usize = kani::any();
863 kani::assume(capacity_bytes <= 1 << 40);
864 let capacity = capacity_bytes - capacity_bytes % class_size;
865 let max_slots = capacity / class_size;
866 let slot: usize = kani::any();
867 kani::assume(slot < max_slots);
868 let offset = slot * class_size;
869 assert!(offset + class_size <= capacity);
870 }
871 }
872}
873
874#[cfg(test)]
875mod tests {
876 use super::*;
877
878 #[mz_ore::test]
879 fn alloc_free_reuses_slots() {
880 let region = Region::new(64 << 10, 1 << 20).expect("mmap");
881 let (a, warm_a) = region.alloc().expect("slot");
882 let (b, _) = region.alloc().expect("slot");
883 assert!(!warm_a, "bump slots are not warm");
884 assert_ne!(a, b);
885 assert_ne!(region.slot_ptr(a), region.slot_ptr(b));
886 let ptr_a = region.slot_ptr(a);
887 // A warm free is preferred by the next alloc and reported warm.
888 region.free(a, true);
889 let (c, warm_c) = region.alloc().expect("slot");
890 assert_eq!(c, a);
891 assert!(warm_c);
892 assert_eq!(region.slot_ptr(c), ptr_a);
893 // A cold free comes back, but not warm.
894 region.free(c, false);
895 let (d, warm_d) = region.alloc().expect("slot");
896 assert_eq!(d, a);
897 assert!(!warm_d);
898 }
899
900 /// Hugepage-class regions get a huge-page-aligned base, so slots tile
901 /// huge-page boundaries exactly.
902 #[mz_ore::test]
903 fn hugepage_class_base_is_aligned() {
904 let region = Region::new(2 << 20, 16 << 20).expect("mmap");
905 let (slot, _) = region.alloc().expect("slot");
906 if cfg!(target_os = "linux") {
907 assert_eq!(
908 region.slot_ptr(slot).addr() % HUGE_PAGE,
909 0,
910 "hugepage-class slots must be huge-page aligned",
911 );
912 }
913 }
914
915 #[mz_ore::test]
916 fn exhaustion_returns_none() {
917 let region = Region::new(64 << 10, 128 << 10).expect("mmap");
918 assert!(region.alloc().is_some());
919 assert!(region.alloc().is_some());
920 assert!(region.alloc().is_none(), "third slot exceeds capacity");
921 }
922
923 #[mz_ore::test]
924 fn slots_are_writable_and_advice_is_accepted() {
925 let region = Region::new(64 << 10, 1 << 20).expect("mmap");
926 let (slot, _) = region.alloc().expect("slot");
927 let ptr = region.slot_ptr(slot);
928 // SAFETY: freshly allocated slot, exclusively owned by this test.
929 unsafe {
930 std::ptr::write_bytes(ptr, 0xAB, region.class_size());
931 }
932 pageout(ptr, region.class_size());
933 willneed(ptr, region.class_size());
934 // SAFETY: the slot is exclusively owned and is not read again before
935 // being overwritten (it is not read again at all).
936 unsafe {
937 dontneed(ptr, region.class_size());
938 }
939 }
940
941 #[mz_ore::test]
942 fn aligned_subrange_agrees_with_examples() {
943 // Fully aligned range: identity.
944 assert_eq!(aligned_subrange(4096, 8192, 4096), Some((0, 8192)));
945 // Unaligned start rounds up, unaligned end rounds down.
946 assert_eq!(aligned_subrange(4097, 8192, 4096), Some((4095, 4096)));
947 // Too short to cover a whole page.
948 assert_eq!(aligned_subrange(4097, 4096, 4096), None);
949 assert_eq!(aligned_subrange(0, 0, 4096), None);
950 }
951
952 #[mz_ore::test]
953 fn align_trim_agrees_with_examples() {
954 // Already aligned: no head, tail is the whole over-map.
955 assert_eq!(
956 align_trim(HUGE_PAGE, 3 * HUGE_PAGE, 2 * HUGE_PAGE, HUGE_PAGE),
957 Some((0, HUGE_PAGE))
958 );
959 // Unaligned base: head consumes the misalignment.
960 assert_eq!(
961 align_trim(HUGE_PAGE + 4096, 3 * HUGE_PAGE, 2 * HUGE_PAGE, HUGE_PAGE),
962 Some((HUGE_PAGE - 4096, 4096)),
963 );
964 }
965}