Skip to main content

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, superblock grids. 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::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    fn from_env() -> Self {
54        let budget_mb: u64 = std::env::var("HERMES_PIN_METADATA_BUDGET_MB")
55            .ok()
56            .and_then(|v| v.parse().ok())
57            .unwrap_or(0);
58        let mode = match std::env::var("HERMES_PIN_MODE").as_deref() {
59            Ok("copy") => PinMode::Copy,
60            Ok("mlock") | Err(_) => PinMode::Mlock,
61            Ok(other) => {
62                log::warn!("HERMES_PIN_MODE '{}' unknown; using mlock", other);
63                PinMode::Mlock
64            }
65        };
66        Self {
67            budget_bytes: budget_mb * 1024 * 1024,
68            mode,
69        }
70    }
71}
72
73static PIN_POLICY: OnceLock<PinPolicy> = OnceLock::new();
74
75/// Override the process-wide pin policy. Must be called before the first
76/// segment is opened; returns false (and warns) if the policy was already
77/// initialized.
78pub fn set_pin_policy(policy: PinPolicy) -> bool {
79    let ok = PIN_POLICY.set(policy).is_ok();
80    if !ok {
81        log::warn!("pin policy already initialized; set_pin_policy ignored");
82    }
83    ok
84}
85
86/// The process-wide pin policy (env-initialized on first use).
87pub fn pin_policy() -> &'static PinPolicy {
88    PIN_POLICY.get_or_init(PinPolicy::from_env)
89}
90
91/// Accumulates pin accounting for one segment.
92#[derive(Debug, Default, Clone, Copy)]
93pub struct PinReport {
94    /// Bytes of pinnable metadata found (regardless of budget/failures)
95    pub intended_bytes: u64,
96    /// Bytes actually pinned
97    pub pinned_bytes: u64,
98    /// Bytes skipped because the budget was exhausted
99    pub skipped_budget_bytes: u64,
100    /// Bytes where mlock failed (RLIMIT_MEMLOCK etc.)
101    pub failed_bytes: u64,
102}
103
104/// Pin one metadata section, updating `remaining` budget and the report.
105///
106/// In `Copy` mode the section is replaced with a heap copy (heap memory is
107/// not page-cache-evictable). In `Mlock` mode the mmap pages are locked in
108/// place. Non-mmap-backed sections (RAM directories) are already resident
109/// and are skipped silently.
110pub(crate) fn pin_section(
111    bytes: &mut OwnedBytes,
112    label: &str,
113    mode: PinMode,
114    remaining: &mut u64,
115    report: &mut PinReport,
116) {
117    if !bytes.is_mmap() || bytes.is_empty() {
118        return;
119    }
120    let len = bytes.len() as u64;
121    report.intended_bytes += len;
122
123    if len > *remaining {
124        report.skipped_budget_bytes += len;
125        log::debug!(
126            "[pin] budget exhausted: skipping {} ({} bytes, {} remaining)",
127            label,
128            len,
129            *remaining
130        );
131        return;
132    }
133
134    match mode {
135        PinMode::Mlock => {
136            if bytes.mlock() {
137                *remaining -= len;
138                report.pinned_bytes += len;
139            } else {
140                report.failed_bytes += len;
141                log::warn!(
142                    "[pin] mlock failed for {} ({} bytes) — check RLIMIT_MEMLOCK; \
143                     continuing unpinned",
144                    label,
145                    len
146                );
147            }
148        }
149        PinMode::Copy => {
150            *bytes = OwnedBytes::new(bytes.to_vec());
151            *remaining -= len;
152            report.pinned_bytes += len;
153        }
154    }
155}