use crate::error::{Result, VantaError};
use crate::query::LogicalPlan;
use std::sync::atomic::{AtomicUsize, Ordering};
pub static ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);
#[derive(Debug, Clone, PartialEq)]
pub enum AllocationStatus {
Granted,
GrantedWithPressure, }
pub struct ResourceGovernor {
pub max_memory_bytes: usize,
pub query_timeout_ms: u64,
}
impl ResourceGovernor {
pub fn new(max_memory_bytes: usize, query_timeout_ms: u64) -> Self {
Self {
max_memory_bytes,
query_timeout_ms,
}
}
pub fn request_allocation(&self, bytes: usize) -> Result<AllocationStatus> {
let current = ALLOCATED_BYTES.load(Ordering::Relaxed);
let new_total = current + bytes;
if new_total > self.max_memory_bytes {
return Err(VantaError::ResourceLimit(
"OOM Guard triggered: query exceeds soft memory limit.".to_string(),
));
}
let pressure_threshold = (self.max_memory_bytes as f64 * 0.9) as usize;
let status = if new_total > pressure_threshold {
AllocationStatus::GrantedWithPressure
} else {
AllocationStatus::Granted
};
ALLOCATED_BYTES.fetch_add(bytes, Ordering::SeqCst);
Ok(status)
}
pub fn free_allocation(&self, bytes: usize) {
ALLOCATED_BYTES.fetch_sub(bytes, Ordering::SeqCst);
}
pub fn apply_temperature_limits(&self, plan: &mut LogicalPlan) {
if plan.temperature > 0.8 {
for op in plan.operators.iter_mut() {
if let crate::query::LogicalOperator::Traverse { max_depth, .. } = op {
if *max_depth > 3 {
*max_depth = 3; }
}
}
}
}
}