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, RunListFilter, RunRecord, RunStatus, RunStore, RunStoreError,
6    SharedInner, 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 set_input_json(&self, id: &RunId, input_json: String) -> Result<(), RunStoreError> {
139        let mut inner = self.inner.lock().unwrap();
140        let record = inner
141            .records
142            .get_mut(id)
143            .ok_or_else(|| RunStoreError::NotFound(id.clone()))?;
144        record.input_json = Some(input_json);
145        record.updated_at = crate::types::now_unix();
146        Ok(())
147    }
148
149    async fn list_running(&self) -> Result<Vec<RunRecord>, RunStoreError> {
150        let inner = self.inner.lock().unwrap();
151        let records: Vec<RunRecord> = inner
152            .order
153            .iter()
154            .filter_map(|id| inner.records.get(id).cloned())
155            .filter(|r| r.status == RunStatus::Running)
156            .collect();
157        Ok(records)
158    }
159
160    async fn list(&self, filter: &RunListFilter) -> Result<Vec<RunRecord>, RunStoreError> {
161        let inner = self.inner.lock().unwrap();
162        let mut records: Vec<RunRecord> = inner
163            .order
164            .iter()
165            .filter_map(|id| inner.records.get(id).cloned())
166            .filter(|r| {
167                filter
168                    .task_id
169                    .as_ref()
170                    .map(|t| &r.task_id == t)
171                    .unwrap_or(true)
172                    && filter.status.map(|s| r.status == s).unwrap_or(true)
173            })
174            .collect();
175        // Newest-first; `order` index breaks `created_at` ties stably
176        // (later insertion sorts first within the same second).
177        records.reverse();
178        records.sort_by_key(|r| std::cmp::Reverse(r.created_at));
179        let offset = filter.offset.unwrap_or(0);
180        let records: Vec<RunRecord> = records
181            .into_iter()
182            .skip(offset)
183            .take(filter.limit.unwrap_or(usize::MAX))
184            .collect();
185        Ok(records)
186    }
187
188    async fn delete(&self, id: &RunId) -> Result<(), RunStoreError> {
189        let mut inner = self.inner.lock().unwrap();
190        if inner.records.remove(id).is_none() {
191            return Err(RunStoreError::NotFound(id.clone()));
192        }
193        inner.order.retain(|r| r != id);
194        Ok(())
195    }
196}
197
198// ──────────────────────────────────────────────────────────────────────────
199// tests
200// ──────────────────────────────────────────────────────────────────────────
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205    use serde_json::json;
206
207    fn mk(id: &str, task_id: &str, created_at: u64) -> RunRecord {
208        RunRecord {
209            id: RunId::parse(id).unwrap(),
210            task_id: TaskId::parse(task_id).unwrap(),
211            status: RunStatus::Pending,
212            step_entries: vec![],
213            degradations: vec![],
214            operator_sid: None,
215            result_ref: None,
216            input_json: None,
217            created_at,
218            updated_at: created_at,
219        }
220    }
221
222    fn mk_degradation(tool: &str, at: u64) -> DegradationEntry {
223        DegradationEntry {
224            tool: tool.to_string(),
225            error: "boom".to_string(),
226            fallback: "cached-default".to_string(),
227            note: None,
228            step_ref: Some("worker".to_string()),
229            attempt: Some(1),
230            at,
231        }
232    }
233
234    #[tokio::test]
235    async fn create_then_get() {
236        let s = InMemoryRunStore::new();
237        s.create(mk("R-1", "T-1", 100)).await.unwrap();
238        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
239        assert_eq!(got.task_id, TaskId::parse("T-1").unwrap());
240        assert_eq!(got.status, RunStatus::Pending);
241        assert!(got.step_entries.is_empty());
242    }
243
244    #[tokio::test]
245    async fn duplicate_create_rejected() {
246        let s = InMemoryRunStore::new();
247        s.create(mk("R-1", "T-1", 100)).await.unwrap();
248        let err = s.create(mk("R-1", "T-1", 200)).await.unwrap_err();
249        assert!(matches!(err, RunStoreError::Duplicate(_)));
250    }
251
252    #[tokio::test]
253    async fn get_missing_returns_not_found() {
254        let s = InMemoryRunStore::new();
255        let err = s.get(&RunId::parse("R-nope").unwrap()).await.unwrap_err();
256        assert!(matches!(err, RunStoreError::NotFound(_)));
257    }
258
259    #[tokio::test]
260    async fn list_by_task_filters_and_orders_ascending() {
261        let s = InMemoryRunStore::new();
262        s.create(mk("R-1", "T-1", 300)).await.unwrap();
263        s.create(mk("R-2", "T-2", 50)).await.unwrap();
264        s.create(mk("R-3", "T-1", 100)).await.unwrap();
265        let list = s
266            .list_by_task(&TaskId::parse("T-1").unwrap())
267            .await
268            .unwrap();
269        let ids: Vec<_> = list.iter().map(|r| r.id.to_string()).collect();
270        assert_eq!(ids, vec!["R-3", "R-1"]);
271    }
272
273    #[tokio::test]
274    async fn append_step_entry_accumulates_in_order() {
275        let s = InMemoryRunStore::new();
276        s.create(mk("R-1", "T-1", 100)).await.unwrap();
277        s.append_step_entry(
278            &RunId::parse("R-1").unwrap(),
279            StepEntry::basic(
280                crate::types::StepId::parse("ST-1").unwrap(),
281                Some("step-a".into()),
282                Some("dispatched".into()),
283                None,
284                101,
285            ),
286        )
287        .await
288        .unwrap();
289        s.append_step_entry(
290            &RunId::parse("R-1").unwrap(),
291            StepEntry::basic(
292                crate::types::StepId::parse("ST-2").unwrap(),
293                Some("step-b".into()),
294                Some("passed".into()),
295                None,
296                102,
297            ),
298        )
299        .await
300        .unwrap();
301        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
302        assert_eq!(got.step_entries.len(), 2);
303        assert_eq!(got.step_entries[0].step_ref, Some("step-a".into()));
304        assert_eq!(got.step_entries[1].step_ref, Some("step-b".into()));
305        assert!(got.updated_at >= got.created_at);
306    }
307
308    #[tokio::test]
309    async fn append_degradation_accumulates_in_order() {
310        let s = InMemoryRunStore::new();
311        s.create(mk("R-1", "T-1", 100)).await.unwrap();
312        s.append_degradation(
313            &RunId::parse("R-1").unwrap(),
314            mk_degradation("web_search", 101),
315        )
316        .await
317        .unwrap();
318        s.append_degradation(
319            &RunId::parse("R-1").unwrap(),
320            mk_degradation("code_exec", 102),
321        )
322        .await
323        .unwrap();
324        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
325        assert_eq!(got.degradations.len(), 2);
326        assert_eq!(got.degradations[0].tool, "web_search");
327        assert_eq!(got.degradations[1].tool, "code_exec");
328        assert!(got.updated_at >= got.created_at);
329    }
330
331    #[tokio::test]
332    async fn append_degradation_unknown_run_fails() {
333        let s = InMemoryRunStore::new();
334        let err = s
335            .append_degradation(
336                &RunId::parse("R-nope").unwrap(),
337                mk_degradation("web_search", 1),
338            )
339            .await
340            .unwrap_err();
341        assert!(matches!(err, RunStoreError::NotFound(_)));
342    }
343
344    #[tokio::test]
345    async fn append_degradation_bumps_updated_at() {
346        let s = InMemoryRunStore::new();
347        s.create(mk("R-1", "T-1", 100)).await.unwrap();
348        s.append_degradation(
349            &RunId::parse("R-1").unwrap(),
350            mk_degradation("web_search", 200),
351        )
352        .await
353        .unwrap();
354        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
355        assert!(got.updated_at > 100);
356    }
357
358    #[tokio::test]
359    async fn append_step_entry_unknown_run_fails() {
360        let s = InMemoryRunStore::new();
361        let err = s
362            .append_step_entry(
363                &RunId::parse("R-nope").unwrap(),
364                StepEntry::basic(
365                    crate::types::StepId::parse("ST-1").unwrap(),
366                    None,
367                    None,
368                    None,
369                    1,
370                ),
371            )
372            .await
373            .unwrap_err();
374        assert!(matches!(err, RunStoreError::NotFound(_)));
375    }
376
377    #[tokio::test]
378    async fn update_status_persists() {
379        let s = InMemoryRunStore::new();
380        s.create(mk("R-1", "T-1", 100)).await.unwrap();
381        s.update_status(&RunId::parse("R-1").unwrap(), RunStatus::Running)
382            .await
383            .unwrap();
384        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
385        assert_eq!(got.status, RunStatus::Running);
386    }
387
388    #[tokio::test]
389    async fn set_result_persists() {
390        let s = InMemoryRunStore::new();
391        s.create(mk("R-1", "T-1", 100)).await.unwrap();
392        s.set_result(&RunId::parse("R-1").unwrap(), json!({"ok": true}))
393            .await
394            .unwrap();
395        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
396        assert_eq!(got.result_ref, Some(json!({"ok": true})));
397    }
398
399    #[tokio::test]
400    async fn name_is_in_memory() {
401        assert_eq!(InMemoryRunStore::new().name(), "in-memory");
402    }
403
404    #[tokio::test]
405    async fn list_running_filters_by_status() {
406        let s = InMemoryRunStore::new();
407        s.create(mk("R-1", "T-1", 100)).await.unwrap();
408        s.create(mk("R-2", "T-2", 200)).await.unwrap();
409        s.create(mk("R-3", "T-3", 300)).await.unwrap();
410        s.update_status(&RunId::parse("R-2").unwrap(), RunStatus::Running)
411            .await
412            .unwrap();
413        s.update_status(&RunId::parse("R-3").unwrap(), RunStatus::Done)
414            .await
415            .unwrap();
416        let running = s.list_running().await.unwrap();
417        assert_eq!(running.len(), 1);
418        assert_eq!(running[0].id, RunId::parse("R-2").unwrap());
419        assert_eq!(running[0].status, RunStatus::Running);
420    }
421
422    #[tokio::test]
423    async fn try_transition_flips_on_match_and_is_idempotent_under_race() {
424        let s = InMemoryRunStore::new();
425        s.create(mk("R-1", "T-1", 100)).await.unwrap();
426        s.update_status(&RunId::parse("R-1").unwrap(), RunStatus::Interrupted)
427            .await
428            .unwrap();
429
430        // First CAS matches `Interrupted` and flips to `Running`.
431        let first = s
432            .try_transition(
433                &RunId::parse("R-1").unwrap(),
434                RunStatus::Interrupted,
435                RunStatus::Running,
436            )
437            .await
438            .unwrap();
439        assert!(first, "first CAS must flip Interrupted -> Running");
440        assert_eq!(
441            s.get(&RunId::parse("R-1").unwrap()).await.unwrap().status,
442            RunStatus::Running
443        );
444
445        // Second CAS (a racing double-resume) no longer sees `Interrupted`
446        // and must report `false` without touching the row.
447        let second = s
448            .try_transition(
449                &RunId::parse("R-1").unwrap(),
450                RunStatus::Interrupted,
451                RunStatus::Running,
452            )
453            .await
454            .unwrap();
455        assert!(!second, "second CAS must not flip a now-Running row");
456    }
457
458    #[tokio::test]
459    async fn try_transition_absent_run_reports_false() {
460        let s = InMemoryRunStore::new();
461        let flipped = s
462            .try_transition(
463                &RunId::parse("R-nope").unwrap(),
464                RunStatus::Interrupted,
465                RunStatus::Running,
466            )
467            .await
468            .unwrap();
469        assert!(!flipped, "an absent Run must report false, not error");
470    }
471
472    #[tokio::test]
473    async fn input_json_roundtrips_through_create_get() {
474        let s = InMemoryRunStore::new();
475        let mut rec = mk("R-1", "T-1", 100);
476        rec.input_json = Some(r#"{"blueprint":"snapshot"}"#.to_string());
477        s.create(rec).await.unwrap();
478        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
479        assert_eq!(
480            got.input_json.as_deref(),
481            Some(r#"{"blueprint":"snapshot"}"#)
482        );
483    }
484}
mlua_swarm/store/run/inmemory.rs

mlua_swarm/store/run/
inmemory.rs
