Skip to main content

mlua_swarm/store/task/
inmemory.rs

1//! `InMemoryTaskStore` — a process-volatile `TaskStore` used by the current
2//! default.
3
4use super::{Inner, SharedInner, TaskId, TaskRecord, TaskRecordStatus, TaskStore, TaskStoreError};
5use async_trait::async_trait;
6use std::sync::Mutex;
7
8/// Process-volatile [`TaskStore`] used as the current default. Entries are
9/// lost on restart; persistent backends (SQLite / Git / mini-app / …) are
10/// future carries.
11#[derive(Default)]
12pub struct InMemoryTaskStore {
13    inner: SharedInner,
14}
15
16impl InMemoryTaskStore {
17    /// Create an empty store.
18    pub fn new() -> Self {
19        Self {
20            inner: Mutex::new(Inner::default()),
21        }
22    }
23}
24
25#[async_trait]
26impl TaskStore for InMemoryTaskStore {
27    fn name(&self) -> &str {
28        "in-memory"
29    }
30
31    async fn create(&self, record: TaskRecord) -> Result<(), TaskStoreError> {
32        let mut inner = self.inner.lock().unwrap();
33        if inner.records.contains_key(&record.id) {
34            return Err(TaskStoreError::Duplicate(record.id));
35        }
36        inner.order.push(record.id.clone());
37        inner.records.insert(record.id.clone(), record);
38        Ok(())
39    }
40
41    async fn get(&self, id: &TaskId) -> Result<TaskRecord, TaskStoreError> {
42        let inner = self.inner.lock().unwrap();
43        inner
44            .records
45            .get(id)
46            .cloned()
47            .ok_or_else(|| TaskStoreError::NotFound(id.clone()))
48    }
49
50    async fn list(&self) -> Result<Vec<TaskRecord>, TaskStoreError> {
51        let inner = self.inner.lock().unwrap();
52        let mut records: Vec<TaskRecord> = inner
53            .order
54            .iter()
55            .filter_map(|id| inner.records.get(id).cloned())
56            .collect();
57        records.sort_by_key(|r| std::cmp::Reverse(r.created_at));
58        Ok(records)
59    }
60
61    async fn update_status(
62        &self,
63        id: &TaskId,
64        status: TaskRecordStatus,
65    ) -> Result<(), TaskStoreError> {
66        let mut inner = self.inner.lock().unwrap();
67        let record = inner
68            .records
69            .get_mut(id)
70            .ok_or_else(|| TaskStoreError::NotFound(id.clone()))?;
71        record.status = status;
72        record.updated_at = crate::types::now_unix();
73        Ok(())
74    }
75}
76
77// ──────────────────────────────────────────────────────────────────────────
78// tests
79// ──────────────────────────────────────────────────────────────────────────
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use serde_json::json;
85
86    fn mk(id: &str, created_at: u64) -> TaskRecord {
87        TaskRecord {
88            id: TaskId::parse(id).unwrap(),
89            goal: format!("goal for {id}"),
90            blueprint_ref: json!({"id": "bp-1"}),
91            input_ctx: json!({"k": "v"}),
92            status: TaskRecordStatus::Pending,
93            created_at,
94            updated_at: created_at,
95        }
96    }
97
98    #[tokio::test]
99    async fn create_then_get() {
100        let s = InMemoryTaskStore::new();
101        s.create(mk("T-1", 100)).await.unwrap();
102        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
103        assert_eq!(got.id, TaskId::parse("T-1").unwrap());
104        assert_eq!(got.goal, "goal for T-1");
105        assert_eq!(got.status, TaskRecordStatus::Pending);
106    }
107
108    #[tokio::test]
109    async fn duplicate_create_rejected() {
110        let s = InMemoryTaskStore::new();
111        s.create(mk("T-1", 100)).await.unwrap();
112        let err = s.create(mk("T-1", 200)).await.unwrap_err();
113        assert!(matches!(err, TaskStoreError::Duplicate(_)));
114    }
115
116    #[tokio::test]
117    async fn get_missing_returns_not_found() {
118        let s = InMemoryTaskStore::new();
119        let err = s.get(&TaskId::parse("T-nope").unwrap()).await.unwrap_err();
120        assert!(matches!(err, TaskStoreError::NotFound(_)));
121    }
122
123    #[tokio::test]
124    async fn list_returns_newest_first() {
125        let s = InMemoryTaskStore::new();
126        s.create(mk("T-1", 100)).await.unwrap();
127        s.create(mk("T-2", 300)).await.unwrap();
128        s.create(mk("T-3", 200)).await.unwrap();
129        let list = s.list().await.unwrap();
130        let ids: Vec<_> = list.iter().map(|r| r.id.to_string()).collect();
131        assert_eq!(ids, vec!["T-2", "T-3", "T-1"]);
132    }
133
134    #[tokio::test]
135    async fn update_status_bumps_updated_at_and_persists_status() {
136        let s = InMemoryTaskStore::new();
137        s.create(mk("T-1", 100)).await.unwrap();
138        s.update_status(&TaskId::parse("T-1").unwrap(), TaskRecordStatus::Running)
139            .await
140            .unwrap();
141        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
142        assert_eq!(got.status, TaskRecordStatus::Running);
143        assert!(got.updated_at >= got.created_at);
144    }
145
146    #[tokio::test]
147    async fn update_status_unknown_fails() {
148        let s = InMemoryTaskStore::new();
149        let err = s
150            .update_status(&TaskId::parse("T-nope").unwrap(), TaskRecordStatus::Done)
151            .await
152            .unwrap_err();
153        assert!(matches!(err, TaskStoreError::NotFound(_)));
154    }
155
156    #[tokio::test]
157    async fn name_is_in_memory() {
158        assert_eq!(InMemoryTaskStore::new().name(), "in-memory");
159    }
160}