summa_core/segment/
pin.rs

1//! Hot-metadata pinning: budgeted residency for per-query-mandatory
2//! structures (a meta/data residency split).
3//!
4//! Queries touch compact ANN directories, dense document maps, BMP block and
5//! hierarchy directories, and MaxScore skip metadata.
6//! Under memory pressure, budgeted pinning keeps eligible metadata resident.
7//! Seismic's compact term and logical-row directories are eligible; its
8//! forward vectors, summaries, and nomination rows stay evictable.
9//!
10//! Design: `docs/hot-metadata-pinning.md`. Corpus-sized vectors and posting
11//! payloads are never pinned by this policy.
12
13use std::sync::{Arc, OnceLock};
14
15use crate::directories::OwnedBytes;
16
17/// How pinned bytes are kept resident.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum PinMode {
20    /// `mlock` the mmap pages in place — zero-copy, but requires
21    /// RLIMIT_MEMLOCK headroom (containers often need CAP_IPC_LOCK or an
22    /// explicit ulimit). Failures are logged and counted, never fatal.
23    Mlock,
24    /// Copy the section to the heap — no permissions needed, duplicates the
25    /// bytes. Immune to page-cache eviction (production runs swapless).
26    Copy,
27}
28
29/// Per-segment metadata pinning policy.
30#[derive(Debug, Clone, Copy)]
31pub struct PinPolicy {
32    /// Metadata bytes to pin per segment. 0 = pinning disabled (default).
33    pub budget_bytes: u64,
34    pub mode: PinMode,
35}
36
37impl PinPolicy {
38    pub const fn disabled() -> Self {
39        Self {
40            budget_bytes: 0,
41            mode: PinMode::Mlock,
42        }
43    }
44
45    pub fn is_enabled(&self) -> bool {
46        self.budget_bytes > 0
47    }
48
49    /// Read policy from environment:
50    /// `SUMMA_PIN_METADATA_BUDGET_MB` (default 0 = off),
51    /// `SUMMA_PIN_MODE` = `mlock` (default) | `copy`.
52    ///
53    /// Used as the default initializer for [`pin_policy`]. Consumers such as
54    /// summa-server expose CLI flags and honor these environment values when
55    /// the corresponding flags are unset.
56    pub fn from_env() -> Self {
57        let budget_bytes = match std::env::var("SUMMA_PIN_METADATA_BUDGET_MB") {
58            Ok(value) => match value
59                .parse::<u64>()
60                .ok()
61                .and_then(|mb| mb.checked_mul(1024 * 1024))
62            {
63                Some(bytes) => bytes,
64                None => {
65                    log::warn!(
66                        "SUMMA_PIN_METADATA_BUDGET_MB '{}' is invalid or too large; \
67                         using budget 0 (pinning disabled); set a non-negative MiB value \
68                         no larger than {}",
69                        value,
70                        u64::MAX / (1024 * 1024),
71                    );
72                    0
73                }
74            },
75            Err(std::env::VarError::NotPresent) => 0,
76            Err(error) => {
77                log::warn!(
78                    "SUMMA_PIN_METADATA_BUDGET_MB cannot be read: {}; pinning disabled",
79                    error
80                );
81                0
82            }
83        };
84        let mode = match std::env::var("SUMMA_PIN_MODE").as_deref() {
85            Ok("copy") => PinMode::Copy,
86            Ok("mlock") | Err(_) => PinMode::Mlock,
87            Ok(other) => {
88                log::warn!("SUMMA_PIN_MODE '{}' unknown; using mlock", other);
89                PinMode::Mlock
90            }
91        };
92        Self { budget_bytes, mode }
93    }
94}
95
96static PIN_POLICY: OnceLock<PinPolicy> = OnceLock::new();
97
98/// Hot-metadata size per segment above which running with pinning disabled
99/// (the default: `SUMMA_PIN_METADATA_BUDGET_MB` unset or 0) is worth one
100/// warning at open. Below this the kernel keeps the pages warm anyway.
101pub const UNPINNED_METADATA_WARN_THRESHOLD_BYTES: u64 = 16 * 1024 * 1024;
102
103static WARNED_PINNING_DISABLED: std::sync::atomic::AtomicBool =
104    std::sync::atomic::AtomicBool::new(false);
105
106/// Pinning is disabled and a segment carries `intended_bytes` of mmap-backed
107/// hot metadata that would have been pinned. Warns once per process (with the
108/// env var to set) when that exceeds
109/// [`UNPINNED_METADATA_WARN_THRESHOLD_BYTES`]; a silent zero default would
110/// otherwise look identical to a healthy configuration.
111pub(crate) fn warn_if_pinning_disabled(index_label: &str, segment_id: u128, intended_bytes: u64) {
112    if intended_bytes < UNPINNED_METADATA_WARN_THRESHOLD_BYTES {
113        return;
114    }
115    if WARNED_PINNING_DISABLED.swap(true, std::sync::atomic::Ordering::Relaxed) {
116        return;
117    }
118    log::warn!(
119        "[pin] index={} segment {:016x}: hot-metadata pinning is disabled (budget 0) but this \
120         segment has {} of mmap-backed per-query metadata; set SUMMA_PIN_METADATA_BUDGET_MB \
121         (and SUMMA_PIN_MODE=mlock|copy) to keep it resident under memory pressure \
122         (reported once per process)",
123        index_label,
124        segment_id,
125        crate::format_bytes(intended_bytes),
126    );
127}
128
129/// Override the process-wide pin policy. Must be called before the first
130/// segment is opened; returns false (and warns) if the policy was already
131/// initialized.
132pub fn set_pin_policy(policy: PinPolicy) -> bool {
133    let ok = PIN_POLICY.set(policy).is_ok();
134    if !ok {
135        log::warn!("pin policy already initialized; set_pin_policy ignored");
136    }
137    ok
138}
139
140/// The process-wide pin policy (env-initialized on first use).
141pub fn pin_policy() -> &'static PinPolicy {
142    PIN_POLICY.get_or_init(PinPolicy::from_env)
143}
144
145/// Accumulates pin accounting for one segment.
146#[derive(Debug, Default, Clone, Copy)]
147pub struct PinReport {
148    /// Bytes of pinnable metadata found (regardless of budget/failures)
149    pub intended_bytes: u64,
150    /// Bytes actually pinned
151    pub pinned_bytes: u64,
152    /// Bytes skipped because the budget was exhausted
153    pub skipped_budget_bytes: u64,
154    /// Bytes where mlock failed (RLIMIT_MEMLOCK etc.)
155    pub failed_bytes: u64,
156    /// Additional heap allocated by `PinMode::Copy`. Already-heap ANN routing
157    /// structures are resident but do not contribute here.
158    pub heap_copy_bytes: u64,
159}
160
161/// RAII owner for heap pages locked on behalf of one immutable ANN artifact
162/// generation. The referenced allocations are owned by the same
163/// `TrainedVectorStructures`; its field order drops this set before the
164/// artifact `Arc`s, so every address remains valid through `munlock`.
165struct HeapPinGuard {
166    page_start: *mut libc::c_void,
167    page_len: usize,
168}
169
170// The guard never dereferences its pointer. The immutable artifact allocations
171// it describes are safe to share, and mlock/munlock operate on process mappings.
172unsafe impl Send for HeapPinGuard {}
173unsafe impl Sync for HeapPinGuard {}
174
175impl Drop for HeapPinGuard {
176    fn drop(&mut self) {
177        if unsafe { libc::munlock(self.page_start, self.page_len) } != 0 {
178            log::warn!(
179                "[pin] munlock failed for {} of ANN heap: {}",
180                crate::format_bytes(self.page_len as u64),
181                std::io::Error::last_os_error()
182            );
183        }
184    }
185}
186
187/// Locked heap allocations associated with one index-global ANN generation.
188/// Segment-local vector/code payloads are intentionally excluded.
189#[derive(Default)]
190pub(crate) struct HeapPinSet {
191    guards: Vec<HeapPinGuard>,
192    /// Keep every allocation owner alive until after its guards are dropped,
193    /// even if a cloned `TrainedVectorStructures` has its public maps mutated.
194    owners: Vec<Arc<dyn std::any::Any + Send + Sync>>,
195    report: PinReport,
196}
197
198impl HeapPinSet {
199    pub(crate) fn report(&self) -> PinReport {
200        self.report
201    }
202
203    pub(crate) fn retain_owner<T: std::any::Any + Send + Sync>(&mut self, owner: Arc<T>) {
204        self.owners.push(owner);
205    }
206
207    /// Keep one immutable heap slice resident, subject to the generation
208    /// budget. `Copy` mode needs no allocation: trained artifacts are already
209    /// heap-owned, which is exactly the residency guarantee that mode provides
210    /// on the supported swapless deployment.
211    pub(crate) fn pin_slice<T>(
212        &mut self,
213        slice: &[T],
214        label: &str,
215        mode: PinMode,
216        remaining: &mut u64,
217    ) {
218        let len = std::mem::size_of_val(slice);
219        if len == 0 {
220            return;
221        }
222        let Ok(len_u64) = u64::try_from(len) else {
223            self.report.failed_bytes = u64::MAX;
224            log::warn!("[pin] ANN region {label} is too large to account");
225            return;
226        };
227        self.report.intended_bytes = self.report.intended_bytes.saturating_add(len_u64);
228        if len_u64 > *remaining {
229            self.report.skipped_budget_bytes =
230                self.report.skipped_budget_bytes.saturating_add(len_u64);
231            log::debug!(
232                "[pin] ANN budget exhausted: skipping {} ({}, {} remaining)",
233                label,
234                crate::format_bytes(len_u64),
235                crate::format_bytes(*remaining)
236            );
237            return;
238        }
239
240        if mode == PinMode::Copy {
241            *remaining -= len_u64;
242            self.report.pinned_bytes = self.report.pinned_bytes.saturating_add(len_u64);
243            return;
244        }
245
246        let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
247        let page_size = usize::try_from(page_size).ok().filter(|&size| size > 0);
248        let Some(page_size) = page_size else {
249            self.report.failed_bytes = self.report.failed_bytes.saturating_add(len_u64);
250            log::warn!("[pin] cannot determine page size while locking {label}");
251            return;
252        };
253        let address = slice.as_ptr() as usize;
254        let page_start = address / page_size * page_size;
255        let Some(end) = address.checked_add(len) else {
256            self.report.failed_bytes = self.report.failed_bytes.saturating_add(len_u64);
257            log::warn!("[pin] ANN region address overflow while locking {label}");
258            return;
259        };
260        let Some(rounded_end) = end
261            .checked_add(page_size - 1)
262            .map(|value| value / page_size * page_size)
263        else {
264            self.report.failed_bytes = self.report.failed_bytes.saturating_add(len_u64);
265            log::warn!("[pin] ANN region page range overflow while locking {label}");
266            return;
267        };
268        let page_len = rounded_end - page_start;
269        let page_start = page_start as *mut libc::c_void;
270        if unsafe { libc::mlock(page_start.cast_const(), page_len) } == 0 {
271            self.guards.push(HeapPinGuard {
272                page_start,
273                page_len,
274            });
275            *remaining -= len_u64;
276            self.report.pinned_bytes = self.report.pinned_bytes.saturating_add(len_u64);
277        } else {
278            self.report.failed_bytes = self.report.failed_bytes.saturating_add(len_u64);
279            log::warn!(
280                "[pin] mlock failed for ANN {} ({}): {} — check RLIMIT_MEMLOCK/CAP_IPC_LOCK; continuing unpinned",
281                label,
282                crate::format_bytes(len_u64),
283                std::io::Error::last_os_error()
284            );
285        }
286    }
287}
288
289/// Pin one metadata section, updating `remaining` budget and the report.
290///
291/// In `Copy` mode the section is replaced with a heap copy (heap memory is
292/// not page-cache-evictable). In `Mlock` mode the mmap pages are locked in
293/// place. Non-mmap-backed sections (RAM directories) are already resident
294/// and are skipped silently.
295pub(crate) fn pin_section(
296    bytes: &mut OwnedBytes,
297    label: &str,
298    mode: PinMode,
299    remaining: &mut u64,
300    report: &mut PinReport,
301) {
302    if !bytes.is_mmap() || bytes.is_empty() {
303        return;
304    }
305    let len = bytes.len() as u64;
306    report.intended_bytes += len;
307
308    if len > *remaining {
309        report.skipped_budget_bytes += len;
310        log::debug!(
311            "[pin] budget exhausted: skipping {} ({}, {} remaining)",
312            label,
313            crate::format_bytes(len),
314            crate::format_bytes(*remaining)
315        );
316        return;
317    }
318
319    match mode {
320        PinMode::Mlock => {
321            if bytes.mlock() {
322                *remaining -= len;
323                report.pinned_bytes += len;
324            } else {
325                report.failed_bytes += len;
326                log::warn!(
327                    "[pin] mlock failed for {} ({}) — check RLIMIT_MEMLOCK; \
328                     continuing unpinned",
329                    label,
330                    crate::format_bytes(len)
331                );
332            }
333        }
334        PinMode::Copy => {
335            *bytes = copy_section(bytes);
336            *remaining -= len;
337            report.pinned_bytes += len;
338            report.heap_copy_bytes += len;
339        }
340    }
341}
342
343/// Copy one admitted metadata section without serial page faults under the
344/// reader's random-access mmap policy. No extra payload scratch or policy toggle.
345fn copy_section(bytes: &OwnedBytes) -> OwnedBytes {
346    #[cfg(target_os = "linux")]
347    {
348        const CHUNK: usize = 128 * 1024;
349        let mut copied = Vec::with_capacity(bytes.len());
350        let mut prefetched = 0;
351        for (i, chunk) in bytes.chunks(CHUNK).enumerate() {
352            let end = (i * CHUNK).saturating_add(2 * CHUNK).min(bytes.len());
353            while prefetched < end {
354                let next = prefetched.saturating_add(CHUNK).min(end);
355                bytes.madvise_range(prefetched..next, libc::MADV_WILLNEED);
356                prefetched = next;
357            }
358            copied.extend_from_slice(chunk);
359        }
360        OwnedBytes::new(copied)
361    }
362    #[cfg(not(target_os = "linux"))]
363    {
364        OwnedBytes::new(bytes.to_vec())
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    #[test]
373    fn copy_pinning_preserves_unaligned_mapped_sections_tails_and_budget() {
374        let len = 5 * 128 * 1024 + 19;
375        let mut mapping = memmap2::MmapMut::map_anon(len).unwrap();
376        for (i, byte) in mapping.iter_mut().enumerate() {
377            *byte = (i % 251) as u8;
378        }
379        let source = OwnedBytes::from_mmap(Arc::new(mapping.make_read_only().unwrap()));
380        let mut section = source.slice(7..len - 5);
381        let size = section.len() as u64;
382        let mut remaining = size - 1;
383        let mut report = PinReport::default();
384        pin_section(
385            &mut section,
386            "test",
387            PinMode::Copy,
388            &mut remaining,
389            &mut report,
390        );
391        assert!(section.is_mmap());
392        assert_eq!(remaining, size - 1);
393        assert_eq!(report.skipped_budget_bytes, size);
394        remaining = size;
395        report = PinReport::default();
396        pin_section(
397            &mut section,
398            "test",
399            PinMode::Copy,
400            &mut remaining,
401            &mut report,
402        );
403        assert!(!section.is_mmap());
404        assert_eq!(section.as_slice(), &source[7..len - 5]);
405        assert_eq!(remaining, 0);
406        assert_eq!(report.pinned_bytes, size);
407        assert_eq!(report.heap_copy_bytes, size);
408        assert_eq!(report.intended_bytes, size);
409        assert_eq!(report.failed_bytes, 0);
410    }
411}
summa_core/segment/pin.rs

summa_core/segment/
pin.rs
