1use anyhow::{bail, Result};
10use lru::LruCache;
11use memmap2::Mmap;
12use oxirs_core::parallel::*;
13use parking_lot::RwLock;
14use std::collections::{HashMap, VecDeque};
15use std::num::NonZeroUsize;
16use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
17use std::sync::Arc;
18use std::time::Instant;
19use tracing::{trace, warn};
20
21const VECTOR_PAGE_SIZE: usize = 16384;
23
24const DEFAULT_MAX_PAGES: usize = 10000;
26
27#[cfg(target_os = "linux")]
42mod numa {
43 use libc::{c_int, c_uint, c_ulong, c_void};
44 use std::sync::OnceLock;
45
46 const NODE_SYSFS_ROOT: &str = "/sys/devices/system/node";
48
49 const CPU_SYSFS_ROOT: &str = "/sys/devices/system/cpu";
51
52 const MAX_SUPPORTED_ID: i32 = 65535;
58
59 const NODEMASK_BITS: usize = std::mem::size_of::<c_ulong>() * 8;
61
62 pub const MPOL_BIND: c_int = libc::MPOL_BIND;
64
65 pub const MPOL_INTERLEAVE: c_int = libc::MPOL_INTERLEAVE;
67
68 pub fn parse_cpuset_list(raw: &str) -> Vec<i32> {
76 let mut ids: Vec<i32> = Vec::new();
77
78 for entry in raw.trim().split(',') {
79 let entry = entry.trim();
80 if entry.is_empty() {
81 continue;
82 }
83
84 let range_part = entry.split(':').next().unwrap_or(entry);
88 let (start, end) = match range_part.split_once('-') {
89 Some((lo, hi)) => {
90 let (lo, hi) = (lo.trim(), hi.trim());
91 match (lo.parse::<i32>(), hi.parse::<i32>()) {
92 (Ok(lo), Ok(hi)) => (lo, hi),
93 _ => continue,
94 }
95 }
96 None => match range_part.parse::<i32>() {
97 Ok(single) => (single, single),
98 Err(_) => continue,
99 },
100 };
101
102 if start < 0 || end < start || start > MAX_SUPPORTED_ID {
104 continue;
105 }
106 let end = end.min(MAX_SUPPORTED_ID);
107 ids.extend(start..=end);
108 }
109
110 ids.sort_unstable();
111 ids.dedup();
112 ids
113 }
114
115 fn read_sysfs(path: &str) -> Option<String> {
117 std::fs::read_to_string(path)
118 .ok()
119 .map(|s| s.trim().to_string())
120 }
121
122 fn discover_nodes() -> Vec<i32> {
129 let raw = read_sysfs(&format!("{NODE_SYSFS_ROOT}/possible"))
130 .or_else(|| read_sysfs(&format!("{NODE_SYSFS_ROOT}/online")));
131
132 let nodes = raw.map(|raw| parse_cpuset_list(&raw)).unwrap_or_default();
133 if nodes.is_empty() {
134 vec![0]
135 } else {
136 nodes
137 }
138 }
139
140 fn topology() -> &'static [i32] {
145 static TOPOLOGY: OnceLock<Vec<i32>> = OnceLock::new();
146 TOPOLOGY.get_or_init(discover_nodes)
147 }
148
149 fn build_cpu_node_map() -> Vec<i32> {
154 let mut map: Vec<i32> = Vec::new();
155
156 for &node in topology() {
157 let path = format!("{NODE_SYSFS_ROOT}/node{node}/cpulist");
158 let Some(raw) = read_sysfs(&path) else {
159 continue;
160 };
161 for cpu in parse_cpuset_list(&raw) {
162 let idx = cpu as usize;
163 if idx >= map.len() {
164 map.resize(idx + 1, 0);
165 }
166 map[idx] = node;
167 }
168 }
169
170 map
171 }
172
173 fn cpu_node_map() -> &'static [i32] {
175 static CPU_NODE_MAP: OnceLock<Vec<i32>> = OnceLock::new();
176 CPU_NODE_MAP.get_or_init(build_cpu_node_map)
177 }
178
179 fn probe_node_of_cpu(cpu: i32) -> Option<i32> {
186 let dir = std::fs::read_dir(format!("{CPU_SYSFS_ROOT}/cpu{cpu}")).ok()?;
187 for entry in dir.flatten() {
188 let name = entry.file_name();
189 let name = name.to_string_lossy();
190 if let Some(rest) = name.strip_prefix("node") {
191 if let Ok(node) = rest.parse::<i32>() {
192 return Some(node);
193 }
194 }
195 }
196 None
197 }
198
199 pub fn is_available() -> bool {
204 static AVAILABLE: OnceLock<bool> = OnceLock::new();
205 *AVAILABLE
206 .get_or_init(|| std::path::Path::new(&format!("{NODE_SYSFS_ROOT}/node0")).exists())
207 }
208
209 pub fn nodes() -> &'static [i32] {
214 topology()
215 }
216
217 pub fn max_node() -> i32 {
219 topology().last().copied().unwrap_or(0)
220 }
221
222 pub fn node_of_cpu(cpu: i32) -> i32 {
224 if cpu < 0 {
225 return 0;
226 }
227 if let Some(node) = cpu_node_map().get(cpu as usize) {
228 return *node;
229 }
230 probe_node_of_cpu(cpu).unwrap_or(0)
231 }
232
233 pub fn nodemask_from_nodes(nodes: &[i32], max_node: i32) -> Option<(Vec<c_ulong>, c_ulong)> {
242 let highest = nodes
243 .iter()
244 .copied()
245 .chain(std::iter::once(max_node))
246 .max()
247 .unwrap_or(0)
248 .clamp(0, MAX_SUPPORTED_ID);
249
250 let words = highest as usize / NODEMASK_BITS + 2;
251 let mut mask = vec![0 as c_ulong; words];
252 let mut any = false;
253
254 for &node in nodes {
255 if node < 0 || node > highest {
256 continue;
257 }
258 let idx = node as usize / NODEMASK_BITS;
259 let bit = node as usize % NODEMASK_BITS;
260 mask[idx] |= (1 as c_ulong) << bit;
261 any = true;
262 }
263
264 if !any {
265 return None;
266 }
267
268 Some((mask, (words * NODEMASK_BITS) as c_ulong))
269 }
270
271 pub unsafe fn mbind(
284 addr: *mut c_void,
285 len: usize,
286 mode: c_int,
287 nodes: &[i32],
288 ) -> std::io::Result<()> {
289 let Some((mask, maxnode)) = nodemask_from_nodes(nodes, max_node()) else {
290 return Err(std::io::Error::new(
291 std::io::ErrorKind::InvalidInput,
292 "empty NUMA node mask",
293 ));
294 };
295
296 let rc = unsafe {
301 libc::syscall(
302 libc::SYS_mbind,
303 addr,
304 len as c_ulong,
305 mode,
306 mask.as_ptr(),
307 maxnode,
308 0 as c_uint,
309 )
310 };
311
312 if rc == 0 {
313 Ok(())
314 } else {
315 Err(std::io::Error::last_os_error())
316 }
317 }
318}
319
320#[cfg(not(target_os = "linux"))]
324mod numa {
325 use std::ffi::c_void;
326
327 pub const MPOL_BIND: i32 = 2;
329
330 pub const MPOL_INTERLEAVE: i32 = 3;
332
333 pub fn is_available() -> bool {
334 false
335 }
336
337 pub fn nodes() -> &'static [i32] {
338 &[0]
339 }
340
341 pub fn max_node() -> i32 {
342 0
343 }
344
345 pub fn node_of_cpu(_cpu: i32) -> i32 {
346 0
347 }
348
349 pub unsafe fn mbind(
354 _addr: *mut c_void,
355 _len: usize,
356 _mode: i32,
357 _nodes: &[i32],
358 ) -> std::io::Result<()> {
359 Err(std::io::Error::new(
360 std::io::ErrorKind::Unsupported,
361 "mbind is only available on Linux",
362 ))
363 }
364}
365
366fn page_size() -> usize {
368 static PAGE_SIZE: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
369 *PAGE_SIZE.get_or_init(|| {
370 #[cfg(unix)]
371 {
372 let value = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
375 if value > 0 {
376 return value as usize;
377 }
378 }
379 4096
382 })
383}
384
385fn page_aligned_subrange(
393 start_addr: usize,
394 len: usize,
395 page_size: usize,
396) -> Option<(usize, usize)> {
397 if len == 0 || page_size == 0 || !page_size.is_power_of_two() {
398 return None;
399 }
400
401 let end_addr = start_addr.checked_add(len)?;
402 let aligned_start = start_addr.checked_add(page_size - 1)? & !(page_size - 1);
404 let aligned_end = end_addr & !(page_size - 1);
405
406 if aligned_end > aligned_start {
407 Some((aligned_start, aligned_end - aligned_start))
408 } else {
409 None
410 }
411}
412
413#[derive(Debug, Clone)]
415struct AccessPattern {
416 page_id: usize,
417 access_time: Instant,
418 access_count: usize,
419}
420
421#[derive(Debug)]
423pub struct PageCacheEntry {
424 data: Vec<u8>,
425 page_id: usize,
426 last_access: Instant,
427 inserted_at: Instant,
430 access_count: AtomicUsize,
431 reference_bit: AtomicBool,
434 dirty: bool,
435 numa_node: i32,
436}
437
438impl PageCacheEntry {
439 pub fn data(&self) -> &[u8] {
441 &self.data
442 }
443
444 pub fn numa_node(&self) -> i32 {
446 self.numa_node
447 }
448}
449
450#[derive(Debug, Clone, Copy)]
452pub enum EvictionPolicy {
453 LRU, LFU, FIFO, Clock, ARC, }
459
460#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
462pub enum MemoryPressure {
463 Low,
464 Medium,
465 High,
466 Critical,
467}
468
469pub struct AdvancedMemoryMap {
471 mmap: Option<Mmap>,
473
474 file_path: Option<std::path::PathBuf>,
476
477 page_cache: Arc<RwLock<LruCache<usize, Arc<PageCacheEntry>>>>,
479
480 access_patterns: Arc<RwLock<VecDeque<AccessPattern>>>,
482
483 page_frequency: Arc<RwLock<HashMap<usize, usize>>>,
485
486 eviction_policy: EvictionPolicy,
488
489 total_memory: AtomicUsize,
491 cache_hits: AtomicU64,
492 cache_misses: AtomicU64,
493
494 numa_enabled: bool,
496 numa_nodes: Vec<i32>,
497
498 memory_pressure: Arc<RwLock<MemoryPressure>>,
500
501 max_pages: usize,
503 page_size: usize,
504 prefetch_distance: usize,
505}
506
507impl AdvancedMemoryMap {
508 pub fn new(mmap: Option<Mmap>, max_pages: usize) -> Self {
510 let numa_enabled = numa::is_available();
511 let numa_nodes = if numa_enabled {
512 numa::nodes().to_vec()
513 } else {
514 vec![0]
515 };
516
517 let cache_size = NonZeroUsize::new(max_pages)
518 .unwrap_or(NonZeroUsize::new(1).expect("constant 1 is non-zero"));
519
520 Self {
521 mmap,
522 file_path: None,
523 page_cache: Arc::new(RwLock::new(LruCache::new(cache_size))),
524 access_patterns: Arc::new(RwLock::new(VecDeque::with_capacity(1000))),
525 page_frequency: Arc::new(RwLock::new(HashMap::new())),
526 eviction_policy: EvictionPolicy::ARC,
527 total_memory: AtomicUsize::new(0),
528 cache_hits: AtomicU64::new(0),
529 cache_misses: AtomicU64::new(0),
530 numa_enabled,
531 numa_nodes,
532 memory_pressure: Arc::new(RwLock::new(MemoryPressure::Low)),
533 max_pages,
534 page_size: VECTOR_PAGE_SIZE,
535 prefetch_distance: 3,
536 }
537 }
538
539 pub fn new_with_path(
541 mmap: Option<Mmap>,
542 max_pages: usize,
543 file_path: Option<std::path::PathBuf>,
544 ) -> Self {
545 let mut s = Self::new(mmap, max_pages);
546 s.file_path = file_path;
547 s
548 }
549
550 pub fn get_page(&self, page_id: usize) -> Result<Arc<PageCacheEntry>> {
552 {
554 let mut cache = self.page_cache.write();
555 if let Some(entry) = cache.get(&page_id) {
556 self.cache_hits.fetch_add(1, Ordering::Relaxed);
557 entry.access_count.fetch_add(1, Ordering::Relaxed);
558 entry.reference_bit.store(true, Ordering::Relaxed);
560 self.record_access(page_id);
561 return Ok(Arc::clone(entry));
562 }
563 }
564
565 self.cache_misses.fetch_add(1, Ordering::Relaxed);
567 self.load_page(page_id)
568 }
569
570 fn load_page(&self, page_id: usize) -> Result<Arc<PageCacheEntry>> {
572 let mmap = self
573 .mmap
574 .as_ref()
575 .ok_or_else(|| anyhow::anyhow!("No memory mapping available"))?;
576
577 let start = page_id * self.page_size;
578 let end = (start + self.page_size).min(mmap.len());
579
580 if start >= mmap.len() {
581 bail!("Page {} out of bounds", page_id);
582 }
583
584 let page_data = mmap[start..end].to_vec();
586
587 let numa_node = if self.numa_enabled {
589 let cpu = sched_getcpu();
590 numa::node_of_cpu(cpu)
591 } else {
592 0
593 };
594
595 let now = Instant::now();
596 let entry = Arc::new(PageCacheEntry {
597 data: page_data,
598 page_id,
599 last_access: now,
600 inserted_at: now,
601 access_count: AtomicUsize::new(1),
602 reference_bit: AtomicBool::new(true),
603 dirty: false,
604 numa_node,
605 });
606
607 self.check_memory_pressure();
609 if *self.memory_pressure.read() >= MemoryPressure::High {
610 self.evict_pages(1)?;
611 }
612
613 {
615 let mut cache = self.page_cache.write();
616 cache.put(page_id, Arc::clone(&entry));
617 }
618
619 self.total_memory
620 .fetch_add(entry.data.len(), Ordering::Relaxed);
621 self.record_access(page_id);
622
623 self.prefetch_pages(page_id);
625
626 Ok(entry)
627 }
628
629 fn record_access(&self, page_id: usize) {
631 let mut patterns = self.access_patterns.write();
632 patterns.push_back(AccessPattern {
633 page_id,
634 access_time: Instant::now(),
635 access_count: 1,
636 });
637
638 while patterns.len() > 1000 {
640 patterns.pop_front();
641 }
642
643 let mut freq = self.page_frequency.write();
645 *freq.entry(page_id).or_insert(0) += 1;
646 }
647
648 fn prefetch_pages(&self, current_page: usize) {
650 let patterns = self.access_patterns.read();
651 let freq = self.page_frequency.read();
652
653 let recent_patterns: Vec<_> = patterns.iter().rev().take(10).collect();
655
656 let is_sequential = recent_patterns
658 .windows(2)
659 .all(|w| w[0].page_id > 0 && w[0].page_id == w[1].page_id + 1);
660
661 let stride = if recent_patterns.len() >= 3 {
663 let diff1 = recent_patterns[0]
664 .page_id
665 .saturating_sub(recent_patterns[1].page_id);
666 let diff2 = recent_patterns[1]
667 .page_id
668 .saturating_sub(recent_patterns[2].page_id);
669 if diff1 == diff2 && diff1 > 0 && diff1 <= 10 {
670 Some(diff1)
671 } else {
672 None
673 }
674 } else {
675 None
676 };
677
678 if is_sequential {
680 for i in 1..=(self.prefetch_distance * 2) {
682 let prefetch_page = current_page + i;
683 self.async_prefetch(prefetch_page);
684 }
685 } else if let Some(stride) = stride {
686 for i in 1..=self.prefetch_distance {
688 let prefetch_page = current_page + (i * stride);
689 self.async_prefetch(prefetch_page);
690 }
691 } else {
692 for i in 1..=self.prefetch_distance {
694 let prefetch_page = current_page + i;
695
696 let frequency = *freq.get(&prefetch_page).unwrap_or(&0);
698 if frequency > 0 {
699 self.async_prefetch(prefetch_page);
700 }
701 }
702 }
703
704 let nearby_range = current_page.saturating_sub(3)..=(current_page + 3);
706 for page_id in nearby_range {
707 let frequency = *freq.get(&page_id).unwrap_or(&0);
708 if frequency > 2 && page_id != current_page {
709 self.async_prefetch(page_id);
710 }
711 }
712 }
713
714 pub fn async_prefetch(&self, page_id: usize) {
716 {
718 let cache = self.page_cache.read();
719 if cache.contains(&page_id) {
720 return;
721 }
722 }
723
724 if *self.memory_pressure.read() >= MemoryPressure::High {
726 return;
727 }
728
729 let self_clone = self.clone_ref();
730 spawn(move || {
731 let _ = self_clone.get_page(page_id);
732 });
733 }
734
735 fn check_memory_pressure(&self) {
737 let total_memory = self.total_memory.load(Ordering::Relaxed);
738 let max_memory = self.max_pages * self.page_size;
739
740 let pressure = if total_memory < max_memory / 2 {
741 MemoryPressure::Low
742 } else if total_memory < max_memory * 3 / 4 {
743 MemoryPressure::Medium
744 } else if total_memory < max_memory * 9 / 10 {
745 MemoryPressure::High
746 } else {
747 MemoryPressure::Critical
748 };
749
750 *self.memory_pressure.write() = pressure;
751 }
752
753 fn evict_pages(&self, num_pages: usize) -> Result<()> {
755 match self.eviction_policy {
756 EvictionPolicy::LRU => self.evict_lru(num_pages),
757 EvictionPolicy::LFU => self.evict_lfu(num_pages),
758 EvictionPolicy::FIFO => self.evict_fifo(num_pages),
759 EvictionPolicy::Clock => self.evict_clock(num_pages),
760 EvictionPolicy::ARC => self.evict_arc(num_pages),
761 }
762 }
763
764 fn evict_lru(&self, num_pages: usize) -> Result<()> {
766 let mut cache = self.page_cache.write();
767
768 for _ in 0..num_pages {
770 if let Some((_, entry)) = cache.pop_lru() {
771 self.total_memory
772 .fetch_sub(entry.data.len(), Ordering::Relaxed);
773
774 if entry.dirty {
776 if let Err(e) = self.write_back_page(entry.page_id, &entry.data) {
777 warn!("Failed to write back page {}: {}", entry.page_id, e);
778 }
779 }
780 }
781 }
782
783 Ok(())
784 }
785
786 fn evict_lfu(&self, num_pages: usize) -> Result<()> {
788 let cache = self.page_cache.read();
789 let freq = self.page_frequency.read();
790
791 let mut pages_by_freq: Vec<(usize, usize)> = cache
793 .iter()
794 .map(|(page_id, _)| (*page_id, *freq.get(page_id).unwrap_or(&0)))
795 .collect();
796 pages_by_freq.sort_by_key(|(_, freq)| *freq);
797
798 drop(cache);
800 drop(freq);
801
802 let mut cache = self.page_cache.write();
803 for (page_id, _) in pages_by_freq.iter().take(num_pages) {
804 if let Some(entry) = cache.pop(page_id) {
805 self.total_memory
806 .fetch_sub(entry.data.len(), Ordering::Relaxed);
807 if entry.dirty {
808 if let Err(e) = self.write_back_page(entry.page_id, &entry.data) {
809 warn!("Failed to write back dirty page {}: {}", entry.page_id, e);
810 }
811 }
812 }
813 }
814
815 Ok(())
816 }
817
818 fn evict_fifo(&self, num_pages: usize) -> Result<()> {
823 let mut pages_by_age: Vec<(usize, Instant)> = {
826 let cache = self.page_cache.read();
827 cache
828 .iter()
829 .map(|(page_id, entry)| (*page_id, entry.inserted_at))
830 .collect()
831 };
832 pages_by_age.sort_by_key(|(_, inserted_at)| *inserted_at);
833
834 let mut cache = self.page_cache.write();
835 for (page_id, _) in pages_by_age.iter().take(num_pages) {
836 if let Some(entry) = cache.pop(page_id) {
837 self.total_memory
838 .fetch_sub(entry.data.len(), Ordering::Relaxed);
839 if entry.dirty {
840 if let Err(e) = self.write_back_page(entry.page_id, &entry.data) {
841 warn!("Failed to write back page {}: {}", entry.page_id, e);
842 }
843 }
844 }
845 }
846
847 Ok(())
848 }
849
850 fn evict_clock(&self, num_pages: usize) -> Result<()> {
855 if num_pages == 0 {
856 return Ok(());
857 }
858
859 let mut cache = self.page_cache.write();
860
861 let mut order: Vec<usize> = cache.iter().map(|(page_id, _)| *page_id).collect();
863 order.sort_unstable();
864 if order.is_empty() {
865 return Ok(());
866 }
867
868 let mut to_evict: Vec<usize> = Vec::with_capacity(num_pages);
869 let max_steps = order.len() * 3 + num_pages;
872 let mut hand = 0usize;
873 let mut steps = 0usize;
874
875 while to_evict.len() < num_pages && steps < max_steps {
876 let page_id = order[hand % order.len()];
877 hand += 1;
878 steps += 1;
879
880 if let Some(entry) = cache.peek(&page_id) {
881 if entry.reference_bit.swap(false, Ordering::Relaxed) {
882 continue;
884 }
885 to_evict.push(page_id);
887 }
888 }
889
890 for page_id in to_evict {
891 if let Some(entry) = cache.pop(&page_id) {
892 self.total_memory
893 .fetch_sub(entry.data.len(), Ordering::Relaxed);
894 if entry.dirty {
895 if let Err(e) = self.write_back_page(entry.page_id, &entry.data) {
896 warn!("Failed to write back page {}: {}", entry.page_id, e);
897 }
898 }
899 }
900 }
901
902 Ok(())
903 }
904
905 fn evict_arc(&self, num_pages: usize) -> Result<()> {
907 let cache = self.page_cache.read();
909 let freq = self.page_frequency.read();
910
911 let now = Instant::now();
913 let mut scored_pages: Vec<(usize, f64)> = cache
914 .iter()
915 .map(|(page_id, entry)| {
916 let recency_score =
917 1.0 / (now.duration_since(entry.last_access).as_secs_f64() + 1.0);
918 let frequency_score = *freq.get(page_id).unwrap_or(&0) as f64;
919 let combined_score = recency_score * 0.5 + frequency_score * 0.5;
920 (*page_id, combined_score)
921 })
922 .collect();
923
924 scored_pages.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
925
926 drop(cache);
927 drop(freq);
928
929 let mut cache = self.page_cache.write();
930 for (page_id, _) in scored_pages.iter().take(num_pages) {
931 if let Some(entry) = cache.pop(page_id) {
932 self.total_memory
933 .fetch_sub(entry.data.len(), Ordering::Relaxed);
934 if entry.dirty {
935 if let Err(e) = self.write_back_page(entry.page_id, &entry.data) {
936 warn!("Failed to write back dirty page {}: {}", entry.page_id, e);
937 }
938 }
939 }
940 }
941
942 Ok(())
943 }
944
945 pub fn stats(&self) -> MemoryMapStats {
947 let cache = self.page_cache.read();
948
949 MemoryMapStats {
950 total_pages: cache.len(),
951 total_memory: self.total_memory.load(Ordering::Relaxed),
952 cache_hits: self.cache_hits.load(Ordering::Relaxed),
953 cache_misses: self.cache_misses.load(Ordering::Relaxed),
954 hit_rate: self.calculate_hit_rate(),
955 memory_pressure: *self.memory_pressure.read(),
956 numa_enabled: self.numa_enabled,
957 }
958 }
959
960 fn calculate_hit_rate(&self) -> f64 {
961 let hits = self.cache_hits.load(Ordering::Relaxed) as f64;
962 let misses = self.cache_misses.load(Ordering::Relaxed) as f64;
963 let total = hits + misses;
964 if total > 0.0 {
965 hits / total
966 } else {
967 0.0
968 }
969 }
970
971 fn clone_ref(&self) -> Self {
972 Self {
973 mmap: None, file_path: self.file_path.clone(),
975 page_cache: Arc::clone(&self.page_cache),
976 access_patterns: Arc::clone(&self.access_patterns),
977 page_frequency: Arc::clone(&self.page_frequency),
978 eviction_policy: self.eviction_policy,
979 total_memory: AtomicUsize::new(0),
980 cache_hits: AtomicU64::new(0),
981 cache_misses: AtomicU64::new(0),
982 numa_enabled: self.numa_enabled,
983 numa_nodes: self.numa_nodes.clone(),
984 memory_pressure: Arc::clone(&self.memory_pressure),
985 max_pages: self.max_pages,
986 page_size: self.page_size,
987 prefetch_distance: self.prefetch_distance,
988 }
989 }
990
991 fn write_back_page(&self, page_id: usize, data: &[u8]) -> Result<()> {
993 use std::io::{Seek, SeekFrom, Write};
994 let path = match &self.file_path {
995 Some(p) => p,
996 None => return Ok(()), };
998 let mut file = std::fs::OpenOptions::new()
999 .write(true)
1000 .open(path)
1001 .map_err(|e| anyhow::anyhow!("Failed to open file for write-back: {}", e))?;
1002 let offset = (page_id * self.page_size) as u64;
1003 file.seek(SeekFrom::Start(offset))
1004 .map_err(|e| anyhow::anyhow!("Failed to seek to page {}: {}", page_id, e))?;
1005 file.write_all(data)
1006 .map_err(|e| anyhow::anyhow!("Failed to write page {}: {}", page_id, e))?;
1007 Ok(())
1008 }
1009
1010 pub fn flush_dirty_pages(&self) -> Result<()> {
1012 if self.file_path.is_none() {
1013 return Ok(());
1014 }
1015 let cache = self.page_cache.read();
1016 for (_, entry) in cache.iter() {
1017 if entry.dirty {
1018 self.write_back_page(entry.page_id, &entry.data)?;
1019 }
1020 }
1021 Ok(())
1022 }
1023}
1024
1025#[derive(Debug, Clone)]
1027pub struct MemoryMapStats {
1028 pub total_pages: usize,
1029 pub total_memory: usize,
1030 pub cache_hits: u64,
1031 pub cache_misses: u64,
1032 pub hit_rate: f64,
1033 pub memory_pressure: MemoryPressure,
1034 pub numa_enabled: bool,
1035}
1036
1037#[cfg(target_os = "linux")]
1039fn sched_getcpu() -> i32 {
1040 unsafe { libc::sched_getcpu() }
1041}
1042
1043#[cfg(not(target_os = "linux"))]
1044fn sched_getcpu() -> i32 {
1045 0
1046}
1047
1048pub struct NumaVectorAllocator {
1050 numa_nodes: Vec<i32>,
1051 current_node: AtomicUsize,
1052}
1053
1054impl Default for NumaVectorAllocator {
1055 fn default() -> Self {
1056 Self::new()
1057 }
1058}
1059
1060impl NumaVectorAllocator {
1061 pub fn new() -> Self {
1062 let numa_nodes = if numa::is_available() {
1063 numa::nodes().to_vec()
1064 } else {
1065 vec![0]
1066 };
1067
1068 Self {
1069 numa_nodes,
1070 current_node: AtomicUsize::new(0),
1071 }
1072 }
1073
1074 pub fn allocate_on_node(&self, size: usize, node: Option<i32>) -> Vec<u8> {
1076 if !numa::is_available() || size == 0 {
1077 return vec![0u8; size];
1078 }
1079
1080 let mut buffer: Vec<u8> = Vec::with_capacity(size);
1084 self.apply_memory_policy(buffer.as_mut_ptr().cast(), size, node);
1085 buffer.resize(size, 0u8);
1086 buffer
1087 }
1088
1089 pub fn allocate_vector_on_node(&self, dimensions: usize, node: Option<i32>) -> Vec<f32> {
1091 let mut vec: Vec<f32> = Vec::with_capacity(dimensions);
1092
1093 if numa::is_available() && dimensions > 0 {
1094 if let Some(byte_len) = dimensions.checked_mul(std::mem::size_of::<f32>()) {
1095 let target = self
1099 .explicit_node(node)
1100 .or_else(|| Some(self.preferred_node()));
1101 self.apply_memory_policy(vec.as_mut_ptr().cast(), byte_len, target);
1102 }
1103 }
1104
1105 vec.resize(dimensions, 0.0f32);
1106 vec
1107 }
1108
1109 fn apply_memory_policy(&self, ptr: *mut std::ffi::c_void, byte_len: usize, node: Option<i32>) {
1117 let page = page_size();
1118 let Some((addr, len)) = page_aligned_subrange(ptr as usize, byte_len, page) else {
1119 return;
1122 };
1123
1124 let rr;
1127 let (mode, target_nodes): (i32, &[i32]) = match self.explicit_node(node) {
1128 Some(explicit) => {
1129 rr = [explicit];
1130 (numa::MPOL_BIND, &rr)
1131 }
1132 None if self.numa_nodes.len() > 1 && len > page => {
1136 (numa::MPOL_INTERLEAVE, self.numa_nodes.as_slice())
1137 }
1138 None => {
1139 rr = [self.next_round_robin_node()];
1140 (numa::MPOL_BIND, &rr)
1141 }
1142 };
1143
1144 if let Err(err) =
1149 unsafe { numa::mbind(addr as *mut std::ffi::c_void, len, mode, target_nodes) }
1150 {
1151 trace!(
1152 "mbind({} bytes, mode {}) failed, falling back to default placement: {}",
1153 len,
1154 mode,
1155 err
1156 );
1157 }
1158 }
1159
1160 fn explicit_node(&self, node: Option<i32>) -> Option<i32> {
1166 let node = node?;
1167 if node >= 0 && node <= numa::max_node() && self.numa_nodes.contains(&node) {
1168 Some(node)
1169 } else {
1170 None
1171 }
1172 }
1173
1174 fn next_round_robin_node(&self) -> i32 {
1176 if self.numa_nodes.is_empty() {
1177 return 0;
1178 }
1179 let idx = self.current_node.fetch_add(1, Ordering::Relaxed) % self.numa_nodes.len();
1180 self.numa_nodes[idx]
1181 }
1182
1183 pub fn preferred_node(&self) -> i32 {
1185 if numa::is_available() {
1186 numa::node_of_cpu(sched_getcpu())
1187 } else {
1188 0
1189 }
1190 }
1191}
1192
1193#[cfg(test)]
1194mod tests {
1195 use super::*;
1196
1197 #[test]
1198 fn test_memory_pressure() {
1199 let mmap = AdvancedMemoryMap::new(None, 100);
1200
1201 assert_eq!(*mmap.memory_pressure.read(), MemoryPressure::Low);
1202
1203 mmap.total_memory
1205 .store(50 * VECTOR_PAGE_SIZE, Ordering::Relaxed);
1206 mmap.check_memory_pressure();
1207 assert_eq!(*mmap.memory_pressure.read(), MemoryPressure::Medium);
1208
1209 mmap.total_memory
1210 .store(90 * VECTOR_PAGE_SIZE, Ordering::Relaxed);
1211 mmap.check_memory_pressure();
1212 assert_eq!(*mmap.memory_pressure.read(), MemoryPressure::Critical);
1213 }
1214
1215 #[test]
1216 fn test_cache_stats() {
1217 let mmap = AdvancedMemoryMap::new(None, 100);
1218
1219 mmap.cache_hits.store(75, Ordering::Relaxed);
1220 mmap.cache_misses.store(25, Ordering::Relaxed);
1221
1222 let stats = mmap.stats();
1223 assert_eq!(stats.cache_hits, 75);
1224 assert_eq!(stats.cache_misses, 25);
1225 assert_eq!(stats.hit_rate, 0.75);
1226 }
1227
1228 fn insert_test_page(map: &AdvancedMemoryMap, page_id: usize, referenced: bool) {
1231 let now = Instant::now();
1232 let entry = Arc::new(PageCacheEntry {
1233 data: vec![0u8; 8],
1234 page_id,
1235 last_access: now,
1236 inserted_at: now,
1237 access_count: AtomicUsize::new(1),
1238 reference_bit: AtomicBool::new(referenced),
1239 dirty: false,
1240 numa_node: 0,
1241 });
1242 map.page_cache.write().put(page_id, entry);
1243 }
1244
1245 #[test]
1246 fn regression_fifo_evicts_oldest_not_lru() {
1247 let map = AdvancedMemoryMap::new(None, 100);
1248 insert_test_page(&map, 0, false);
1250 insert_test_page(&map, 1, false);
1251 insert_test_page(&map, 2, false);
1252
1253 {
1256 let mut cache = map.page_cache.write();
1257 let _ = cache.get(&0);
1258 }
1259
1260 map.evict_fifo(1).expect("fifo eviction");
1261
1262 let cache = map.page_cache.read();
1263 assert!(
1264 cache.peek(&0).is_none(),
1265 "FIFO must evict the first-inserted page (0)"
1266 );
1267 assert!(cache.peek(&1).is_some());
1268 assert!(cache.peek(&2).is_some());
1269 }
1270
1271 #[test]
1272 fn regression_clock_gives_second_chance() {
1273 let map = AdvancedMemoryMap::new(None, 100);
1274 insert_test_page(&map, 0, true);
1277 insert_test_page(&map, 1, false);
1278 insert_test_page(&map, 2, false);
1279
1280 map.evict_clock(1).expect("clock eviction");
1281
1282 let cache = map.page_cache.read();
1283 assert!(
1284 cache.peek(&0).is_some(),
1285 "referenced page 0 must survive one Clock sweep (second chance)"
1286 );
1287 assert!(
1288 cache.peek(&1).is_none(),
1289 "unreferenced page 1 must be the Clock victim"
1290 );
1291 assert!(
1293 !cache
1294 .peek(&0)
1295 .expect("page 0 present")
1296 .reference_bit
1297 .load(Ordering::Relaxed),
1298 "Clock sweep must clear the reference bit it consumed"
1299 );
1300 }
1301
1302 #[test]
1307 fn test_numa_topology_is_sane() {
1308 let nodes = numa::nodes();
1309 assert!(!nodes.is_empty(), "node list must never be empty");
1310 assert!(
1311 nodes.iter().all(|&n| n >= 0),
1312 "node ids must be non-negative, got {nodes:?}"
1313 );
1314 assert!(
1315 nodes.windows(2).all(|w| w[0] < w[1]),
1316 "node list must be sorted and de-duplicated, got {nodes:?}"
1317 );
1318 assert_eq!(
1319 numa::max_node(),
1320 nodes.iter().copied().max().unwrap_or(0),
1321 "max_node must agree with the node list"
1322 );
1323
1324 let allocator = NumaVectorAllocator::new();
1326 assert!(!allocator.numa_nodes.is_empty());
1327 assert!(allocator.numa_nodes.iter().all(|&n| n >= 0));
1328 }
1329
1330 #[test]
1331 fn test_node_of_cpu_within_topology() {
1332 let node = numa::node_of_cpu(0);
1333 assert!(
1334 (0..=numa::max_node()).contains(&node),
1335 "node_of_cpu(0) = {node} must lie in 0..={}",
1336 numa::max_node()
1337 );
1338
1339 assert_eq!(numa::node_of_cpu(-1), 0);
1341 assert_eq!(numa::node_of_cpu(i32::MAX), 0);
1342
1343 let allocator = NumaVectorAllocator::new();
1345 let preferred = allocator.preferred_node();
1346 assert!((0..=numa::max_node()).contains(&preferred));
1347 }
1348
1349 #[cfg(target_os = "linux")]
1350 #[test]
1351 fn test_parse_cpuset_list() {
1352 use super::numa::parse_cpuset_list;
1353
1354 assert_eq!(parse_cpuset_list("0"), vec![0]);
1355 assert_eq!(parse_cpuset_list("0-1"), vec![0, 1]);
1356 assert_eq!(parse_cpuset_list("0-1,4"), vec![0, 1, 4]);
1357 assert_eq!(parse_cpuset_list("0-3\n"), vec![0, 1, 2, 3]);
1358 assert_eq!(parse_cpuset_list("4,0-1,1"), vec![0, 1, 4]);
1360 assert_eq!(parse_cpuset_list("0-2:1/2"), vec![0, 1, 2]);
1362
1363 assert!(parse_cpuset_list("").is_empty());
1365 assert!(parse_cpuset_list(" \n").is_empty());
1366 assert!(parse_cpuset_list(",,").is_empty());
1367
1368 assert!(parse_cpuset_list("abc").is_empty());
1370 assert!(parse_cpuset_list("-").is_empty());
1371 assert!(parse_cpuset_list("3-1").is_empty(), "inverted range");
1372 assert!(parse_cpuset_list("-5").is_empty(), "negative id");
1373 assert!(parse_cpuset_list("99999999999999999999").is_empty());
1374 assert_eq!(parse_cpuset_list("0,bogus,2"), vec![0, 2]);
1376 assert!(parse_cpuset_list("0-4294967295").is_empty());
1378 }
1379
1380 #[cfg(target_os = "linux")]
1381 #[test]
1382 fn test_nodemask_from_nodes() {
1383 use super::numa::nodemask_from_nodes;
1384
1385 let (mask, maxnode) = nodemask_from_nodes(&[0], 0).expect("node 0 mask");
1386 assert_eq!(mask[0] & 1, 1, "bit 0 must be set");
1387 assert!(
1388 maxnode as usize >= 64,
1389 "mask must be at least one word wide"
1390 );
1391 assert_eq!(
1392 maxnode as usize,
1393 mask.len() * std::mem::size_of::<libc::c_ulong>() * 8,
1394 "maxnode must describe the whole mask in bits"
1395 );
1396
1397 let (mask, _) = nodemask_from_nodes(&[1, 3], 3).expect("sparse mask");
1398 assert_eq!(mask[0] & 0b1111, 0b1010);
1399
1400 let (mask, maxnode) = nodemask_from_nodes(&[65], 65).expect("wide mask");
1402 assert!(mask.len() >= 2);
1403 assert_eq!(mask[1] & 0b10, 0b10);
1404 assert!((maxnode as usize) > 65);
1405
1406 assert!(nodemask_from_nodes(&[], 0).is_none());
1408 assert!(nodemask_from_nodes(&[-1], 0).is_none());
1409 }
1410
1411 #[test]
1412 fn test_page_aligned_subrange() {
1413 let page = 4096usize;
1414
1415 assert_eq!(
1417 page_aligned_subrange(page, 2 * page, page),
1418 Some((page, 2 * page))
1419 );
1420
1421 assert_eq!(
1423 page_aligned_subrange(page + 100, 3 * page, page),
1424 Some((2 * page, 2 * page))
1425 );
1426
1427 assert_eq!(
1429 page_aligned_subrange(page, page + 7, page),
1430 Some((page, page))
1431 );
1432
1433 assert_eq!(page_aligned_subrange(page + 1, 16, page), None);
1435 assert_eq!(page_aligned_subrange(page, page - 1, page), None);
1436
1437 assert_eq!(page_aligned_subrange(page + 1, page, page), None);
1439
1440 assert_eq!(page_aligned_subrange(0, 0, page), None);
1442 assert_eq!(page_aligned_subrange(page, page, 0), None);
1443 assert_eq!(
1444 page_aligned_subrange(page, page, 4095),
1445 None,
1446 "not a power of two"
1447 );
1448 assert_eq!(page_aligned_subrange(usize::MAX, 4, page), None, "overflow");
1449
1450 let real = page_size();
1452 assert!(real.is_power_of_two() && real >= 4096);
1453 assert!(page_aligned_subrange(real, 4 * real, real).is_some());
1454 }
1455
1456 #[test]
1457 fn test_numa_allocation_shapes_and_contents() {
1458 let allocator = NumaVectorAllocator::new();
1459
1460 for node in [None, Some(0), Some(numa::max_node()), Some(-7), Some(9999)] {
1463 let buf = allocator.allocate_on_node(3 * page_size(), node);
1465 assert_eq!(buf.len(), 3 * page_size());
1466 assert!(buf.iter().all(|&b| b == 0));
1467 }
1468 assert!(allocator.allocate_on_node(0, None).is_empty());
1469
1470 for node in [None, Some(0), Some(-1)] {
1472 let vec = allocator.allocate_vector_on_node(2048, node);
1473 assert_eq!(vec.len(), 2048);
1474 assert!(vec.iter().all(|&v| v == 0.0));
1475 }
1476 assert!(allocator.allocate_vector_on_node(0, None).is_empty());
1477 }
1478
1479 #[test]
1480 fn test_explicit_node_validation_and_round_robin() {
1481 let allocator = NumaVectorAllocator::new();
1482 let valid = allocator.numa_nodes[0];
1483
1484 assert_eq!(allocator.explicit_node(Some(valid)), Some(valid));
1485 assert_eq!(allocator.explicit_node(None), None);
1486 assert_eq!(allocator.explicit_node(Some(-1)), None);
1487 assert_eq!(
1488 allocator.explicit_node(Some(numa::max_node() + 1)),
1489 None,
1490 "out-of-range hints must be rejected, not passed to the kernel"
1491 );
1492
1493 for _ in 0..(allocator.numa_nodes.len() * 3 + 1) {
1495 let node = allocator.next_round_robin_node();
1496 assert!(allocator.numa_nodes.contains(&node));
1497 }
1498 }
1499}