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