mlua_swarm/store/task/
inmemory.rs1use super::{Inner, SharedInner, TaskId, TaskRecord, TaskRecordStatus, TaskStore, TaskStoreError};
5use async_trait::async_trait;
6use std::sync::Mutex;
7
8#[derive(Default)]
12pub struct InMemoryTaskStore {
13 inner: SharedInner,
14}
15
16impl InMemoryTaskStore {
17 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#[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}