impl MemoryManager {
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn new() -> Result<Arc<Self>> {
Self::with_config(MemoryConfig::default())
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn with_config(config: MemoryConfig) -> Result<Arc<Self>> {
let mut pools = FxHashMap::default();
for (&pool_type, &max_size) in &config.pool_limits {
pools.insert(pool_type, MemoryPool::new(max_size));
}
let string_interner = StringInterner::new(
config
.pool_limits
.get(&PoolType::StringIntern)
.copied()
.unwrap_or(16 * 1024 * 1024),
);
Ok(Arc::new(Self {
config,
pools,
string_interner,
total_allocated: Mutex::new(0),
peak_usage: Mutex::new(0),
last_cleanup: Mutex::new(Instant::now()),
}))
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn configure_pool(&self, pool_type: PoolType, _max_size: usize) -> Result<()> {
if let Some(_pool) = self.pools.get(&pool_type) {
warn!("Pool reconfiguration not supported in current implementation");
}
Ok(())
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn allocate_buffer(
self: &Arc<Self>,
pool_type: PoolType,
size: usize,
) -> Result<PooledBuffer> {
let strategy = self.determine_strategy(size);
match strategy {
AllocationStrategy::Pooled => {
if let Some(pool) = self.pools.get(&pool_type) {
let buffer = pool.get_buffer(size);
self.track_allocation(buffer.capacity());
Ok(PooledBuffer::new(buffer, pool_type, Arc::clone(self)))
} else {
Err(anyhow!("Pool type {pool_type:?} not configured"))
}
}
AllocationStrategy::Direct => {
let buffer = vec![0; size];
self.track_allocation(buffer.capacity());
Ok(PooledBuffer::new(buffer, pool_type, Arc::clone(self)))
}
AllocationStrategy::MemoryMapped => {
let buffer = vec![0; size];
self.track_allocation(buffer.capacity());
Ok(PooledBuffer::new(buffer, pool_type, Arc::clone(self)))
}
}
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn intern_string(&self, s: &str) -> Result<Arc<str>> {
self.string_interner.intern(s)
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn stats(&self) -> MemoryStats {
let total_allocated = *self.total_allocated.lock();
let peak_usage = *self.peak_usage.lock();
let mut pool_stats = FxHashMap::default();
for (&pool_type, pool) in &self.pools {
pool_stats.insert(pool_type, pool.stats());
}
let string_intern_size = self.string_interner.memory_usage();
let allocation_pressure = total_allocated as f64 / self.config.max_total_memory as f64;
MemoryStats {
total_allocated,
pool_stats,
string_intern_size,
peak_usage,
allocation_pressure,
}
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn cleanup(&self) -> Result<usize> {
let mut cleaned = 0;
let now = Instant::now();
let mut last_cleanup = self.last_cleanup.lock();
if now.duration_since(*last_cleanup) < Duration::from_secs(30) {
return Ok(0); }
*last_cleanup = now;
let stats = self.stats();
if stats.allocation_pressure < self.config.cache_pressure_threshold {
return Ok(0); }
if stats.allocation_pressure > 0.9 {
self.string_interner.clear();
cleaned += stats.string_intern_size;
info!(
"Cleared string interner: {} bytes",
stats.string_intern_size
);
}
for (pool_type, pool) in &self.pools {
let pool_stats = pool.stats();
if pool_stats.total_size > 0 {
pool.clear();
cleaned += pool_stats.total_size;
debug!(
"Cleared pool {:?}: {} bytes",
pool_type, pool_stats.total_size
);
}
}
if cleaned > 0 {
info!("Memory cleanup freed {} bytes", cleaned);
}
Ok(cleaned)
}
fn determine_strategy(&self, size: usize) -> AllocationStrategy {
if size < self.config.small_allocation_threshold {
AllocationStrategy::Pooled
} else if size > self.config.large_allocation_threshold {
AllocationStrategy::MemoryMapped
} else {
AllocationStrategy::Direct
}
}
fn track_allocation(&self, size: usize) {
let mut total = self.total_allocated.lock();
*total += size;
let mut peak = self.peak_usage.lock();
if *total > *peak {
*peak = *total;
}
if *total as f64 / self.config.max_total_memory as f64
> self.config.cache_pressure_threshold
{
trace!(
"Memory pressure detected: {:.1}%",
*total as f64 / self.config.max_total_memory as f64 * 100.0
);
}
}
fn return_buffer(&self, pool_type: PoolType, buffer: Vec<u8>) {
if let Some(pool) = self.pools.get(&pool_type) {
let capacity = buffer.capacity();
pool.return_buffer(buffer);
let mut total = self.total_allocated.lock();
*total = total.saturating_sub(capacity);
}
}
}
static GLOBAL_MEMORY_MANAGER: std::sync::OnceLock<Arc<MemoryManager>> = std::sync::OnceLock::new();
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn global_memory_manager() -> Result<Arc<MemoryManager>> {
GLOBAL_MEMORY_MANAGER
.get()
.cloned()
.ok_or_else(|| anyhow!("Global memory manager not initialized"))
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn init_global_memory_manager() -> Result<()> {
let manager = MemoryManager::new()?;
GLOBAL_MEMORY_MANAGER
.set(manager)
.map_err(|_| anyhow!("Global memory manager already initialized"))?;
Ok(())
}
#[provable_contracts_macros::contract("pmat-core.yaml", equation = "check_compliance")]
pub fn init_global_memory_manager_with_config(config: MemoryConfig) -> Result<()> {
let manager = MemoryManager::with_config(config)?;
GLOBAL_MEMORY_MANAGER
.set(manager)
.map_err(|_| anyhow!("Global memory manager already initialized"))?;
Ok(())
}