mecha10_core/state/
memory_impl.rs

1//! `StateManager` implementation for [`super::MemoryStateManager`].
2
3use super::{MemoryStateManager, StateManager, StateOptions, StateSnapshot, VersionedState};
4use crate::error::{Mecha10Error, Result};
5use crate::prelude::now_micros;
6use async_trait::async_trait;
7use serde::{de::DeserializeOwned, Serialize};
8
9#[async_trait]
10impl StateManager for MemoryStateManager {
11    async fn save_with_options<T: Serialize + Send + Sync>(
12        &self,
13        node_id: &str,
14        state: &T,
15        options: StateOptions,
16    ) -> Result<()> {
17        let versioned = if let Some(ttl) = options.ttl {
18            VersionedState::with_ttl(state, options.version, ttl)
19        } else {
20            VersionedState::new(state, options.version)
21        };
22
23        let data = serde_json::to_vec(&versioned).map_err(|e| Mecha10Error::SerializationError {
24            message: format!("Failed to serialize state: {}", e),
25            suggestion: "Ensure state type implements Serialize".to_string(),
26        })?;
27
28        let mut states = self.states.write().await;
29        states.insert(node_id.to_string(), data);
30
31        Ok(())
32    }
33
34    async fn load_versioned<T: DeserializeOwned + Send + Sync>(
35        &self,
36        node_id: &str,
37    ) -> Result<Option<VersionedState<T>>> {
38        let states = self.states.read().await;
39
40        if let Some(data) = states.get(node_id) {
41            let versioned: VersionedState<T> =
42                serde_json::from_slice(data).map_err(|e| Mecha10Error::SerializationError {
43                    message: format!("Failed to deserialize state: {}", e),
44                    suggestion: "State may be corrupted or type mismatch".to_string(),
45                })?;
46
47            if versioned.is_expired() {
48                return Ok(None);
49            }
50
51            Ok(Some(versioned))
52        } else {
53            Ok(None)
54        }
55    }
56
57    async fn delete(&self, node_id: &str) -> Result<()> {
58        let mut states = self.states.write().await;
59        states.remove(node_id);
60        Ok(())
61    }
62
63    async fn exists(&self, node_id: &str) -> Result<bool> {
64        let states = self.states.read().await;
65        Ok(states.contains_key(node_id))
66    }
67
68    async fn list_nodes(&self) -> Result<Vec<String>> {
69        let states = self.states.read().await;
70        Ok(states.keys().cloned().collect())
71    }
72
73    async fn size(&self, node_id: &str) -> Result<Option<usize>> {
74        let states = self.states.read().await;
75        Ok(states.get(node_id).map(|data| data.len()))
76    }
77
78    async fn create_snapshot(&self, node_id: &str, description: Option<String>) -> Result<StateSnapshot> {
79        let states = self.states.read().await;
80
81        let data = states
82            .get(node_id)
83            .ok_or_else(|| Mecha10Error::Configuration(format!("No state for node: {}", node_id)))?
84            .clone();
85
86        let snapshot = StateSnapshot {
87            id: uuid::Uuid::new_v4().to_string(),
88            node_id: node_id.to_string(),
89            timestamp: now_micros(),
90            description,
91            size_bytes: data.len(),
92        };
93
94        let mut snapshots = self.snapshots.write().await;
95        snapshots.insert(snapshot.id.clone(), (data, snapshot.clone()));
96
97        Ok(snapshot)
98    }
99
100    async fn restore_snapshot(&self, snapshot_id: &str) -> Result<()> {
101        let snapshots = self.snapshots.read().await;
102
103        let (data, snapshot) = snapshots
104            .get(snapshot_id)
105            .ok_or_else(|| Mecha10Error::Configuration(format!("Snapshot not found: {}", snapshot_id)))?;
106
107        let mut states = self.states.write().await;
108        states.insert(snapshot.node_id.clone(), data.clone());
109
110        Ok(())
111    }
112
113    async fn list_snapshots(&self, node_id: &str) -> Result<Vec<StateSnapshot>> {
114        let snapshots = self.snapshots.read().await;
115
116        let mut result: Vec<StateSnapshot> = snapshots
117            .values()
118            .filter_map(|(_, snapshot)| {
119                if snapshot.node_id == node_id {
120                    Some(snapshot.clone())
121                } else {
122                    None
123                }
124            })
125            .collect();
126
127        result.sort_by_key(|s| s.timestamp);
128        result.reverse();
129
130        Ok(result)
131    }
132
133    async fn delete_snapshot(&self, snapshot_id: &str) -> Result<()> {
134        let mut snapshots = self.snapshots.write().await;
135        snapshots.remove(snapshot_id);
136        Ok(())
137    }
138
139    async fn compare_and_swap<T: Serialize + DeserializeOwned + Send + Sync>(
140        &self,
141        node_id: &str,
142        expected: &T,
143        new: &T,
144    ) -> Result<bool> {
145        let mut states = self.states.write().await;
146
147        if let Some(data) = states.get(node_id) {
148            let current: VersionedState<T> =
149                serde_json::from_slice(data).map_err(|e| Mecha10Error::SerializationError {
150                    message: format!("Failed to deserialize for CAS: {}", e),
151                    suggestion: "State may be corrupted".to_string(),
152                })?;
153
154            let expected_json = serde_json::to_string(expected).unwrap();
155            let current_json = serde_json::to_string(&current.data).unwrap();
156
157            if expected_json == current_json {
158                let versioned = VersionedState::new(new, current.version + 1);
159                let new_data = serde_json::to_vec(&versioned).unwrap();
160                states.insert(node_id.to_string(), new_data);
161                Ok(true)
162            } else {
163                Ok(false)
164            }
165        } else {
166            Ok(false)
167        }
168    }
169
170    async fn clear_expired(&self) -> Result<usize> {
171        let mut states = self.states.write().await;
172        let mut expired_keys = Vec::new();
173
174        for (node_id, data) in states.iter() {
175            if let Ok(versioned) = serde_json::from_slice::<VersionedState<serde_json::Value>>(data) {
176                if versioned.is_expired() {
177                    expired_keys.push(node_id.clone());
178                }
179            }
180        }
181
182        let count = expired_keys.len();
183        for key in expired_keys {
184            states.remove(&key);
185        }
186
187        Ok(count)
188    }
189}
mecha10_core/state/memory_impl.rs

mecha10_core/state/
memory_impl.rs
