Skip to main content

foxtive_supervisor/persistence/
memory.rs

1use super::{PersistedTaskState, TaskStateStore};
2use crate::error::SupervisorResult;
3use std::collections::HashMap;
4use tokio::sync::Mutex;
5
6/// A default in-memory implementation of TaskStateStore (for reference/testing)
7pub struct InMemoryStateStore {
8    states: Mutex<HashMap<String, PersistedTaskState>>,
9}
10
11impl InMemoryStateStore {
12    pub fn new() -> Self {
13        Self {
14            states: Mutex::new(HashMap::new()),
15        }
16    }
17}
18
19impl Default for InMemoryStateStore {
20    fn default() -> Self {
21        Self::new()
22    }
23}
24
25#[async_trait::async_trait]
26impl TaskStateStore for InMemoryStateStore {
27    async fn save_state(&self, state: PersistedTaskState) -> SupervisorResult<()> {
28        let mut states = self.states.lock().await;
29        states.insert(state.task_id.clone(), state);
30        Ok(())
31    }
32
33    async fn load_state(&self, task_id: &str) -> SupervisorResult<Option<PersistedTaskState>> {
34        let states = self.states.lock().await;
35        Ok(states.get(task_id).cloned())
36    }
37
38    async fn load_all_states(&self) -> SupervisorResult<Vec<PersistedTaskState>> {
39        let states = self.states.lock().await;
40        Ok(states.values().cloned().collect())
41    }
42
43    async fn delete_state(&self, task_id: &str) -> SupervisorResult<()> {
44        let mut states = self.states.lock().await;
45        states.remove(task_id);
46        Ok(())
47    }
48}