use crate::error::Result;
use super::super::resource::allocation::SystemMetrics;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MemoryStrategy {
Standard,
MemoryPool { pool_size: usize },
MemoryMapped { chunk_size: usize },
Streaming { buffer_size: usize },
Hybrid {
small_data_threshold: usize,
memory_pool_size: usize,
streaming_threshold: usize,
},
}
#[derive(Debug)]
pub struct MemoryPoolManager {
pools: HashMap<usize, MemoryPool>,
allocation_strategy: AllocationStrategy,
fragmentation_monitor: FragmentationMonitor,
}
impl MemoryPoolManager {
pub fn new() -> Self {
Self {
pools: HashMap::new(),
allocation_strategy: AllocationStrategy::BestFit,
fragmentation_monitor: FragmentationMonitor::new(),
}
}
pub fn determine_optimal_strategy(
&mut self,
data_size: usize,
system_metrics: &SystemMetrics,
) -> Result<MemoryStrategy> {
let available_memory = system_metrics.memory_usage.available as usize;
let memory_pressure = system_metrics.memory_usage.utilization;
if data_size > available_memory / 2 {
if memory_pressure > 0.8 {
Ok(MemoryStrategy::Streaming {
buffer_size: available_memory / 10,
})
} else {
Ok(MemoryStrategy::MemoryMapped {
chunk_size: available_memory / 4,
})
}
} else if data_size > 1024 * 1024 {
Ok(MemoryStrategy::MemoryPool {
pool_size: data_size * 2,
})
} else {
Ok(MemoryStrategy::Standard)
}
}
pub fn allocate_pool(&mut self, pool_id: usize, pool_size: usize) -> Result<()> {
if !self.pools.contains_key(&pool_id) {
let pool = MemoryPool {
pool_size,
allocated: 0,
free_blocks: vec![MemoryBlock {
address: 0, size: pool_size,
is_free: true,
}],
};
self.pools.insert(pool_id, pool);
}
Ok(())
}
pub fn get_pool_stats(&self, pool_id: usize) -> Option<MemoryPoolStats> {
self.pools.get(&pool_id).map(|pool| MemoryPoolStats {
pool_size: pool.pool_size,
allocated: pool.allocated,
free_blocks: pool.free_blocks.len(),
fragmentation: self.fragmentation_monitor.external_fragmentation,
})
}
pub fn compact_if_needed(&mut self) -> Result<bool> {
if self.fragmentation_monitor.external_fragmentation > self.fragmentation_monitor.compaction_threshold {
self.compact_memory()?;
Ok(true)
} else {
Ok(false)
}
}
fn compact_memory(&mut self) -> Result<()> {
for pool in self.pools.values_mut() {
pool.free_blocks.retain(|block| !block.is_free || block.size > 0);
}
self.fragmentation_monitor.external_fragmentation *= 0.5; Ok(())
}
}
impl Default for MemoryPoolManager {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct MemoryPool {
pub pool_size: usize,
pub allocated: usize,
pub free_blocks: Vec<MemoryBlock>,
}
#[derive(Debug)]
pub struct MemoryBlock {
pub address: usize,
pub size: usize,
pub is_free: bool,
}
#[derive(Debug)]
pub enum AllocationStrategy {
FirstFit,
BestFit,
WorstFit,
NextFit,
}
#[derive(Debug)]
pub struct FragmentationMonitor {
pub internal_fragmentation: f64,
pub external_fragmentation: f64,
pub compaction_threshold: f64,
}
impl Default for FragmentationMonitor {
fn default() -> Self {
Self::new()
}
}
impl FragmentationMonitor {
pub fn new() -> Self {
Self {
internal_fragmentation: 0.0,
external_fragmentation: 0.0,
compaction_threshold: 0.3, }
}
pub fn update_fragmentation(&mut self, internal: f64, external: f64) {
self.internal_fragmentation = internal;
self.external_fragmentation = external;
}
pub fn needs_compaction(&self) -> bool {
self.external_fragmentation > self.compaction_threshold
}
}
#[derive(Debug, Clone)]
pub struct MemoryPoolStats {
pub pool_size: usize,
pub allocated: usize,
pub free_blocks: usize,
pub fragmentation: f64,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pipeline::optimization::resource::allocation::{SystemMetrics, MemoryUsage};
#[test]
fn test_memory_pool_manager_creation() {
let manager = MemoryPoolManager::new();
assert!(manager.pools.is_empty());
}
#[test]
fn test_optimal_strategy_small_data() {
let mut manager = MemoryPoolManager::new();
let system_metrics = SystemMetrics {
memory_usage: MemoryUsage {
total: 8 * 1024 * 1024 * 1024,
available: 4 * 1024 * 1024 * 1024,
used: 4 * 1024 * 1024 * 1024,
utilization: 0.5,
},
..Default::default()
};
let strategy = manager.determine_optimal_strategy(1024, &system_metrics).expect("Operation failed");
matches!(strategy, MemoryStrategy::Standard);
}
#[test]
fn test_memory_pool_allocation() {
let mut manager = MemoryPoolManager::new();
let result = manager.allocate_pool(0, 1024 * 1024);
assert!(result.is_ok());
let stats = manager.get_pool_stats(0);
assert!(stats.is_some());
assert_eq!(stats.expect("Operation failed").pool_size, 1024 * 1024);
}
#[test]
fn test_fragmentation_monitor() {
let mut monitor = FragmentationMonitor::new();
assert_eq!(monitor.compaction_threshold, 0.3);
monitor.update_fragmentation(0.1, 0.4);
assert!(monitor.needs_compaction());
}
}