1use std::collections::hash_map::Entry;
2use std::collections::HashMap;
3use std::sync::Mutex;
4
5use async_trait::async_trait;
6
7use crate::types::{Approval, ApprovalStatus, LogEntry, Mission, Person};
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum Spend {
12 WithinBudget,
14 Approved,
16}
17
18#[derive(Debug)]
21pub struct StoreError(Box<dyn std::error::Error + Send + Sync>);
22
23impl StoreError {
24 pub fn new(error: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> Self {
26 Self(error.into())
27 }
28}
29
30impl std::fmt::Display for StoreError {
31 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32 write!(f, "store unavailable: {}", self.0)
33 }
34}
35
36impl std::error::Error for StoreError {
37 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
38 Some(self.0.as_ref())
39 }
40}
41
42pub type StoreResult<T> = Result<T, StoreError>;
44
45#[async_trait]
61pub trait Store: Send + Sync {
62 async fn binding(&self, tenant: &str, jkt: &str) -> StoreResult<Option<Person>>;
64 async fn bind(&self, tenant: &str, jkt: &str, person: &Person) -> StoreResult<Person>;
66 async fn mission(&self, tenant: &str, s256: &str) -> StoreResult<Option<Mission>>;
68 async fn insert_mission(&self, tenant: &str, mission: &Mission) -> StoreResult<bool>;
71 async fn log_mission(
73 &self,
74 tenant: &str,
75 s256: &str,
76 entry: LogEntry,
77 now: i64,
78 ) -> StoreResult<bool>;
79 async fn end_mission(&self, tenant: &str, s256: &str, at: i64) -> StoreResult<Option<Mission>>;
82 async fn spend(
85 &self,
86 tenant: &str,
87 s256: &str,
88 amount: u64,
89 rule: Spend,
90 now: i64,
91 ) -> StoreResult<Option<Mission>>;
92 async fn missions(&self, tenant: &str) -> StoreResult<Vec<Mission>>;
94 async fn approval(&self, tenant: &str, id: &str) -> StoreResult<Option<Approval>>;
96 async fn insert_approval(&self, tenant: &str, approval: &Approval) -> StoreResult<()>;
98 async fn transition_approval(
101 &self,
102 tenant: &str,
103 expected: ApprovalStatus,
104 next: &Approval,
105 ) -> StoreResult<bool>;
106 async fn approvals(&self, tenant: &str) -> StoreResult<Vec<Approval>>;
108}
109
110type Key = (String, String);
112
113fn key(tenant: &str, id: &str) -> Key {
114 (tenant.to_string(), id.to_string())
115}
116
117#[derive(Default)]
118struct State {
119 bindings: HashMap<Key, Person>,
120 missions: HashMap<Key, Mission>,
121 approvals: HashMap<Key, Approval>,
122}
123
124const APPROVAL_RETENTION_SECS: i64 = 24 * 60 * 60;
127
128#[derive(Default)]
132pub struct MemoryStore(Mutex<State>);
133
134impl MemoryStore {
135 pub fn new() -> Self {
137 Self::default()
138 }
139
140 fn with<T>(&self, f: impl FnOnce(&mut State) -> T) -> StoreResult<T> {
141 let mut state = self
142 .0
143 .lock()
144 .unwrap_or_else(|poisoned| poisoned.into_inner());
145 Ok(f(&mut state))
146 }
147}
148
149#[async_trait]
150impl Store for MemoryStore {
151 async fn binding(&self, tenant: &str, jkt: &str) -> StoreResult<Option<Person>> {
152 self.with(|s| s.bindings.get(&key(tenant, jkt)).cloned())
153 }
154
155 async fn bind(&self, tenant: &str, jkt: &str, person: &Person) -> StoreResult<Person> {
156 self.with(|s| {
157 s.bindings
158 .entry(key(tenant, jkt))
159 .or_insert_with(|| person.clone())
160 .clone()
161 })
162 }
163
164 async fn mission(&self, tenant: &str, s256: &str) -> StoreResult<Option<Mission>> {
165 self.with(|s| s.missions.get(&key(tenant, s256)).cloned())
166 }
167
168 async fn insert_mission(&self, tenant: &str, mission: &Mission) -> StoreResult<bool> {
169 self.with(|s| match s.missions.entry(key(tenant, &mission.s256)) {
170 Entry::Occupied(_) => false,
171 Entry::Vacant(slot) => {
172 slot.insert(mission.clone());
173 true
174 }
175 })
176 }
177
178 async fn log_mission(
179 &self,
180 tenant: &str,
181 s256: &str,
182 entry: LogEntry,
183 now: i64,
184 ) -> StoreResult<bool> {
185 self.with(|s| match s.missions.get_mut(&key(tenant, s256)) {
186 Some(mission) if mission.is_live(now) => {
187 mission.log.push(entry);
188 true
189 }
190 _ => false,
191 })
192 }
193
194 async fn end_mission(&self, tenant: &str, s256: &str, at: i64) -> StoreResult<Option<Mission>> {
195 self.with(|s| {
196 let mission = s.missions.get_mut(&key(tenant, s256))?;
197 mission.completed_at.get_or_insert(at);
198 Some(mission.clone())
199 })
200 }
201
202 async fn spend(
203 &self,
204 tenant: &str,
205 s256: &str,
206 amount: u64,
207 rule: Spend,
208 now: i64,
209 ) -> StoreResult<Option<Mission>> {
210 self.with(|s| {
211 let mission = s.missions.get_mut(&key(tenant, s256))?;
212 if !mission.is_live(now) {
213 return None;
214 }
215 if rule == Spend::WithinBudget {
216 let left = mission.budget.as_ref()?.amount.checked_sub(mission.spent)?;
217 if left < amount {
218 return None;
219 }
220 }
221 mission.spent = mission.spent.saturating_add(amount);
222 Some(mission.clone())
223 })
224 }
225
226 async fn missions(&self, tenant: &str) -> StoreResult<Vec<Mission>> {
227 self.with(|s| {
228 s.missions
229 .iter()
230 .filter(|((t, _), _)| t == tenant)
231 .map(|(_, m)| m.clone())
232 .collect()
233 })
234 }
235
236 async fn approval(&self, tenant: &str, id: &str) -> StoreResult<Option<Approval>> {
237 self.with(|s| s.approvals.get(&key(tenant, id)).cloned())
238 }
239
240 async fn insert_approval(&self, tenant: &str, approval: &Approval) -> StoreResult<()> {
241 self.with(|s| {
242 let now = approval.created_at;
243 s.approvals
244 .retain(|_, a| a.expires_at + APPROVAL_RETENTION_SECS > now);
245 s.approvals
246 .insert(key(tenant, &approval.id), approval.clone());
247 })
248 }
249
250 async fn transition_approval(
251 &self,
252 tenant: &str,
253 expected: ApprovalStatus,
254 next: &Approval,
255 ) -> StoreResult<bool> {
256 self.with(|s| match s.approvals.get_mut(&key(tenant, &next.id)) {
257 Some(current) if current.status == expected => {
258 *current = next.clone();
259 true
260 }
261 _ => false,
262 })
263 }
264
265 async fn approvals(&self, tenant: &str) -> StoreResult<Vec<Approval>> {
266 self.with(|s| {
267 s.approvals
268 .iter()
269 .filter(|((t, _), _)| t == tenant)
270 .map(|(_, a)| a.clone())
271 .collect()
272 })
273 }
274}
275
276#[cfg(test)]
277mod tests {
278 use super::*;
279
280 fn mission(person: &str) -> Mission {
281 Mission {
282 s256: "h".into(),
283 mission: "{}".into(),
284 agent_id: "aauth:travel@agents.example".into(),
285 person: Person::new(person),
286 resources: Vec::new(),
287 budget: None,
288 spent: 0,
289 log: Vec::new(),
290 expires_at: i64::MAX,
291 completed_at: None,
292 tenant: "default".into(),
293 }
294 }
295
296 #[tokio::test]
297 async fn a_mission_is_never_replaced() {
298 let store = MemoryStore::default();
299 assert!(store
300 .insert_mission("default", &mission("sam"))
301 .await
302 .unwrap());
303 store.end_mission("default", "h", 10).await.unwrap();
304 assert!(!store
305 .insert_mission("default", &mission("dana"))
306 .await
307 .unwrap());
308 let kept = store.mission("default", "h").await.unwrap().unwrap();
309 assert_eq!(kept.person.id, "sam");
310 assert_eq!(kept.completed_at, Some(10));
311 }
312
313 #[tokio::test]
314 async fn each_tenant_sees_only_its_own() {
315 let store = MemoryStore::default();
316 assert!(store.insert_mission("acme", &mission("sam")).await.unwrap());
317 assert!(store
318 .insert_mission("globex", &mission("hank"))
319 .await
320 .unwrap());
321 let at = |t| store.mission(t, "h");
322 assert_eq!(at("acme").await.unwrap().unwrap().person.id, "sam");
323 assert_eq!(at("globex").await.unwrap().unwrap().person.id, "hank");
324 assert!(at("default").await.unwrap().is_none());
325 assert_eq!(store.missions("acme").await.unwrap().len(), 1);
326 store.bind("acme", "k", &Person::new("sam")).await.unwrap();
327 assert!(store.binding("globex", "k").await.unwrap().is_none());
328 }
329}