use super::{MemoryStateManager, StateManager, StateOptions, StateSnapshot, VersionedState};
use crate::error::{Mecha10Error, Result};
use crate::prelude::now_micros;
use async_trait::async_trait;
use serde::{de::DeserializeOwned, Serialize};
#[async_trait]
impl StateManager for MemoryStateManager {
async fn save_with_options<T: Serialize + Send + Sync>(
&self,
node_id: &str,
state: &T,
options: StateOptions,
) -> Result<()> {
let versioned = if let Some(ttl) = options.ttl {
VersionedState::with_ttl(state, options.version, ttl)
} else {
VersionedState::new(state, options.version)
};
let data = serde_json::to_vec(&versioned).map_err(|e| Mecha10Error::SerializationError {
message: format!("Failed to serialize state: {}", e),
suggestion: "Ensure state type implements Serialize".to_string(),
})?;
let mut states = self.states.write().await;
states.insert(node_id.to_string(), data);
Ok(())
}
async fn load_versioned<T: DeserializeOwned + Send + Sync>(
&self,
node_id: &str,
) -> Result<Option<VersionedState<T>>> {
let states = self.states.read().await;
if let Some(data) = states.get(node_id) {
let versioned: VersionedState<T> =
serde_json::from_slice(data).map_err(|e| Mecha10Error::SerializationError {
message: format!("Failed to deserialize state: {}", e),
suggestion: "State may be corrupted or type mismatch".to_string(),
})?;
if versioned.is_expired() {
return Ok(None);
}
Ok(Some(versioned))
} else {
Ok(None)
}
}
async fn delete(&self, node_id: &str) -> Result<()> {
let mut states = self.states.write().await;
states.remove(node_id);
Ok(())
}
async fn exists(&self, node_id: &str) -> Result<bool> {
let states = self.states.read().await;
Ok(states.contains_key(node_id))
}
async fn list_nodes(&self) -> Result<Vec<String>> {
let states = self.states.read().await;
Ok(states.keys().cloned().collect())
}
async fn size(&self, node_id: &str) -> Result<Option<usize>> {
let states = self.states.read().await;
Ok(states.get(node_id).map(|data| data.len()))
}
async fn create_snapshot(&self, node_id: &str, description: Option<String>) -> Result<StateSnapshot> {
let states = self.states.read().await;
let data = states
.get(node_id)
.ok_or_else(|| Mecha10Error::Configuration(format!("No state for node: {}", node_id)))?
.clone();
let snapshot = StateSnapshot {
id: uuid::Uuid::new_v4().to_string(),
node_id: node_id.to_string(),
timestamp: now_micros(),
description,
size_bytes: data.len(),
};
let mut snapshots = self.snapshots.write().await;
snapshots.insert(snapshot.id.clone(), (data, snapshot.clone()));
Ok(snapshot)
}
async fn restore_snapshot(&self, snapshot_id: &str) -> Result<()> {
let snapshots = self.snapshots.read().await;
let (data, snapshot) = snapshots
.get(snapshot_id)
.ok_or_else(|| Mecha10Error::Configuration(format!("Snapshot not found: {}", snapshot_id)))?;
let mut states = self.states.write().await;
states.insert(snapshot.node_id.clone(), data.clone());
Ok(())
}
async fn list_snapshots(&self, node_id: &str) -> Result<Vec<StateSnapshot>> {
let snapshots = self.snapshots.read().await;
let mut result: Vec<StateSnapshot> = snapshots
.values()
.filter_map(|(_, snapshot)| {
if snapshot.node_id == node_id {
Some(snapshot.clone())
} else {
None
}
})
.collect();
result.sort_by_key(|s| s.timestamp);
result.reverse();
Ok(result)
}
async fn delete_snapshot(&self, snapshot_id: &str) -> Result<()> {
let mut snapshots = self.snapshots.write().await;
snapshots.remove(snapshot_id);
Ok(())
}
async fn compare_and_swap<T: Serialize + DeserializeOwned + Send + Sync>(
&self,
node_id: &str,
expected: &T,
new: &T,
) -> Result<bool> {
let mut states = self.states.write().await;
if let Some(data) = states.get(node_id) {
let current: VersionedState<T> =
serde_json::from_slice(data).map_err(|e| Mecha10Error::SerializationError {
message: format!("Failed to deserialize for CAS: {}", e),
suggestion: "State may be corrupted".to_string(),
})?;
let expected_json = serde_json::to_string(expected).unwrap();
let current_json = serde_json::to_string(¤t.data).unwrap();
if expected_json == current_json {
let versioned = VersionedState::new(new, current.version + 1);
let new_data = serde_json::to_vec(&versioned).unwrap();
states.insert(node_id.to_string(), new_data);
Ok(true)
} else {
Ok(false)
}
} else {
Ok(false)
}
}
async fn clear_expired(&self) -> Result<usize> {
let mut states = self.states.write().await;
let mut expired_keys = Vec::new();
for (node_id, data) in states.iter() {
if let Ok(versioned) = serde_json::from_slice::<VersionedState<serde_json::Value>>(data) {
if versioned.is_expired() {
expired_keys.push(node_id.clone());
}
}
}
let count = expired_keys.len();
for key in expired_keys {
states.remove(&key);
}
Ok(count)
}
}