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

mlua_swarm/store/run/
inmemory.rs
