oxigdal_streaming/state/
operator_state.rs

1//! Operator state for stream processing.
2
3use crate::error::{Result, StreamingError};
4use std::collections::HashMap;
5use std::future::Future;
6use std::pin::Pin;
7use std::sync::Arc;
8use tokio::sync::RwLock;
9
10/// Operator state trait.
11pub trait OperatorState: Send + Sync {
12    /// Snapshot the state.
13    fn snapshot(&self) -> impl std::future::Future<Output = Result<Vec<u8>>> + Send;
14
15    /// Restore from a snapshot.
16    fn restore(&self, snapshot: &[u8]) -> impl std::future::Future<Output = Result<()>> + Send;
17}
18
19/// Object-safe (`dyn`-compatible) adapter over [`OperatorState`].
20///
21/// [`OperatorState`] uses return-position `impl Future` (RPITIT), which is not
22/// `dyn`-safe, so it cannot be stored behind `Arc<dyn OperatorState>`. This
23/// adapter boxes the returned futures so heterogeneous operator states can be
24/// registered with the [`CheckpointCoordinator`](crate::state::CheckpointCoordinator)
25/// for real state capture and recovery. A blanket implementation is provided
26/// for every `T: OperatorState`, so any concrete operator state works
27/// automatically via `Arc::new(state) as Arc<dyn DynOperatorState>`.
28pub trait DynOperatorState: Send + Sync {
29    /// Snapshot the state, returning a boxed future.
30    fn snapshot_boxed(&self) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + '_>>;
31
32    /// Restore from a snapshot, returning a boxed future.
33    fn restore_boxed<'a>(
34        &'a self,
35        snapshot: &'a [u8],
36    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
37}
38
39impl<T: OperatorState> DynOperatorState for T {
40    fn snapshot_boxed(&self) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + '_>> {
41        Box::pin(self.snapshot())
42    }
43
44    fn restore_boxed<'a>(
45        &'a self,
46        snapshot: &'a [u8],
47    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
48        Box::pin(self.restore(snapshot))
49    }
50}
51
52/// Broadcast state (shared across all parallel instances).
53pub struct BroadcastState {
54    state: Arc<RwLock<HashMap<Vec<u8>, Vec<u8>>>>,
55}
56
57impl BroadcastState {
58    /// Create a new broadcast state.
59    pub fn new() -> Self {
60        Self {
61            state: Arc::new(RwLock::new(HashMap::new())),
62        }
63    }
64
65    /// Get a value.
66    pub async fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
67        self.state.read().await.get(key).cloned()
68    }
69
70    /// Put a value.
71    pub async fn put(&self, key: Vec<u8>, value: Vec<u8>) {
72        self.state.write().await.insert(key, value);
73    }
74
75    /// Remove a value.
76    pub async fn remove(&self, key: &[u8]) {
77        self.state.write().await.remove(key);
78    }
79
80    /// Check if a key exists.
81    pub async fn contains(&self, key: &[u8]) -> bool {
82        self.state.read().await.contains_key(key)
83    }
84
85    /// Clear all state.
86    pub async fn clear(&self) {
87        self.state.write().await.clear();
88    }
89
90    /// Get all keys.
91    pub async fn keys(&self) -> Vec<Vec<u8>> {
92        self.state.read().await.keys().cloned().collect()
93    }
94}
95
96impl Default for BroadcastState {
97    fn default() -> Self {
98        Self::new()
99    }
100}
101
102impl OperatorState for BroadcastState {
103    async fn snapshot(&self) -> Result<Vec<u8>> {
104        let state = self.state.read().await;
105        // Use oxicode for binary serialization since JSON requires string keys
106        oxicode::encode_to_vec(&*state)
107            .map_err(|e| StreamingError::SerializationError(e.to_string()))
108    }
109
110    async fn restore(&self, snapshot: &[u8]) -> Result<()> {
111        let (restored, _): (HashMap<Vec<u8>, Vec<u8>>, _) = oxicode::decode_from_slice(snapshot)
112            .map_err(|e| StreamingError::SerializationError(e.to_string()))?;
113        *self.state.write().await = restored;
114        Ok(())
115    }
116}
117
118/// Union list state (list that is distributed across parallel instances).
119pub struct UnionListState {
120    state: Arc<RwLock<Vec<Vec<u8>>>>,
121}
122
123impl UnionListState {
124    /// Create a new union list state.
125    pub fn new() -> Self {
126        Self {
127            state: Arc::new(RwLock::new(Vec::new())),
128        }
129    }
130
131    /// Get all values.
132    pub async fn get(&self) -> Vec<Vec<u8>> {
133        self.state.read().await.clone()
134    }
135
136    /// Add a value.
137    pub async fn add(&self, value: Vec<u8>) {
138        self.state.write().await.push(value);
139    }
140
141    /// Add multiple values.
142    pub async fn add_all(&self, values: Vec<Vec<u8>>) {
143        self.state.write().await.extend(values);
144    }
145
146    /// Update with new values.
147    pub async fn update(&self, values: Vec<Vec<u8>>) {
148        *self.state.write().await = values;
149    }
150
151    /// Clear all values.
152    pub async fn clear(&self) {
153        self.state.write().await.clear();
154    }
155
156    /// Get the number of values.
157    pub async fn len(&self) -> usize {
158        self.state.read().await.len()
159    }
160
161    /// Check if empty.
162    pub async fn is_empty(&self) -> bool {
163        self.state.read().await.is_empty()
164    }
165}
166
167impl Default for UnionListState {
168    fn default() -> Self {
169        Self::new()
170    }
171}
172
173impl OperatorState for UnionListState {
174    async fn snapshot(&self) -> Result<Vec<u8>> {
175        let state = self.state.read().await;
176        Ok(serde_json::to_vec(&*state)?)
177    }
178
179    async fn restore(&self, snapshot: &[u8]) -> Result<()> {
180        let restored: Vec<Vec<u8>> = serde_json::from_slice(snapshot)?;
181        *self.state.write().await = restored;
182        Ok(())
183    }
184}
185
186/// Trait for checkpointable list state.
187pub trait ListCheckpointed {
188    /// Get the current state as a list.
189    fn snapshot_state(&self) -> impl std::future::Future<Output = Vec<Vec<u8>>> + Send;
190
191    /// Restore state from a list.
192    fn restore_state(
193        &self,
194        state: Vec<Vec<u8>>,
195    ) -> impl std::future::Future<Output = Result<()>> + Send;
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201
202    #[tokio::test]
203    async fn test_broadcast_state() {
204        let state = BroadcastState::new();
205
206        state.put(vec![1], vec![42]).await;
207        assert_eq!(state.get(&[1]).await, Some(vec![42]));
208
209        assert!(state.contains(&[1]).await);
210        assert!(!state.contains(&[2]).await);
211
212        state.remove(&[1]).await;
213        assert_eq!(state.get(&[1]).await, None);
214    }
215
216    #[tokio::test]
217    async fn test_broadcast_state_snapshot() {
218        let state = BroadcastState::new();
219
220        state.put(vec![1], vec![42]).await;
221        state.put(vec![2], vec![43]).await;
222
223        let snapshot = state
224            .snapshot()
225            .await
226            .expect("Failed to create snapshot of broadcast state");
227
228        let state2 = BroadcastState::new();
229        state2
230            .restore(&snapshot)
231            .await
232            .expect("Failed to restore broadcast state from snapshot");
233
234        assert_eq!(state2.get(&[1]).await, Some(vec![42]));
235        assert_eq!(state2.get(&[2]).await, Some(vec![43]));
236    }
237
238    #[tokio::test]
239    async fn test_union_list_state() {
240        let state = UnionListState::new();
241
242        state.add(vec![1]).await;
243        state.add(vec![2]).await;
244        state.add(vec![3]).await;
245
246        let values = state.get().await;
247        assert_eq!(values, vec![vec![1], vec![2], vec![3]]);
248
249        assert_eq!(state.len().await, 3);
250        assert!(!state.is_empty().await);
251    }
252
253    #[tokio::test]
254    async fn test_union_list_state_snapshot() {
255        let state = UnionListState::new();
256
257        state.add(vec![1]).await;
258        state.add(vec![2]).await;
259
260        let snapshot = state
261            .snapshot()
262            .await
263            .expect("Failed to create snapshot of union list state");
264
265        let state2 = UnionListState::new();
266        state2
267            .restore(&snapshot)
268            .await
269            .expect("Failed to restore union list state from snapshot");
270
271        assert_eq!(state2.get().await, vec![vec![1], vec![2]]);
272    }
273}
oxigdal_streaming/state/operator_state.rs

oxigdal_streaming/state/
operator_state.rs
