maple_runtime/allocator/
mod.rs1use dashmap::DashMap;
6use crate::types::*;
7use crate::types::AttentionConfig;
8
9pub struct AttentionAllocator {
13 budgets: DashMap<ResonatorId, AttentionBudget>,
15
16 config: AttentionConfig,
18}
19
20impl AttentionAllocator {
21 pub fn new(config: &AttentionConfig) -> Self {
22 Self {
23 budgets: DashMap::new(),
24 config: config.clone(),
25 }
26 }
27
28 pub async fn allocate_budget(
30 &self,
31 resonator: &ResonatorId,
32 spec: &AttentionBudgetSpec,
33 ) -> Result<(), String> {
34 let mut budget = AttentionBudget::new(spec.total_capacity);
35
36 if let Some(reserve) = spec.safety_reserve {
37 budget.safety_reserve = reserve;
38 }
39
40 if let Some(threshold) = spec.exhaustion_threshold {
41 budget.exhaustion_threshold = threshold;
42 }
43
44 self.budgets.insert(*resonator, budget);
45
46 tracing::debug!(
47 "Allocated attention budget for {}: {} total",
48 resonator,
49 spec.total_capacity
50 );
51
52 Ok(())
53 }
54
55 pub async fn allocate(
59 &self,
60 resonator: &ResonatorId,
61 amount: u64,
62 ) -> Result<AllocationToken, AttentionError> {
63 let mut budget = self
64 .budgets
65 .get_mut(resonator)
66 .ok_or(AttentionError::ResonatorNotFound)?;
67
68 let available = budget.available();
69
70 if amount > available {
71 return Err(AttentionError::InsufficientAttention {
72 requested: amount,
73 available,
74 });
75 }
76
77 let token = AllocationToken::new(*resonator, amount);
78
79 tracing::trace!(
82 "Allocated {} attention for {} (available: {})",
83 amount,
84 resonator,
85 available - amount
86 );
87
88 Ok(token)
89 }
90
91 pub async fn allocate_more(
93 &self,
94 resonator: &ResonatorId,
95 amount: u64,
96 ) -> Result<(), AttentionError> {
97 let mut budget = self
98 .budgets
99 .get_mut(resonator)
100 .ok_or(AttentionError::ResonatorNotFound)?;
101
102 let available = budget.available();
103
104 if amount > available {
105 return Err(AttentionError::InsufficientAttention {
106 requested: amount,
107 available,
108 });
109 }
110
111 Ok(())
114 }
115
116 pub async fn release_partial(&self, resonator: &ResonatorId, amount: u64) {
118 if let Some(_budget) = self.budgets.get_mut(resonator) {
119 tracing::trace!("Released {} attention for {}", amount, resonator);
121 }
122 }
123
124 pub async fn release_all(&self, resonator: &ResonatorId, amount: u64) {
126 if let Some(_budget) = self.budgets.get_mut(resonator) {
127 tracing::trace!("Released all {} attention for {}", amount, resonator);
129 }
130 }
131
132 pub fn is_exhaustion_imminent(&self, resonator: &ResonatorId) -> bool {
134 if let Some(budget) = self.budgets.get(resonator) {
135 return budget.is_exhaustion_imminent();
136 }
137 false
138 }
139
140 pub async fn available_for_coupling(
142 &self,
143 resonator: &ResonatorId,
144 ) -> Result<u64, AttentionError> {
145 let budget = self
146 .budgets
147 .get(resonator)
148 .ok_or(AttentionError::ResonatorNotFound)?;
149
150 Ok(budget.available())
151 }
152
153 pub async fn get_budget(&self, resonator: &ResonatorId) -> Option<AttentionBudget> {
155 self.budgets.get(resonator).map(|r| r.clone())
156 }
157
158 pub async fn restore_budget(
160 &self,
161 resonator: &ResonatorId,
162 budget: &AttentionBudget,
163 ) -> Result<(), String> {
164 self.budgets.insert(*resonator, budget.clone());
165 tracing::debug!("Restored attention budget for {}", resonator);
166 Ok(())
167 }
168
169 pub async fn rebalance(&self, _resonator: &ResonatorId) {
173 tracing::debug!("Rebalancing attention (placeholder)");
178 }
179
180 pub fn remove_budget(&self, resonator: &ResonatorId) {
182 self.budgets.remove(resonator);
183 }
184}