hermes_core/segment/
pin.rs

1//! Hot-metadata pinning: budgeted residency for per-query-mandatory
2//! structures (a meta/data residency split).
3//!
4//! Every query must touch certain small metadata sections — BMP block-offset
5//! tables, sparse skip sections, doc-id maps, and the coarse BMP hierarchy. Under memory
6//! pressure the kernel evicts them like bulk data, and queries then pay major
7//! faults on structures they cannot skip. This module pins them, in priority
8//! order (smallest/hottest first), until a per-segment budget is exhausted.
9//!
10//! Design: `docs/hot-metadata-pinning.md`. Bulk data (BMP 4-bit grid, block
11//! data, raw vectors) is never pinned — it is covered by the
12//! `MADV_RANDOM`/`MADV_WILLNEED` discipline instead.
13
14use std::sync::{Arc, OnceLock};
15
16use crate::directories::OwnedBytes;
17
18/// How pinned bytes are kept resident.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum PinMode {
21    /// `mlock` the mmap pages in place — zero-copy, but requires
22    /// RLIMIT_MEMLOCK headroom (containers often need CAP_IPC_LOCK or an
23    /// explicit ulimit). Failures are logged and counted, never fatal.
24    Mlock,
25    /// Copy the section to the heap — no permissions needed, duplicates the
26    /// bytes. Immune to page-cache eviction (production runs swapless).
27    Copy,
28}
29
30/// Per-segment metadata pinning policy.
31#[derive(Debug, Clone, Copy)]
32pub struct PinPolicy {
33    /// Metadata bytes to pin per segment. 0 = pinning disabled (default).
34    pub budget_bytes: u64,
35    pub mode: PinMode,
36}
37
38impl PinPolicy {
39    pub const fn disabled() -> Self {
40        Self {
41            budget_bytes: 0,
42            mode: PinMode::Mlock,
43        }
44    }
45
46    pub fn is_enabled(&self) -> bool {
47        self.budget_bytes > 0
48    }
49
50    /// Read policy from environment:
51    /// `HERMES_PIN_METADATA_BUDGET_MB` (default 0 = off),
52    /// `HERMES_PIN_MODE` = `mlock` (default) | `copy`.
53    ///
54    /// Used as the default initializer for [`pin_policy`]. Consumers such as
55    /// hermes-server expose CLI flags and honor these environment values when
56    /// the corresponding flags are unset.
57    pub fn from_env() -> Self {
58        let budget_mb: u64 = std::env::var("HERMES_PIN_METADATA_BUDGET_MB")
59            .ok()
60            .and_then(|v| v.parse().ok())
61            .unwrap_or(0);
62        let mode = match std::env::var("HERMES_PIN_MODE").as_deref() {
63            Ok("copy") => PinMode::Copy,
64            Ok("mlock") | Err(_) => PinMode::Mlock,
65            Ok(other) => {
66                log::warn!("HERMES_PIN_MODE '{}' unknown; using mlock", other);
67                PinMode::Mlock
68            }
69        };
70        Self {
71            budget_bytes: budget_mb * 1024 * 1024,
72            mode,
73        }
74    }
75}
76
77static PIN_POLICY: OnceLock<PinPolicy> = OnceLock::new();
78
79/// Override the process-wide pin policy. Must be called before the first
80/// segment is opened; returns false (and warns) if the policy was already
81/// initialized.
82pub fn set_pin_policy(policy: PinPolicy) -> bool {
83    let ok = PIN_POLICY.set(policy).is_ok();
84    if !ok {
85        log::warn!("pin policy already initialized; set_pin_policy ignored");
86    }
87    ok
88}
89
90/// The process-wide pin policy (env-initialized on first use).
91pub fn pin_policy() -> &'static PinPolicy {
92    PIN_POLICY.get_or_init(PinPolicy::from_env)
93}
94
95/// Accumulates pin accounting for one segment.
96#[derive(Debug, Default, Clone, Copy)]
97pub struct PinReport {
98    /// Bytes of pinnable metadata found (regardless of budget/failures)
99    pub intended_bytes: u64,
100    /// Bytes actually pinned
101    pub pinned_bytes: u64,
102    /// Bytes skipped because the budget was exhausted
103    pub skipped_budget_bytes: u64,
104    /// Bytes where mlock failed (RLIMIT_MEMLOCK etc.)
105    pub failed_bytes: u64,
106    /// Additional heap allocated by `PinMode::Copy`. Already-heap ANN routing
107    /// structures are resident but do not contribute here.
108    pub heap_copy_bytes: u64,
109}
110
111/// RAII owner for heap pages locked on behalf of one immutable ANN artifact
112/// generation. The referenced allocations are owned by the same
113/// `TrainedVectorStructures`; its field order drops this set before the
114/// artifact `Arc`s, so every address remains valid through `munlock`.
115struct HeapPinGuard {
116    page_start: *mut libc::c_void,
117    page_len: usize,
118}
119
120// The guard never dereferences its pointer. The immutable artifact allocations
121// it describes are safe to share, and mlock/munlock operate on process mappings.
122unsafe impl Send for HeapPinGuard {}
123unsafe impl Sync for HeapPinGuard {}
124
125impl Drop for HeapPinGuard {
126    fn drop(&mut self) {
127        if unsafe { libc::munlock(self.page_start, self.page_len) } != 0 {
128            log::warn!(
129                "[pin] munlock failed for {} of ANN heap: {}",
130                crate::format_bytes(self.page_len as u64),
131                std::io::Error::last_os_error()
132            );
133        }
134    }
135}
136
137/// Locked heap allocations associated with one index-global ANN generation.
138/// Segment-local vector/code payloads are intentionally excluded.
139#[derive(Default)]
140pub(crate) struct HeapPinSet {
141    guards: Vec<HeapPinGuard>,
142    /// Keep every allocation owner alive until after its guards are dropped,
143    /// even if a cloned `TrainedVectorStructures` has its public maps mutated.
144    owners: Vec<Arc<dyn std::any::Any + Send + Sync>>,
145    report: PinReport,
146}
147
148impl HeapPinSet {
149    pub(crate) fn report(&self) -> PinReport {
150        self.report
151    }
152
153    pub(crate) fn retain_owner<T: std::any::Any + Send + Sync>(&mut self, owner: Arc<T>) {
154        self.owners.push(owner);
155    }
156
157    /// Keep one immutable heap slice resident, subject to the generation
158    /// budget. `Copy` mode needs no allocation: trained artifacts are already
159    /// heap-owned, which is exactly the residency guarantee that mode provides
160    /// on the supported swapless deployment.
161    pub(crate) fn pin_slice<T>(
162        &mut self,
163        slice: &[T],
164        label: &str,
165        mode: PinMode,
166        remaining: &mut u64,
167    ) {
168        let len = std::mem::size_of_val(slice);
169        if len == 0 {
170            return;
171        }
172        let Ok(len_u64) = u64::try_from(len) else {
173            self.report.failed_bytes = u64::MAX;
174            log::warn!("[pin] ANN region {label} is too large to account");
175            return;
176        };
177        self.report.intended_bytes = self.report.intended_bytes.saturating_add(len_u64);
178        if len_u64 > *remaining {
179            self.report.skipped_budget_bytes =
180                self.report.skipped_budget_bytes.saturating_add(len_u64);
181            log::debug!(
182                "[pin] ANN budget exhausted: skipping {} ({}, {} remaining)",
183                label,
184                crate::format_bytes(len_u64),
185                crate::format_bytes(*remaining)
186            );
187            return;
188        }
189
190        if mode == PinMode::Copy {
191            *remaining -= len_u64;
192            self.report.pinned_bytes = self.report.pinned_bytes.saturating_add(len_u64);
193            return;
194        }
195
196        let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
197        let page_size = usize::try_from(page_size).ok().filter(|&size| size > 0);
198        let Some(page_size) = page_size else {
199            self.report.failed_bytes = self.report.failed_bytes.saturating_add(len_u64);
200            log::warn!("[pin] cannot determine page size while locking {label}");
201            return;
202        };
203        let address = slice.as_ptr() as usize;
204        let page_start = address / page_size * page_size;
205        let Some(end) = address.checked_add(len) else {
206            self.report.failed_bytes = self.report.failed_bytes.saturating_add(len_u64);
207            log::warn!("[pin] ANN region address overflow while locking {label}");
208            return;
209        };
210        let Some(rounded_end) = end
211            .checked_add(page_size - 1)
212            .map(|value| value / page_size * page_size)
213        else {
214            self.report.failed_bytes = self.report.failed_bytes.saturating_add(len_u64);
215            log::warn!("[pin] ANN region page range overflow while locking {label}");
216            return;
217        };
218        let page_len = rounded_end - page_start;
219        let page_start = page_start as *mut libc::c_void;
220        if unsafe { libc::mlock(page_start.cast_const(), page_len) } == 0 {
221            self.guards.push(HeapPinGuard {
222                page_start,
223                page_len,
224            });
225            *remaining -= len_u64;
226            self.report.pinned_bytes = self.report.pinned_bytes.saturating_add(len_u64);
227        } else {
228            self.report.failed_bytes = self.report.failed_bytes.saturating_add(len_u64);
229            log::warn!(
230                "[pin] mlock failed for ANN {} ({}): {} — check RLIMIT_MEMLOCK/CAP_IPC_LOCK; continuing unpinned",
231                label,
232                crate::format_bytes(len_u64),
233                std::io::Error::last_os_error()
234            );
235        }
236    }
237}
238
239/// Pin one metadata section, updating `remaining` budget and the report.
240///
241/// In `Copy` mode the section is replaced with a heap copy (heap memory is
242/// not page-cache-evictable). In `Mlock` mode the mmap pages are locked in
243/// place. Non-mmap-backed sections (RAM directories) are already resident
244/// and are skipped silently.
245pub(crate) fn pin_section(
246    bytes: &mut OwnedBytes,
247    label: &str,
248    mode: PinMode,
249    remaining: &mut u64,
250    report: &mut PinReport,
251) {
252    if !bytes.is_mmap() || bytes.is_empty() {
253        return;
254    }
255    let len = bytes.len() as u64;
256    report.intended_bytes += len;
257
258    if len > *remaining {
259        report.skipped_budget_bytes += len;
260        log::debug!(
261            "[pin] budget exhausted: skipping {} ({}, {} remaining)",
262            label,
263            crate::format_bytes(len),
264            crate::format_bytes(*remaining)
265        );
266        return;
267    }
268
269    match mode {
270        PinMode::Mlock => {
271            if bytes.mlock() {
272                *remaining -= len;
273                report.pinned_bytes += len;
274            } else {
275                report.failed_bytes += len;
276                log::warn!(
277                    "[pin] mlock failed for {} ({}) — check RLIMIT_MEMLOCK; \
278                     continuing unpinned",
279                    label,
280                    crate::format_bytes(len)
281                );
282            }
283        }
284        PinMode::Copy => {
285            *bytes = OwnedBytes::new(bytes.to_vec());
286            *remaining -= len;
287            report.pinned_bytes += len;
288            report.heap_copy_bytes += len;
289        }
290    }
291}
hermes_core/segment/pin.rs

hermes_core/segment/
pin.rs
