somatize_core/
state.rs

1//! Trained-state storage — authoritative data produced by `fit()`.
2//!
3//! States are distinct from [`CacheStore`](crate::CacheStore) entries:
4//! - Cache entries are **discardable** — the system can recompute them.
5//! - States are **authoritative** — they are the product of training and
6//!   belong to the Graph that produced them. They must not be evicted
7//!   arbitrarily.
8//!
9//! [`StateStore`] is the trait; implementations may keep states in memory,
10//! on local disk, or in object storage. States are returned as
11//! `Arc<Value>` so the hot forward path can borrow them (`&*arc`) without
12//! cloning potentially-large tensors.
13
14use crate::error::Result;
15use crate::value::Value;
16use std::collections::HashMap;
17use std::sync::{Arc, Mutex};
18
19/// Storage for trained filter states, keyed by node id.
20///
21/// Implementations must be `Send + Sync` and use interior mutability so
22/// the store can be shared (via `Arc`) across the executor and the
23/// graph session.
24pub trait StateStore: Send + Sync {
25    /// Fetch the state for `node_id`, if present.
26    fn get(&self, node_id: &str) -> Result<Option<Arc<Value>>>;
27
28    /// Store `state` under `node_id`, replacing any previous value.
29    fn set(&self, node_id: &str, state: Value) -> Result<()>;
30
31    /// Remove the state for `node_id`, if present.
32    fn remove(&self, node_id: &str) -> Result<()>;
33
34    /// Drop all stored states.
35    fn clear(&self) -> Result<()>;
36
37    /// List all node ids that currently have a stored state.
38    fn keys(&self) -> Result<Vec<String>>;
39}
40
41/// In-memory [`StateStore`] — the default backend.
42///
43/// States live as `Arc<Value>` so reads are zero-copy (just `Arc::clone`)
44/// and multiple consumers can hold references concurrently.
45#[derive(Default)]
46pub struct MemoryStateStore {
47    inner: Mutex<HashMap<String, Arc<Value>>>,
48}
49
50impl MemoryStateStore {
51    /// Create an empty store.
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// Lock the map, tolerating poisoning.
57    ///
58    /// The runtime catches panics from user code and keeps going, so a
59    /// recovered panic must not leave the store permanently unusable —
60    /// which is exactly what propagating the poison would do. The map's
61    /// invariants do not span a lock acquisition, so the data behind a
62    /// poisoned lock is still sound. Same policy as the LRU cache.
63    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<String, Arc<Value>>> {
64        self.inner.lock().unwrap_or_else(|e| e.into_inner())
65    }
66}
67
68impl StateStore for MemoryStateStore {
69    fn get(&self, node_id: &str) -> Result<Option<Arc<Value>>> {
70        Ok(self.lock().get(node_id).cloned())
71    }
72
73    fn set(&self, node_id: &str, state: Value) -> Result<()> {
74        self.lock().insert(node_id.to_string(), Arc::new(state));
75        Ok(())
76    }
77
78    fn remove(&self, node_id: &str) -> Result<()> {
79        self.lock().remove(node_id);
80        Ok(())
81    }
82
83    fn clear(&self) -> Result<()> {
84        self.lock().clear();
85        Ok(())
86    }
87
88    fn keys(&self) -> Result<Vec<String>> {
89        Ok(self.lock().keys().cloned().collect())
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn memory_store_roundtrip() {
99        let store = MemoryStateStore::new();
100        assert!(store.get("a").unwrap().is_none());
101
102        store
103            .set("a", Value::json(serde_json::json!({"mean": 5.0})))
104            .unwrap();
105        let state = store.get("a").unwrap().unwrap();
106        assert_eq!(state.as_json().unwrap()["mean"], 5.0);
107
108        // Same Arc returned on subsequent reads
109        let s1 = store.get("a").unwrap().unwrap();
110        let s2 = store.get("a").unwrap().unwrap();
111        assert!(Arc::ptr_eq(&s1, &s2));
112    }
113
114    #[test]
115    fn memory_store_remove_and_clear() {
116        let store = MemoryStateStore::new();
117        store.set("a", Value::Empty).unwrap();
118        store.set("b", Value::Empty).unwrap();
119        assert_eq!(store.keys().unwrap().len(), 2);
120
121        store.remove("a").unwrap();
122        assert!(store.get("a").unwrap().is_none());
123        assert!(store.get("b").unwrap().is_some());
124
125        store.clear().unwrap();
126        assert!(store.keys().unwrap().is_empty());
127    }
128
129    #[test]
130    fn memory_store_overwrites() {
131        let store = MemoryStateStore::new();
132        store
133            .set("a", Value::json(serde_json::json!({"v": 1})))
134            .unwrap();
135        store
136            .set("a", Value::json(serde_json::json!({"v": 2})))
137            .unwrap();
138        let state = store.get("a").unwrap().unwrap();
139        assert_eq!(state.as_json().unwrap()["v"], 2);
140    }
141}
somatize_core/state.rs

somatize_core/
state.rs
