use crate::types::AttentionConfig;
use crate::types::*;
use dashmap::DashMap;
pub struct AttentionAllocator {
budgets: DashMap<ResonatorId, AttentionBudget>,
#[allow(dead_code)]
config: AttentionConfig,
}
impl AttentionAllocator {
pub fn new(config: &AttentionConfig) -> Self {
Self {
budgets: DashMap::new(),
config: config.clone(),
}
}
pub async fn allocate_budget(
&self,
resonator: &ResonatorId,
spec: &AttentionBudgetSpec,
) -> Result<(), String> {
let mut budget = AttentionBudget::new(spec.total_capacity);
if let Some(reserve) = spec.safety_reserve {
budget.safety_reserve = reserve;
}
if let Some(threshold) = spec.exhaustion_threshold {
budget.exhaustion_threshold = threshold;
}
self.budgets.insert(*resonator, budget);
tracing::debug!(
"Allocated attention budget for {}: {} total",
resonator,
spec.total_capacity
);
Ok(())
}
pub async fn allocate(
&self,
resonator: &ResonatorId,
amount: u64,
) -> Result<AllocationToken, AttentionError> {
let budget = self
.budgets
.get_mut(resonator)
.ok_or(AttentionError::ResonatorNotFound)?;
let available = budget.available();
if amount > available {
return Err(AttentionError::InsufficientAttention {
requested: amount,
available,
});
}
let token = AllocationToken::new(*resonator, amount);
tracing::trace!(
"Allocated {} attention for {} (available: {})",
amount,
resonator,
available - amount
);
Ok(token)
}
pub async fn allocate_more(
&self,
resonator: &ResonatorId,
amount: u64,
) -> Result<(), AttentionError> {
let budget = self
.budgets
.get_mut(resonator)
.ok_or(AttentionError::ResonatorNotFound)?;
let available = budget.available();
if amount > available {
return Err(AttentionError::InsufficientAttention {
requested: amount,
available,
});
}
Ok(())
}
pub async fn release_partial(&self, resonator: &ResonatorId, amount: u64) {
if let Some(_budget) = self.budgets.get_mut(resonator) {
tracing::trace!("Released {} attention for {}", amount, resonator);
}
}
pub async fn release_all(&self, resonator: &ResonatorId, amount: u64) {
if let Some(_budget) = self.budgets.get_mut(resonator) {
tracing::trace!("Released all {} attention for {}", amount, resonator);
}
}
pub fn is_exhaustion_imminent(&self, resonator: &ResonatorId) -> bool {
if let Some(budget) = self.budgets.get(resonator) {
return budget.is_exhaustion_imminent();
}
false
}
pub async fn available_for_coupling(
&self,
resonator: &ResonatorId,
) -> Result<u64, AttentionError> {
let budget = self
.budgets
.get(resonator)
.ok_or(AttentionError::ResonatorNotFound)?;
Ok(budget.available())
}
pub async fn get_budget(&self, resonator: &ResonatorId) -> Option<AttentionBudget> {
self.budgets.get(resonator).map(|r| r.clone())
}
pub async fn restore_budget(
&self,
resonator: &ResonatorId,
budget: &AttentionBudget,
) -> Result<(), String> {
self.budgets.insert(*resonator, budget.clone());
tracing::debug!("Restored attention budget for {}", resonator);
Ok(())
}
pub async fn rebalance(&self, _resonator: &ResonatorId) {
tracing::debug!("Rebalancing attention (placeholder)");
}
pub fn remove_budget(&self, resonator: &ResonatorId) {
self.budgets.remove(resonator);
}
}