oxirs_vec/
mmap_advanced.rs

1//! Advanced memory mapping features for large datasets
2//!
3//! This module provides advanced memory mapping capabilities including:
4//! - Lazy loading with page-level access
5//! - Smart caching and eviction policies
6//! - NUMA-aware memory allocation
7//! - Swapping policies for memory pressure
8
9use 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::{AtomicU64, AtomicUsize, Ordering};
17use std::sync::Arc;
18use std::time::Instant;
19use tracing::warn;
20
21/// Page size for lazy loading (16KB for better vector alignment)
22const VECTOR_PAGE_SIZE: usize = 16384;
23
24/// Maximum number of pages to keep in memory
25const DEFAULT_MAX_PAGES: usize = 10000;
26
27/// NUMA node information
28#[cfg(target_os = "linux")]
29mod numa {
30    use libc::{c_ulong, c_void};
31
32    extern "C" {
33        fn numa_available() -> i32;
34        fn numa_max_node() -> i32;
35        fn numa_node_of_cpu(cpu: i32) -> i32;
36        fn numa_alloc_onnode(size: usize, node: i32) -> *mut c_void;
37        fn numa_free(ptr: *mut c_void, size: usize);
38        fn mbind(
39            addr: *mut c_void,
40            len: c_ulong,
41            mode: i32,
42            nodemask: *const c_ulong,
43            maxnode: c_ulong,
44            flags: u32,
45        ) -> i32;
46    }
47
48    pub const MPOL_BIND: i32 = 2;
49    pub const MPOL_INTERLEAVE: i32 = 3;
50
51    pub fn is_available() -> bool {
52        unsafe { numa_available() >= 0 }
53    }
54
55    pub fn max_node() -> i32 {
56        unsafe { numa_max_node() }
57    }
58
59    pub fn node_of_cpu(cpu: i32) -> i32 {
60        unsafe { numa_node_of_cpu(cpu) }
61    }
62}
63
64#[cfg(not(target_os = "linux"))]
65mod numa {
66    pub fn is_available() -> bool {
67        false
68    }
69    pub fn max_node() -> i32 {
70        0
71    }
72    pub fn node_of_cpu(_cpu: i32) -> i32 {
73        0
74    }
75}
76
77/// Page access pattern for predictive prefetching
78#[derive(Debug, Clone)]
79struct AccessPattern {
80    page_id: usize,
81    access_time: Instant,
82    access_count: usize,
83}
84
85/// Page cache entry with metadata
86#[derive(Debug)]
87pub struct PageCacheEntry {
88    data: Vec<u8>,
89    page_id: usize,
90    last_access: Instant,
91    access_count: AtomicUsize,
92    dirty: bool,
93    numa_node: i32,
94}
95
96impl PageCacheEntry {
97    /// Get the data slice
98    pub fn data(&self) -> &[u8] {
99        &self.data
100    }
101
102    /// Get the NUMA node
103    pub fn numa_node(&self) -> i32 {
104        self.numa_node
105    }
106}
107
108/// Eviction policy for page cache
109#[derive(Debug, Clone, Copy)]
110pub enum EvictionPolicy {
111    LRU,   // Least Recently Used
112    LFU,   // Least Frequently Used
113    FIFO,  // First In First Out
114    Clock, // Clock algorithm
115    ARC,   // Adaptive Replacement Cache
116}
117
118/// Memory pressure levels
119#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
120pub enum MemoryPressure {
121    Low,
122    Medium,
123    High,
124    Critical,
125}
126
127/// Advanced memory-mapped vector storage
128pub struct AdvancedMemoryMap {
129    /// Base file mapping
130    mmap: Option<Mmap>,
131
132    /// Path to the backing file for dirty-page write-back
133    file_path: Option<std::path::PathBuf>,
134
135    /// Page cache
136    page_cache: Arc<RwLock<LruCache<usize, Arc<PageCacheEntry>>>>,
137
138    /// Access pattern tracking
139    access_patterns: Arc<RwLock<VecDeque<AccessPattern>>>,
140
141    /// Page access frequency
142    page_frequency: Arc<RwLock<HashMap<usize, usize>>>,
143
144    /// Eviction policy
145    eviction_policy: EvictionPolicy,
146
147    /// Memory statistics
148    total_memory: AtomicUsize,
149    cache_hits: AtomicU64,
150    cache_misses: AtomicU64,
151
152    /// NUMA configuration
153    numa_enabled: bool,
154    numa_nodes: Vec<i32>,
155
156    /// Memory pressure monitor
157    memory_pressure: Arc<RwLock<MemoryPressure>>,
158
159    /// Configuration
160    max_pages: usize,
161    page_size: usize,
162    prefetch_distance: usize,
163}
164
165impl AdvancedMemoryMap {
166    /// Create a new advanced memory map
167    pub fn new(mmap: Option<Mmap>, max_pages: usize) -> Self {
168        let numa_enabled = numa::is_available();
169        let numa_nodes = if numa_enabled {
170            (0..=numa::max_node()).collect()
171        } else {
172            vec![0]
173        };
174
175        let cache_size = NonZeroUsize::new(max_pages)
176            .unwrap_or(NonZeroUsize::new(1).expect("constant 1 is non-zero"));
177
178        Self {
179            mmap,
180            file_path: None,
181            page_cache: Arc::new(RwLock::new(LruCache::new(cache_size))),
182            access_patterns: Arc::new(RwLock::new(VecDeque::with_capacity(1000))),
183            page_frequency: Arc::new(RwLock::new(HashMap::new())),
184            eviction_policy: EvictionPolicy::ARC,
185            total_memory: AtomicUsize::new(0),
186            cache_hits: AtomicU64::new(0),
187            cache_misses: AtomicU64::new(0),
188            numa_enabled,
189            numa_nodes,
190            memory_pressure: Arc::new(RwLock::new(MemoryPressure::Low)),
191            max_pages,
192            page_size: VECTOR_PAGE_SIZE,
193            prefetch_distance: 3,
194        }
195    }
196
197    /// Create a new advanced memory map with a backing file path for dirty-page write-back
198    pub fn new_with_path(
199        mmap: Option<Mmap>,
200        max_pages: usize,
201        file_path: Option<std::path::PathBuf>,
202    ) -> Self {
203        let mut s = Self::new(mmap, max_pages);
204        s.file_path = file_path;
205        s
206    }
207
208    /// Get a page with lazy loading
209    pub fn get_page(&self, page_id: usize) -> Result<Arc<PageCacheEntry>> {
210        // Check cache first
211        {
212            let mut cache = self.page_cache.write();
213            if let Some(entry) = cache.get(&page_id) {
214                self.cache_hits.fetch_add(1, Ordering::Relaxed);
215                entry.access_count.fetch_add(1, Ordering::Relaxed);
216                self.record_access(page_id);
217                return Ok(Arc::clone(entry));
218            }
219        }
220
221        // Cache miss - load from mmap
222        self.cache_misses.fetch_add(1, Ordering::Relaxed);
223        self.load_page(page_id)
224    }
225
226    /// Load a page from memory-mapped file
227    fn load_page(&self, page_id: usize) -> Result<Arc<PageCacheEntry>> {
228        let mmap = self
229            .mmap
230            .as_ref()
231            .ok_or_else(|| anyhow::anyhow!("No memory mapping available"))?;
232
233        let start = page_id * self.page_size;
234        let end = (start + self.page_size).min(mmap.len());
235
236        if start >= mmap.len() {
237            bail!("Page {} out of bounds", page_id);
238        }
239
240        // Copy page data
241        let page_data = mmap[start..end].to_vec();
242
243        // Determine NUMA node for allocation
244        let numa_node = if self.numa_enabled {
245            let cpu = sched_getcpu();
246            numa::node_of_cpu(cpu)
247        } else {
248            0
249        };
250
251        let entry = Arc::new(PageCacheEntry {
252            data: page_data,
253            page_id,
254            last_access: Instant::now(),
255            access_count: AtomicUsize::new(1),
256            dirty: false,
257            numa_node,
258        });
259
260        // Check memory pressure and evict if needed
261        self.check_memory_pressure();
262        if *self.memory_pressure.read() >= MemoryPressure::High {
263            self.evict_pages(1)?;
264        }
265
266        // Insert into cache
267        {
268            let mut cache = self.page_cache.write();
269            cache.put(page_id, Arc::clone(&entry));
270        }
271
272        self.total_memory
273            .fetch_add(entry.data.len(), Ordering::Relaxed);
274        self.record_access(page_id);
275
276        // Predictive prefetching
277        self.prefetch_pages(page_id);
278
279        Ok(entry)
280    }
281
282    /// Record page access for pattern analysis
283    fn record_access(&self, page_id: usize) {
284        let mut patterns = self.access_patterns.write();
285        patterns.push_back(AccessPattern {
286            page_id,
287            access_time: Instant::now(),
288            access_count: 1,
289        });
290
291        // Keep only recent patterns
292        while patterns.len() > 1000 {
293            patterns.pop_front();
294        }
295
296        // Update frequency map
297        let mut freq = self.page_frequency.write();
298        *freq.entry(page_id).or_insert(0) += 1;
299    }
300
301    /// Predictive prefetching based on access patterns
302    fn prefetch_pages(&self, current_page: usize) {
303        let patterns = self.access_patterns.read();
304        let freq = self.page_frequency.read();
305
306        // Analyze recent access patterns for intelligent prefetching
307        let recent_patterns: Vec<_> = patterns.iter().rev().take(10).collect();
308
309        // Check for sequential access pattern
310        let is_sequential = recent_patterns
311            .windows(2)
312            .all(|w| w[0].page_id > 0 && w[0].page_id == w[1].page_id + 1);
313
314        // Check for strided access pattern
315        let stride = if recent_patterns.len() >= 3 {
316            let diff1 = recent_patterns[0]
317                .page_id
318                .saturating_sub(recent_patterns[1].page_id);
319            let diff2 = recent_patterns[1]
320                .page_id
321                .saturating_sub(recent_patterns[2].page_id);
322            if diff1 == diff2 && diff1 > 0 && diff1 <= 10 {
323                Some(diff1)
324            } else {
325                None
326            }
327        } else {
328            None
329        };
330
331        // Adaptive prefetching based on patterns
332        if is_sequential {
333            // Aggressive sequential prefetching
334            for i in 1..=(self.prefetch_distance * 2) {
335                let prefetch_page = current_page + i;
336                self.async_prefetch(prefetch_page);
337            }
338        } else if let Some(stride) = stride {
339            // Strided prefetching
340            for i in 1..=self.prefetch_distance {
341                let prefetch_page = current_page + (i * stride);
342                self.async_prefetch(prefetch_page);
343            }
344        } else {
345            // Conservative prefetching with frequency-based hints
346            for i in 1..=self.prefetch_distance {
347                let prefetch_page = current_page + i;
348
349                // Check if this page has been accessed frequently
350                let frequency = *freq.get(&prefetch_page).unwrap_or(&0);
351                if frequency > 0 {
352                    self.async_prefetch(prefetch_page);
353                }
354            }
355        }
356
357        // Prefetch frequently accessed pages near current page
358        let nearby_range = current_page.saturating_sub(3)..=(current_page + 3);
359        for page_id in nearby_range {
360            let frequency = *freq.get(&page_id).unwrap_or(&0);
361            if frequency > 2 && page_id != current_page {
362                self.async_prefetch(page_id);
363            }
364        }
365    }
366
367    /// Asynchronous prefetch with throttling
368    pub fn async_prefetch(&self, page_id: usize) {
369        // Check if page is already in cache
370        {
371            let cache = self.page_cache.read();
372            if cache.contains(&page_id) {
373                return;
374            }
375        }
376
377        // Check memory pressure before prefetching
378        if *self.memory_pressure.read() >= MemoryPressure::High {
379            return;
380        }
381
382        let self_clone = self.clone_ref();
383        spawn(move || {
384            let _ = self_clone.get_page(page_id);
385        });
386    }
387
388    /// Check system memory pressure
389    fn check_memory_pressure(&self) {
390        let total_memory = self.total_memory.load(Ordering::Relaxed);
391        let max_memory = self.max_pages * self.page_size;
392
393        let pressure = if total_memory < max_memory / 2 {
394            MemoryPressure::Low
395        } else if total_memory < max_memory * 3 / 4 {
396            MemoryPressure::Medium
397        } else if total_memory < max_memory * 9 / 10 {
398            MemoryPressure::High
399        } else {
400            MemoryPressure::Critical
401        };
402
403        *self.memory_pressure.write() = pressure;
404    }
405
406    /// Evict pages based on eviction policy
407    fn evict_pages(&self, num_pages: usize) -> Result<()> {
408        match self.eviction_policy {
409            EvictionPolicy::LRU => self.evict_lru(num_pages),
410            EvictionPolicy::LFU => self.evict_lfu(num_pages),
411            EvictionPolicy::FIFO => self.evict_fifo(num_pages),
412            EvictionPolicy::Clock => self.evict_clock(num_pages),
413            EvictionPolicy::ARC => self.evict_arc(num_pages),
414        }
415    }
416
417    /// LRU eviction
418    fn evict_lru(&self, num_pages: usize) -> Result<()> {
419        let mut cache = self.page_cache.write();
420
421        // LruCache automatically evicts least recently used
422        for _ in 0..num_pages {
423            if let Some((_, entry)) = cache.pop_lru() {
424                self.total_memory
425                    .fetch_sub(entry.data.len(), Ordering::Relaxed);
426
427                // Write back if dirty
428                if entry.dirty {
429                    if let Err(e) = self.write_back_page(entry.page_id, &entry.data) {
430                        warn!("Failed to write back page {}: {}", entry.page_id, e);
431                    }
432                }
433            }
434        }
435
436        Ok(())
437    }
438
439    /// LFU eviction
440    fn evict_lfu(&self, num_pages: usize) -> Result<()> {
441        let cache = self.page_cache.read();
442        let freq = self.page_frequency.read();
443
444        // Sort pages by frequency
445        let mut pages_by_freq: Vec<(usize, usize)> = cache
446            .iter()
447            .map(|(page_id, _)| (*page_id, *freq.get(page_id).unwrap_or(&0)))
448            .collect();
449        pages_by_freq.sort_by_key(|(_, freq)| *freq);
450
451        // Evict least frequently used
452        drop(cache);
453        drop(freq);
454
455        let mut cache = self.page_cache.write();
456        for (page_id, _) in pages_by_freq.iter().take(num_pages) {
457            if let Some(entry) = cache.pop(page_id) {
458                self.total_memory
459                    .fetch_sub(entry.data.len(), Ordering::Relaxed);
460                if entry.dirty {
461                    if let Err(e) = self.write_back_page(entry.page_id, &entry.data) {
462                        warn!("Failed to write back dirty page {}: {}", entry.page_id, e);
463                    }
464                }
465            }
466        }
467
468        Ok(())
469    }
470
471    /// FIFO eviction (not implemented - uses LRU as fallback)
472    fn evict_fifo(&self, num_pages: usize) -> Result<()> {
473        self.evict_lru(num_pages)
474    }
475
476    /// Clock algorithm eviction (not implemented - uses LRU as fallback)
477    fn evict_clock(&self, num_pages: usize) -> Result<()> {
478        self.evict_lru(num_pages)
479    }
480
481    /// ARC (Adaptive Replacement Cache) eviction
482    fn evict_arc(&self, num_pages: usize) -> Result<()> {
483        // Simplified ARC - combines recency and frequency
484        let cache = self.page_cache.read();
485        let freq = self.page_frequency.read();
486
487        // Score = recency * 0.5 + frequency * 0.5
488        let now = Instant::now();
489        let mut scored_pages: Vec<(usize, f64)> = cache
490            .iter()
491            .map(|(page_id, entry)| {
492                let recency_score =
493                    1.0 / (now.duration_since(entry.last_access).as_secs_f64() + 1.0);
494                let frequency_score = *freq.get(page_id).unwrap_or(&0) as f64;
495                let combined_score = recency_score * 0.5 + frequency_score * 0.5;
496                (*page_id, combined_score)
497            })
498            .collect();
499
500        scored_pages.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
501
502        drop(cache);
503        drop(freq);
504
505        let mut cache = self.page_cache.write();
506        for (page_id, _) in scored_pages.iter().take(num_pages) {
507            if let Some(entry) = cache.pop(page_id) {
508                self.total_memory
509                    .fetch_sub(entry.data.len(), Ordering::Relaxed);
510                if entry.dirty {
511                    if let Err(e) = self.write_back_page(entry.page_id, &entry.data) {
512                        warn!("Failed to write back dirty page {}: {}", entry.page_id, e);
513                    }
514                }
515            }
516        }
517
518        Ok(())
519    }
520
521    /// Get cache statistics
522    pub fn stats(&self) -> MemoryMapStats {
523        let cache = self.page_cache.read();
524
525        MemoryMapStats {
526            total_pages: cache.len(),
527            total_memory: self.total_memory.load(Ordering::Relaxed),
528            cache_hits: self.cache_hits.load(Ordering::Relaxed),
529            cache_misses: self.cache_misses.load(Ordering::Relaxed),
530            hit_rate: self.calculate_hit_rate(),
531            memory_pressure: *self.memory_pressure.read(),
532            numa_enabled: self.numa_enabled,
533        }
534    }
535
536    fn calculate_hit_rate(&self) -> f64 {
537        let hits = self.cache_hits.load(Ordering::Relaxed) as f64;
538        let misses = self.cache_misses.load(Ordering::Relaxed) as f64;
539        let total = hits + misses;
540        if total > 0.0 {
541            hits / total
542        } else {
543            0.0
544        }
545    }
546
547    fn clone_ref(&self) -> Self {
548        Self {
549            mmap: None, // Don't clone the mmap
550            file_path: self.file_path.clone(),
551            page_cache: Arc::clone(&self.page_cache),
552            access_patterns: Arc::clone(&self.access_patterns),
553            page_frequency: Arc::clone(&self.page_frequency),
554            eviction_policy: self.eviction_policy,
555            total_memory: AtomicUsize::new(0),
556            cache_hits: AtomicU64::new(0),
557            cache_misses: AtomicU64::new(0),
558            numa_enabled: self.numa_enabled,
559            numa_nodes: self.numa_nodes.clone(),
560            memory_pressure: Arc::clone(&self.memory_pressure),
561            max_pages: self.max_pages,
562            page_size: self.page_size,
563            prefetch_distance: self.prefetch_distance,
564        }
565    }
566
567    /// Write a dirty page back to the backing file.
568    fn write_back_page(&self, page_id: usize, data: &[u8]) -> Result<()> {
569        use std::io::{Seek, SeekFrom, Write};
570        let path = match &self.file_path {
571            Some(p) => p,
572            None => return Ok(()), // No file path configured — skip write-back
573        };
574        let mut file = std::fs::OpenOptions::new()
575            .write(true)
576            .open(path)
577            .map_err(|e| anyhow::anyhow!("Failed to open file for write-back: {}", e))?;
578        let offset = (page_id * self.page_size) as u64;
579        file.seek(SeekFrom::Start(offset))
580            .map_err(|e| anyhow::anyhow!("Failed to seek to page {}: {}", page_id, e))?;
581        file.write_all(data)
582            .map_err(|e| anyhow::anyhow!("Failed to write page {}: {}", page_id, e))?;
583        Ok(())
584    }
585
586    /// Flush all dirty pages back to the backing file.
587    pub fn flush_dirty_pages(&self) -> Result<()> {
588        if self.file_path.is_none() {
589            return Ok(());
590        }
591        let cache = self.page_cache.read();
592        for (_, entry) in cache.iter() {
593            if entry.dirty {
594                self.write_back_page(entry.page_id, &entry.data)?;
595            }
596        }
597        Ok(())
598    }
599}
600
601/// Statistics for memory-mapped storage
602#[derive(Debug, Clone)]
603pub struct MemoryMapStats {
604    pub total_pages: usize,
605    pub total_memory: usize,
606    pub cache_hits: u64,
607    pub cache_misses: u64,
608    pub hit_rate: f64,
609    pub memory_pressure: MemoryPressure,
610    pub numa_enabled: bool,
611}
612
613/// Get current CPU for NUMA operations
614#[cfg(target_os = "linux")]
615fn sched_getcpu() -> i32 {
616    unsafe { libc::sched_getcpu() }
617}
618
619#[cfg(not(target_os = "linux"))]
620fn sched_getcpu() -> i32 {
621    0
622}
623
624/// NUMA-aware vector allocator
625pub struct NumaVectorAllocator {
626    numa_nodes: Vec<i32>,
627    current_node: AtomicUsize,
628}
629
630impl Default for NumaVectorAllocator {
631    fn default() -> Self {
632        Self::new()
633    }
634}
635
636impl NumaVectorAllocator {
637    pub fn new() -> Self {
638        let numa_nodes = if numa::is_available() {
639            (0..=numa::max_node()).collect()
640        } else {
641            vec![0]
642        };
643
644        Self {
645            numa_nodes,
646            current_node: AtomicUsize::new(0),
647        }
648    }
649
650    /// Allocate vector memory on specific NUMA node
651    pub fn allocate_on_node(&self, size: usize, node: Option<i32>) -> Vec<u8> {
652        if !numa::is_available() {
653            return vec![0u8; size];
654        }
655
656        let _target_node = node.unwrap_or_else(|| {
657            // Round-robin allocation across NUMA nodes
658            let idx = self.current_node.fetch_add(1, Ordering::Relaxed) % self.numa_nodes.len();
659            self.numa_nodes[idx]
660        });
661
662        // For now, just use standard allocation
663        // TODO: Implement actual NUMA allocation when libc bindings are available
664        vec![0u8; size]
665    }
666
667    /// Allocate optimized vector with NUMA awareness (specialized for f32 vectors)
668    pub fn allocate_vector_on_node(&self, dimensions: usize, node: Option<i32>) -> Vec<f32> {
669        if !numa::is_available() {
670            // Pre-allocate with optimal alignment for SIMD operations
671            let mut vec = Vec::with_capacity(dimensions);
672            vec.resize(dimensions, 0.0f32);
673            return vec;
674        }
675
676        let _target_node = node.unwrap_or_else(|| {
677            // Use current CPU's NUMA node for better locality
678            self.preferred_node()
679        });
680
681        // For better performance, use aligned allocation
682        let mut vec = Vec::with_capacity(dimensions);
683        vec.resize(dimensions, 0.0f32);
684
685        // TODO: When NUMA bindings are available, use numa_alloc_onnode
686        // and bind the memory to the specific node
687
688        vec
689    }
690
691    /// Get preferred NUMA node for current thread
692    pub fn preferred_node(&self) -> i32 {
693        if numa::is_available() {
694            numa::node_of_cpu(sched_getcpu())
695        } else {
696            0
697        }
698    }
699}
700
701#[cfg(test)]
702mod tests {
703    use super::*;
704
705    #[test]
706    fn test_memory_pressure() {
707        let mmap = AdvancedMemoryMap::new(None, 100);
708
709        assert_eq!(*mmap.memory_pressure.read(), MemoryPressure::Low);
710
711        // Simulate memory usage
712        mmap.total_memory
713            .store(50 * VECTOR_PAGE_SIZE, Ordering::Relaxed);
714        mmap.check_memory_pressure();
715        assert_eq!(*mmap.memory_pressure.read(), MemoryPressure::Medium);
716
717        mmap.total_memory
718            .store(90 * VECTOR_PAGE_SIZE, Ordering::Relaxed);
719        mmap.check_memory_pressure();
720        assert_eq!(*mmap.memory_pressure.read(), MemoryPressure::Critical);
721    }
722
723    #[test]
724    fn test_cache_stats() {
725        let mmap = AdvancedMemoryMap::new(None, 100);
726
727        mmap.cache_hits.store(75, Ordering::Relaxed);
728        mmap.cache_misses.store(25, Ordering::Relaxed);
729
730        let stats = mmap.stats();
731        assert_eq!(stats.cache_hits, 75);
732        assert_eq!(stats.cache_misses, 25);
733        assert_eq!(stats.hit_rate, 0.75);
734    }
735}
oxirs_vec/mmap_advanced.rs

oxirs_vec/
mmap_advanced.rs
