meerkat_mobkit/memory/
guards.rs1use std::collections::HashMap;
16use std::sync::{Arc, Mutex};
17use std::time::{Duration, Instant};
18
19use crate::memory::events::{MemoryEventSink, MemoryTimelineEvent};
20
21pub const DEFAULT_RUNS_PER_HOUR: u32 = 12;
25pub const DEFAULT_MAX_CONCURRENT: u32 = 1;
27
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct BackgroundBudgetConfig {
30 pub runs_per_window: u32,
31 pub window: Duration,
32 pub max_concurrent: u32,
33}
34
35impl Default for BackgroundBudgetConfig {
36 fn default() -> Self {
37 Self {
38 runs_per_window: DEFAULT_RUNS_PER_HOUR,
39 window: Duration::from_hours(1),
40 max_concurrent: DEFAULT_MAX_CONCURRENT,
41 }
42 }
43}
44
45#[derive(Debug, Default)]
46struct RealmBudgetState {
47 starts: Vec<Instant>,
49 concurrent: u32,
50}
51
52struct BudgetInner {
53 config: BackgroundBudgetConfig,
54 realms: HashMap<String, RealmBudgetState>,
55 event_sink: Option<Arc<dyn MemoryEventSink>>,
56}
57
58#[derive(Clone)]
61pub struct BackgroundBudget {
62 inner: Arc<Mutex<BudgetInner>>,
63}
64
65pub struct BudgetPermit {
68 inner: Arc<Mutex<BudgetInner>>,
69 realm: String,
70}
71
72impl std::fmt::Debug for BudgetPermit {
73 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74 f.debug_struct("BudgetPermit")
75 .field("realm", &self.realm)
76 .finish_non_exhaustive()
77 }
78}
79
80impl Drop for BudgetPermit {
81 fn drop(&mut self) {
82 let mut inner = self
83 .inner
84 .lock()
85 .unwrap_or_else(std::sync::PoisonError::into_inner);
86 if let Some(state) = inner.realms.get_mut(&self.realm) {
87 state.concurrent = state.concurrent.saturating_sub(1);
88 }
89 }
90}
91
92#[derive(Debug, Clone, PartialEq, Eq)]
94pub enum BudgetDenied {
95 WindowExhausted { used: u32, cap: u32 },
96 ConcurrencyCeiling { cap: u32 },
97}
98
99impl std::fmt::Display for BudgetDenied {
100 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101 match self {
102 Self::WindowExhausted { used, cap } => {
103 write!(f, "window budget exhausted ({used}/{cap} runs)")
104 }
105 Self::ConcurrencyCeiling { cap } => {
106 write!(f, "concurrency ceiling reached ({cap} in flight)")
107 }
108 }
109 }
110}
111
112impl BackgroundBudget {
113 pub fn new(config: BackgroundBudgetConfig) -> Self {
114 Self {
115 inner: Arc::new(Mutex::new(BudgetInner {
116 config,
117 realms: HashMap::new(),
118 event_sink: None,
119 })),
120 }
121 }
122
123 pub fn set_event_sink(&self, sink: Arc<dyn MemoryEventSink>) {
126 self.inner
127 .lock()
128 .unwrap_or_else(std::sync::PoisonError::into_inner)
129 .event_sink = Some(sink);
130 }
131
132 pub fn try_acquire(&self, realm: &str, stage: &str) -> Result<BudgetPermit, BudgetDenied> {
135 let mut inner = self
136 .inner
137 .lock()
138 .unwrap_or_else(std::sync::PoisonError::into_inner);
139 let window = inner.config.window;
140 let runs_cap = inner.config.runs_per_window;
141 let concurrent_cap = inner.config.max_concurrent;
142 let state = inner.realms.entry(realm.to_string()).or_default();
143 let now = Instant::now();
144 state
145 .starts
146 .retain(|start| now.duration_since(*start) < window);
147 let denied = if state.concurrent >= concurrent_cap {
148 Some(BudgetDenied::ConcurrencyCeiling {
149 cap: concurrent_cap,
150 })
151 } else if state.starts.len() as u32 >= runs_cap {
152 Some(BudgetDenied::WindowExhausted {
153 used: state.starts.len() as u32,
154 cap: runs_cap,
155 })
156 } else {
157 None
158 };
159 if let Some(denied) = denied {
160 tracing::warn!(
161 realm,
162 stage,
163 reason = %denied,
164 "agent memory background budget: run skipped"
165 );
166 if let Some(sink) = inner.event_sink.as_ref() {
167 sink.emit(MemoryTimelineEvent::BudgetDenied {
168 realm: realm.to_string(),
169 stage: stage.to_string(),
170 reason: denied.to_string(),
171 });
172 }
173 return Err(denied);
174 }
175 state.starts.push(now);
176 state.concurrent += 1;
177 Ok(BudgetPermit {
178 inner: self.inner.clone(),
179 realm: realm.to_string(),
180 })
181 }
182}
183
184#[cfg(test)]
185#[allow(clippy::expect_used, clippy::unwrap_used)]
186mod tests {
187 use super::*;
188
189 fn config(runs: u32, concurrent: u32) -> BackgroundBudgetConfig {
190 BackgroundBudgetConfig {
191 runs_per_window: runs,
192 window: Duration::from_hours(1),
193 max_concurrent: concurrent,
194 }
195 }
196
197 #[test]
198 fn window_cap_denies_after_budget_spent() {
199 let budget = BackgroundBudget::new(config(2, 10));
200 let p1 = budget.try_acquire("realm-a", "distiller").expect("run 1");
201 drop(p1);
202 let p2 = budget.try_acquire("realm-a", "distiller").expect("run 2");
203 drop(p2);
204 let denied = budget
205 .try_acquire("realm-a", "distiller")
206 .expect_err("third run in window must deny");
207 assert!(
208 matches!(denied, BudgetDenied::WindowExhausted { used: 2, cap: 2 }),
209 "{denied:?}"
210 );
211 budget
213 .try_acquire("realm-b", "distiller")
214 .expect("other realm has its own window");
215 }
216
217 #[test]
218 fn denial_emits_timeline_event_when_sink_wired() {
219 let budget = BackgroundBudget::new(config(1, 10));
220 let sink = Arc::new(crate::memory::events::CollectingEventSink::new());
221 budget.set_event_sink(sink.clone());
222 let _permit = budget.try_acquire("realm-a", "steward").expect("first");
223 let _ = budget
224 .try_acquire("realm-a", "steward")
225 .expect_err("window spent");
226 assert_eq!(sink.types(), vec!["memory.budget.denied"]);
227 let events = sink.events.lock().unwrap();
228 assert!(matches!(
229 &events[0],
230 MemoryTimelineEvent::BudgetDenied { realm, stage, .. }
231 if realm == "realm-a" && stage == "steward"
232 ));
233 }
234
235 #[test]
236 fn concurrency_ceiling_releases_on_drop() {
237 let budget = BackgroundBudget::new(config(10, 1));
238 let permit = budget.try_acquire("realm-a", "distiller").expect("first");
239 let denied = budget
240 .try_acquire("realm-a", "distiller")
241 .expect_err("second concurrent run must deny");
242 assert!(
243 matches!(denied, BudgetDenied::ConcurrencyCeiling { cap: 1 }),
244 "{denied:?}"
245 );
246 drop(permit);
247 budget
248 .try_acquire("realm-a", "distiller")
249 .expect("slot released on drop");
250 }
251}