mlua_swarm/store/run/
inmemory.rs

1//! `InMemoryRunStore` — a process-volatile `RunStore` used by the current
2//! default.
3
4use super::{
5    DegradationEntry, Inner, RunId, RunRecord, RunStatus, RunStore, RunStoreError, SharedInner,
6    StepEntry, TaskId,
7};
8use async_trait::async_trait;
9use std::sync::Mutex;
10
11/// Process-volatile [`RunStore`] used as the current default. Entries are
12/// lost on restart; persistent backends (SQLite / Git / mini-app / …) are
13/// future carries.
14#[derive(Default)]
15pub struct InMemoryRunStore {
16    inner: SharedInner,
17}
18
19impl InMemoryRunStore {
20    /// Create an empty store.
21    pub fn new() -> Self {
22        Self {
23            inner: Mutex::new(Inner::default()),
24        }
25    }
26}
27
28#[async_trait]
29impl RunStore for InMemoryRunStore {
30    fn name(&self) -> &str {
31        "in-memory"
32    }
33
34    async fn create(&self, record: RunRecord) -> Result<(), RunStoreError> {
35        let mut inner = self.inner.lock().unwrap();
36        if inner.records.contains_key(&record.id) {
37            return Err(RunStoreError::Duplicate(record.id));
38        }
39        inner.order.push(record.id.clone());
40        inner.records.insert(record.id.clone(), record);
41        Ok(())
42    }
43
44    async fn get(&self, id: &RunId) -> Result<RunRecord, RunStoreError> {
45        let inner = self.inner.lock().unwrap();
46        inner
47            .records
48            .get(id)
49            .cloned()
50            .ok_or_else(|| RunStoreError::NotFound(id.clone()))
51    }
52
53    async fn list_by_task(&self, task_id: &TaskId) -> Result<Vec<RunRecord>, RunStoreError> {
54        let inner = self.inner.lock().unwrap();
55        let mut records: Vec<RunRecord> = inner
56            .order
57            .iter()
58            .filter_map(|id| inner.records.get(id).cloned())
59            .filter(|r| &r.task_id == task_id)
60            .collect();
61        records.sort_by_key(|r| r.created_at);
62        Ok(records)
63    }
64
65    async fn append_step_entry(&self, id: &RunId, entry: StepEntry) -> Result<(), RunStoreError> {
66        let mut inner = self.inner.lock().unwrap();
67        let record = inner
68            .records
69            .get_mut(id)
70            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
71        record.step_entries.push(entry);
72        record.updated_at = crate::types::now_unix();
73        Ok(())
74    }
75
76    async fn append_degradation(
77        &self,
78        id: &RunId,
79        entry: DegradationEntry,
80    ) -> Result<(), RunStoreError> {
81        let mut inner = self.inner.lock().unwrap();
82        let record = inner
83            .records
84            .get_mut(id)
85            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
86        record.degradations.push(entry);
87        record.updated_at = crate::types::now_unix();
88        Ok(())
89    }
90
91    async fn update_status(&self, id: &RunId, status: RunStatus) -> Result<(), RunStoreError> {
92        let mut inner = self.inner.lock().unwrap();
93        let record = inner
94            .records
95            .get_mut(id)
96            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
97        record.status = status;
98        record.updated_at = crate::types::now_unix();
99        Ok(())
100    }
101
102    async fn try_transition(
103        &self,
104        id: &RunId,
105        from: RunStatus,
106        to: RunStatus,
107    ) -> Result<bool, RunStoreError> {
108        let mut inner = self.inner.lock().unwrap();
109        // Held under the single `inner` mutex, so the read + compare + set
110        // is atomic against any other appender/transition. An absent row or
111        // a status mismatch both report `false` (the caller's race signal),
112        // not an error.
113        match inner.records.get_mut(id) {
114            Some(record) if record.status == from => {
115                record.status = to;
116                record.updated_at = crate::types::now_unix();
117                Ok(true)
118            }
119            _ => Ok(false),
120        }
121    }
122
123    async fn set_result(
124        &self,
125        id: &RunId,
126        result_ref: serde_json::Value,
127    ) -> Result<(), RunStoreError> {
128        let mut inner = self.inner.lock().unwrap();
129        let record = inner
130            .records
131            .get_mut(id)
132            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
133        record.result_ref = Some(result_ref);
134        record.updated_at = crate::types::now_unix();
135        Ok(())
136    }
137
138    async fn list_running(&self) -> Result<Vec<RunRecord>, RunStoreError> {
139        let inner = self.inner.lock().unwrap();
140        let records: Vec<RunRecord> = inner
141            .order
142            .iter()
143            .filter_map(|id| inner.records.get(id).cloned())
144            .filter(|r| r.status == RunStatus::Running)
145            .collect();
146        Ok(records)
147    }
148}
149
150// ──────────────────────────────────────────────────────────────────────────
151// tests
152// ──────────────────────────────────────────────────────────────────────────
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use serde_json::json;
158
159    fn mk(id: &str, task_id: &str, created_at: u64) -> RunRecord {
160        RunRecord {
161            id: RunId::parse(id).unwrap(),
162            task_id: TaskId::parse(task_id).unwrap(),
163            status: RunStatus::Pending,
164            step_entries: vec![],
165            degradations: vec![],
166            operator_sid: None,
167            result_ref: None,
168            input_json: None,
169            created_at,
170            updated_at: created_at,
171        }
172    }
173
174    fn mk_degradation(tool: &str, at: u64) -> DegradationEntry {
175        DegradationEntry {
176            tool: tool.to_string(),
177            error: "boom".to_string(),
178            fallback: "cached-default".to_string(),
179            note: None,
180            step_ref: Some("worker".to_string()),
181            attempt: Some(1),
182            at,
183        }
184    }
185
186    #[tokio::test]
187    async fn create_then_get() {
188        let s = InMemoryRunStore::new();
189        s.create(mk("R-1", "T-1", 100)).await.unwrap();
190        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
191        assert_eq!(got.task_id, TaskId::parse("T-1").unwrap());
192        assert_eq!(got.status, RunStatus::Pending);
193        assert!(got.step_entries.is_empty());
194    }
195
196    #[tokio::test]
197    async fn duplicate_create_rejected() {
198        let s = InMemoryRunStore::new();
199        s.create(mk("R-1", "T-1", 100)).await.unwrap();
200        let err = s.create(mk("R-1", "T-1", 200)).await.unwrap_err();
201        assert!(matches!(err, RunStoreError::Duplicate(_)));
202    }
203
204    #[tokio::test]
205    async fn get_missing_returns_not_found() {
206        let s = InMemoryRunStore::new();
207        let err = s.get(&RunId::parse("R-nope").unwrap()).await.unwrap_err();
208        assert!(matches!(err, RunStoreError::NotFound(_)));
209    }
210
211    #[tokio::test]
212    async fn list_by_task_filters_and_orders_ascending() {
213        let s = InMemoryRunStore::new();
214        s.create(mk("R-1", "T-1", 300)).await.unwrap();
215        s.create(mk("R-2", "T-2", 50)).await.unwrap();
216        s.create(mk("R-3", "T-1", 100)).await.unwrap();
217        let list = s
218            .list_by_task(&TaskId::parse("T-1").unwrap())
219            .await
220            .unwrap();
221        let ids: Vec<_> = list.iter().map(|r| r.id.to_string()).collect();
222        assert_eq!(ids, vec!["R-3", "R-1"]);
223    }
224
225    #[tokio::test]
226    async fn append_step_entry_accumulates_in_order() {
227        let s = InMemoryRunStore::new();
228        s.create(mk("R-1", "T-1", 100)).await.unwrap();
229        s.append_step_entry(
230            &RunId::parse("R-1").unwrap(),
231            StepEntry {
232                step_id: crate::types::StepId::parse("ST-1").unwrap(),
233                step_ref: Some("step-a".into()),
234                status: Some("dispatched".into()),
235                at: 101,
236            },
237        )
238        .await
239        .unwrap();
240        s.append_step_entry(
241            &RunId::parse("R-1").unwrap(),
242            StepEntry {
243                step_id: crate::types::StepId::parse("ST-2").unwrap(),
244                step_ref: Some("step-b".into()),
245                status: Some("passed".into()),
246                at: 102,
247            },
248        )
249        .await
250        .unwrap();
251        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
252        assert_eq!(got.step_entries.len(), 2);
253        assert_eq!(got.step_entries[0].step_ref, Some("step-a".into()));
254        assert_eq!(got.step_entries[1].step_ref, Some("step-b".into()));
255        assert!(got.updated_at >= got.created_at);
256    }
257
258    #[tokio::test]
259    async fn append_degradation_accumulates_in_order() {
260        let s = InMemoryRunStore::new();
261        s.create(mk("R-1", "T-1", 100)).await.unwrap();
262        s.append_degradation(
263            &RunId::parse("R-1").unwrap(),
264            mk_degradation("web_search", 101),
265        )
266        .await
267        .unwrap();
268        s.append_degradation(
269            &RunId::parse("R-1").unwrap(),
270            mk_degradation("code_exec", 102),
271        )
272        .await
273        .unwrap();
274        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
275        assert_eq!(got.degradations.len(), 2);
276        assert_eq!(got.degradations[0].tool, "web_search");
277        assert_eq!(got.degradations[1].tool, "code_exec");
278        assert!(got.updated_at >= got.created_at);
279    }
280
281    #[tokio::test]
282    async fn append_degradation_unknown_run_fails() {
283        let s = InMemoryRunStore::new();
284        let err = s
285            .append_degradation(
286                &RunId::parse("R-nope").unwrap(),
287                mk_degradation("web_search", 1),
288            )
289            .await
290            .unwrap_err();
291        assert!(matches!(err, RunStoreError::NotFound(_)));
292    }
293
294    #[tokio::test]
295    async fn append_degradation_bumps_updated_at() {
296        let s = InMemoryRunStore::new();
297        s.create(mk("R-1", "T-1", 100)).await.unwrap();
298        s.append_degradation(
299            &RunId::parse("R-1").unwrap(),
300            mk_degradation("web_search", 200),
301        )
302        .await
303        .unwrap();
304        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
305        assert!(got.updated_at > 100);
306    }
307
308    #[tokio::test]
309    async fn append_step_entry_unknown_run_fails() {
310        let s = InMemoryRunStore::new();
311        let err = s
312            .append_step_entry(
313                &RunId::parse("R-nope").unwrap(),
314                StepEntry {
315                    step_id: crate::types::StepId::parse("ST-1").unwrap(),
316                    step_ref: None,
317                    status: None,
318                    at: 1,
319                },
320            )
321            .await
322            .unwrap_err();
323        assert!(matches!(err, RunStoreError::NotFound(_)));
324    }
325
326    #[tokio::test]
327    async fn update_status_persists() {
328        let s = InMemoryRunStore::new();
329        s.create(mk("R-1", "T-1", 100)).await.unwrap();
330        s.update_status(&RunId::parse("R-1").unwrap(), RunStatus::Running)
331            .await
332            .unwrap();
333        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
334        assert_eq!(got.status, RunStatus::Running);
335    }
336
337    #[tokio::test]
338    async fn set_result_persists() {
339        let s = InMemoryRunStore::new();
340        s.create(mk("R-1", "T-1", 100)).await.unwrap();
341        s.set_result(&RunId::parse("R-1").unwrap(), json!({"ok": true}))
342            .await
343            .unwrap();
344        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
345        assert_eq!(got.result_ref, Some(json!({"ok": true})));
346    }
347
348    #[tokio::test]
349    async fn name_is_in_memory() {
350        assert_eq!(InMemoryRunStore::new().name(), "in-memory");
351    }
352
353    #[tokio::test]
354    async fn list_running_filters_by_status() {
355        let s = InMemoryRunStore::new();
356        s.create(mk("R-1", "T-1", 100)).await.unwrap();
357        s.create(mk("R-2", "T-2", 200)).await.unwrap();
358        s.create(mk("R-3", "T-3", 300)).await.unwrap();
359        s.update_status(&RunId::parse("R-2").unwrap(), RunStatus::Running)
360            .await
361            .unwrap();
362        s.update_status(&RunId::parse("R-3").unwrap(), RunStatus::Done)
363            .await
364            .unwrap();
365        let running = s.list_running().await.unwrap();
366        assert_eq!(running.len(), 1);
367        assert_eq!(running[0].id, RunId::parse("R-2").unwrap());
368        assert_eq!(running[0].status, RunStatus::Running);
369    }
370
371    #[tokio::test]
372    async fn try_transition_flips_on_match_and_is_idempotent_under_race() {
373        let s = InMemoryRunStore::new();
374        s.create(mk("R-1", "T-1", 100)).await.unwrap();
375        s.update_status(&RunId::parse("R-1").unwrap(), RunStatus::Interrupted)
376            .await
377            .unwrap();
378
379        // First CAS matches `Interrupted` and flips to `Running`.
380        let first = s
381            .try_transition(
382                &RunId::parse("R-1").unwrap(),
383                RunStatus::Interrupted,
384                RunStatus::Running,
385            )
386            .await
387            .unwrap();
388        assert!(first, "first CAS must flip Interrupted -> Running");
389        assert_eq!(
390            s.get(&RunId::parse("R-1").unwrap()).await.unwrap().status,
391            RunStatus::Running
392        );
393
394        // Second CAS (a racing double-resume) no longer sees `Interrupted`
395        // and must report `false` without touching the row.
396        let second = s
397            .try_transition(
398                &RunId::parse("R-1").unwrap(),
399                RunStatus::Interrupted,
400                RunStatus::Running,
401            )
402            .await
403            .unwrap();
404        assert!(!second, "second CAS must not flip a now-Running row");
405    }
406
407    #[tokio::test]
408    async fn try_transition_absent_run_reports_false() {
409        let s = InMemoryRunStore::new();
410        let flipped = s
411            .try_transition(
412                &RunId::parse("R-nope").unwrap(),
413                RunStatus::Interrupted,
414                RunStatus::Running,
415            )
416            .await
417            .unwrap();
418        assert!(!flipped, "an absent Run must report false, not error");
419    }
420
421    #[tokio::test]
422    async fn input_json_roundtrips_through_create_get() {
423        let s = InMemoryRunStore::new();
424        let mut rec = mk("R-1", "T-1", 100);
425        rec.input_json = Some(r#"{"blueprint":"snapshot"}"#.to_string());
426        s.create(rec).await.unwrap();
427        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
428        assert_eq!(
429            got.input_json.as_deref(),
430            Some(r#"{"blueprint":"snapshot"}"#)
431        );
432    }
433}
mlua_swarm/store/run/inmemory.rs

mlua_swarm/store/run/
inmemory.rs
