use crate::error::Result;
use crate::storage::page::{PageId, PAGE_SIZE};
use parking_lot::RwLock;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
#[derive(Debug, Clone)]
pub struct MemoryOptimizationConfig {
pub enable_prefetch: bool,
pub prefetch_pages: usize,
pub enable_memory_pool: bool,
pub pool_size: usize,
pub readahead_strategy: ReadaheadStrategy,
}
impl Default for MemoryOptimizationConfig {
fn default() -> Self {
Self {
enable_prefetch: true,
prefetch_pages: 8,
enable_memory_pool: true,
pool_size: 256,
readahead_strategy: ReadaheadStrategy::Adaptive,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReadaheadStrategy {
None,
Fixed,
Adaptive,
Aggressive,
}
#[derive(Debug)]
struct AccessPatternDetector {
recent_accesses: RwLock<VecDeque<PageId>>,
sequential_count: AtomicU64,
random_count: AtomicU64,
history_size: usize,
}
impl AccessPatternDetector {
fn new(history_size: usize) -> Self {
Self {
recent_accesses: RwLock::new(VecDeque::with_capacity(history_size)),
sequential_count: AtomicU64::new(0),
random_count: AtomicU64::new(0),
history_size,
}
}
fn record_access(&self, page_id: PageId) {
let mut accesses = self.recent_accesses.write();
if let Some(&last_page) = accesses.back() {
if page_id == last_page + 1 || page_id == last_page.saturating_sub(1) {
self.sequential_count.fetch_add(1, Ordering::Relaxed);
} else {
self.random_count.fetch_add(1, Ordering::Relaxed);
}
}
accesses.push_back(page_id);
while accesses.len() > self.history_size {
accesses.pop_front();
}
}
fn is_sequential(&self) -> bool {
let sequential = self.sequential_count.load(Ordering::Relaxed);
let random = self.random_count.load(Ordering::Relaxed);
let total = sequential + random;
if total == 0 {
false
} else {
(sequential as f64 / total as f64) > 0.7
}
}
fn recommended_prefetch_distance(&self) -> usize {
if self.is_sequential() {
16
} else {
4
}
}
fn reset(&self) {
self.recent_accesses.write().clear();
self.sequential_count.store(0, Ordering::Relaxed);
self.random_count.store(0, Ordering::Relaxed);
}
}
#[derive(Debug)]
struct MemoryPool {
buffers: RwLock<Vec<Vec<u8>>>,
capacity: usize,
allocations: AtomicU64,
pool_hits: AtomicU64,
pool_misses: AtomicU64,
}
impl MemoryPool {
fn new(capacity: usize) -> Self {
let mut buffers = Vec::with_capacity(capacity);
for _ in 0..capacity {
let buffer = vec![0; PAGE_SIZE];
buffers.push(buffer);
}
Self {
buffers: RwLock::new(buffers),
capacity,
allocations: AtomicU64::new(0),
pool_hits: AtomicU64::new(0),
pool_misses: AtomicU64::new(0),
}
}
fn acquire(&self) -> Vec<u8> {
self.allocations.fetch_add(1, Ordering::Relaxed);
let mut buffers = self.buffers.write();
if let Some(buffer) = buffers.pop() {
self.pool_hits.fetch_add(1, Ordering::Relaxed);
buffer
} else {
self.pool_misses.fetch_add(1, Ordering::Relaxed);
vec![0; PAGE_SIZE]
}
}
fn release(&self, mut buffer: Vec<u8>) {
buffer.fill(0);
let mut buffers = self.buffers.write();
if buffers.len() < self.capacity {
buffers.push(buffer);
}
}
fn hit_rate(&self) -> f64 {
let hits = self.pool_hits.load(Ordering::Relaxed) as f64;
let total = self.allocations.load(Ordering::Relaxed) as f64;
if total == 0.0 {
0.0
} else {
hits / total
}
}
fn stats(&self) -> MemoryPoolStats {
MemoryPoolStats {
capacity: self.capacity,
available: self.buffers.read().len(),
allocations: self.allocations.load(Ordering::Relaxed),
hits: self.pool_hits.load(Ordering::Relaxed),
misses: self.pool_misses.load(Ordering::Relaxed),
hit_rate: self.hit_rate(),
}
}
}
#[derive(Debug, Clone)]
pub struct MemoryPoolStats {
pub capacity: usize,
pub available: usize,
pub allocations: u64,
pub hits: u64,
pub misses: u64,
pub hit_rate: f64,
}
pub struct MemoryOptimizer {
config: MemoryOptimizationConfig,
pattern_detector: Arc<AccessPatternDetector>,
memory_pool: Option<Arc<MemoryPool>>,
stats: MemoryOptimizerStats,
}
#[derive(Debug, Default)]
pub struct MemoryOptimizerStats {
pub prefetch_requests: AtomicU64,
pub prefetch_hits: AtomicU64,
pub sequential_accesses: AtomicU64,
pub random_accesses: AtomicU64,
}
impl MemoryOptimizerStats {
pub fn prefetch_hit_rate(&self) -> f64 {
let requests = self.prefetch_requests.load(Ordering::Relaxed) as f64;
if requests == 0.0 {
0.0
} else {
let hits = self.prefetch_hits.load(Ordering::Relaxed) as f64;
hits / requests
}
}
pub fn sequential_percentage(&self) -> f64 {
let sequential = self.sequential_accesses.load(Ordering::Relaxed) as f64;
let random = self.random_accesses.load(Ordering::Relaxed) as f64;
let total = sequential + random;
if total == 0.0 {
0.0
} else {
(sequential / total) * 100.0
}
}
}
impl MemoryOptimizer {
pub fn new(config: MemoryOptimizationConfig) -> Self {
let pattern_detector = Arc::new(AccessPatternDetector::new(100));
let memory_pool = if config.enable_memory_pool {
Some(Arc::new(MemoryPool::new(config.pool_size)))
} else {
None
};
Self {
config,
pattern_detector,
memory_pool,
stats: MemoryOptimizerStats::default(),
}
}
pub fn record_access(&self, page_id: PageId) {
self.pattern_detector.record_access(page_id);
if self.pattern_detector.is_sequential() {
self.stats
.sequential_accesses
.fetch_add(1, Ordering::Relaxed);
} else {
self.stats.random_accesses.fetch_add(1, Ordering::Relaxed);
}
}
pub fn get_prefetch_recommendations(&self, current_page: PageId) -> Vec<PageId> {
if !self.config.enable_prefetch {
return Vec::new();
}
let prefetch_distance = match self.config.readahead_strategy {
ReadaheadStrategy::None => return Vec::new(),
ReadaheadStrategy::Fixed => self.config.prefetch_pages,
ReadaheadStrategy::Adaptive => self.pattern_detector.recommended_prefetch_distance(),
ReadaheadStrategy::Aggressive => 32, };
let mut recommendations = Vec::with_capacity(prefetch_distance);
for i in 1..=prefetch_distance {
recommendations.push(current_page + i as u64);
}
self.stats
.prefetch_requests
.fetch_add(prefetch_distance as u64, Ordering::Relaxed);
recommendations
}
pub fn acquire_buffer(&self) -> Vec<u8> {
if let Some(ref pool) = self.memory_pool {
pool.acquire()
} else {
vec![0; PAGE_SIZE]
}
}
pub fn release_buffer(&self, buffer: Vec<u8>) {
if let Some(ref pool) = self.memory_pool {
pool.release(buffer);
}
}
pub fn pool_stats(&self) -> Option<MemoryPoolStats> {
self.memory_pool.as_ref().map(|pool| pool.stats())
}
pub fn stats(&self) -> &MemoryOptimizerStats {
&self.stats
}
pub fn is_sequential_access(&self) -> bool {
self.pattern_detector.is_sequential()
}
pub fn reset_pattern_detection(&self) {
self.pattern_detector.reset();
self.stats.sequential_accesses.store(0, Ordering::Relaxed);
self.stats.random_accesses.store(0, Ordering::Relaxed);
}
pub fn recommended_readahead(&self) -> usize {
self.pattern_detector.recommended_prefetch_distance()
}
}
impl Default for MemoryOptimizer {
fn default() -> Self {
Self::new(MemoryOptimizationConfig::default())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_memory_optimizer_creation() {
let config = MemoryOptimizationConfig::default();
let optimizer = MemoryOptimizer::new(config);
assert!(!optimizer.is_sequential_access());
}
#[test]
fn test_sequential_access_detection() {
let optimizer = MemoryOptimizer::default();
for page_id in 0..20 {
optimizer.record_access(page_id);
}
assert!(optimizer.is_sequential_access());
assert!(optimizer.stats().sequential_percentage() > 50.0);
}
#[test]
fn test_random_access_detection() {
let optimizer = MemoryOptimizer::default();
let random_pages = [5, 100, 3, 200, 7, 150, 9, 250];
for &page_id in &random_pages {
optimizer.record_access(page_id);
}
assert!(!optimizer.is_sequential_access());
}
#[test]
fn test_prefetch_recommendations() {
let config = MemoryOptimizationConfig {
enable_prefetch: true,
prefetch_pages: 4,
readahead_strategy: ReadaheadStrategy::Fixed,
..Default::default()
};
let optimizer = MemoryOptimizer::new(config);
let recommendations = optimizer.get_prefetch_recommendations(100);
assert_eq!(recommendations.len(), 4);
assert_eq!(recommendations[0], 101);
assert_eq!(recommendations[1], 102);
assert_eq!(recommendations[2], 103);
assert_eq!(recommendations[3], 104);
}
#[test]
fn test_adaptive_prefetch() {
let config = MemoryOptimizationConfig {
readahead_strategy: ReadaheadStrategy::Adaptive,
..Default::default()
};
let optimizer = MemoryOptimizer::new(config);
for page_id in 0..10 {
optimizer.record_access(page_id);
}
let recommendations = optimizer.get_prefetch_recommendations(10);
assert!(recommendations.len() >= 4);
}
#[test]
fn test_memory_pool_acquire_release() {
let config = MemoryOptimizationConfig {
enable_memory_pool: true,
pool_size: 10,
..Default::default()
};
let optimizer = MemoryOptimizer::new(config);
let buffer1 = optimizer.acquire_buffer();
assert_eq!(buffer1.len(), PAGE_SIZE);
optimizer.release_buffer(buffer1);
let buffer2 = optimizer.acquire_buffer();
assert_eq!(buffer2.len(), PAGE_SIZE);
let stats = optimizer.pool_stats().unwrap();
assert!(stats.allocations >= 2);
assert!(stats.hit_rate > 0.0);
}
#[test]
fn test_memory_pool_statistics() {
let config = MemoryOptimizationConfig {
enable_memory_pool: true,
pool_size: 5,
..Default::default()
};
let optimizer = MemoryOptimizer::new(config);
for _ in 0..10 {
let buffer = optimizer.acquire_buffer();
optimizer.release_buffer(buffer);
}
let stats = optimizer.pool_stats().unwrap();
assert_eq!(stats.capacity, 5);
assert!(stats.hits > 0);
assert!(stats.hit_rate > 0.0);
}
#[test]
fn test_disabled_prefetch() {
let config = MemoryOptimizationConfig {
enable_prefetch: false,
..Default::default()
};
let optimizer = MemoryOptimizer::new(config);
let recommendations = optimizer.get_prefetch_recommendations(10);
assert!(recommendations.is_empty());
}
#[test]
fn test_aggressive_readahead() {
let config = MemoryOptimizationConfig {
readahead_strategy: ReadaheadStrategy::Aggressive,
..Default::default()
};
let optimizer = MemoryOptimizer::new(config);
let recommendations = optimizer.get_prefetch_recommendations(10);
assert_eq!(recommendations.len(), 32);
}
#[test]
fn test_pattern_reset() {
let optimizer = MemoryOptimizer::default();
for page_id in 0..10 {
optimizer.record_access(page_id);
}
assert!(optimizer.is_sequential_access());
optimizer.reset_pattern_detection();
assert_eq!(optimizer.stats().sequential_percentage(), 0.0);
}
#[test]
fn test_recommended_readahead_distance() {
let optimizer = MemoryOptimizer::default();
for page_id in 0..20 {
optimizer.record_access(page_id);
}
let distance = optimizer.recommended_readahead();
assert!(distance >= 8); }
#[test]
fn test_memory_pool_capacity_limit() {
let config = MemoryOptimizationConfig {
enable_memory_pool: true,
pool_size: 2,
..Default::default()
};
let optimizer = MemoryOptimizer::new(config);
let buf1 = optimizer.acquire_buffer();
let buf2 = optimizer.acquire_buffer();
let buf3 = optimizer.acquire_buffer();
optimizer.release_buffer(buf1);
optimizer.release_buffer(buf2);
optimizer.release_buffer(buf3);
let stats = optimizer.pool_stats().unwrap();
assert_eq!(stats.available, 2); }
#[test]
fn test_prefetch_statistics() {
let optimizer = MemoryOptimizer::default();
for i in 0..5 {
optimizer.get_prefetch_recommendations(i * 10);
}
let stats = optimizer.stats();
assert!(stats.prefetch_requests.load(Ordering::Relaxed) > 0);
}
#[test]
fn test_mixed_access_pattern() {
let optimizer = MemoryOptimizer::default();
for i in 0..5 {
optimizer.record_access(i);
}
optimizer.record_access(100);
optimizer.record_access(200);
let stats = optimizer.stats();
assert!(stats.sequential_accesses.load(Ordering::Relaxed) > 0);
assert!(stats.random_accesses.load(Ordering::Relaxed) > 0);
}
}