Skip to main content

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