tinyflow_api/state/
mod.rs

1use std::collections::HashMap;
2use std::path::Path;
3use std::sync::{Arc, Mutex};
4
5use async_trait::async_trait;
6
7use crate::error::{StreamResult, task_failed};
8
9#[async_trait]
10pub trait StateStore: Send + Sync {
11    async fn put(&self, namespace: &str, key: &[u8], value: &[u8]) -> StreamResult<()>;
12    async fn get(&self, namespace: &str, key: &[u8]) -> StreamResult<Option<Vec<u8>>>;
13    async fn delete(&self, namespace: &str, key: &[u8]) -> StreamResult<()>;
14    async fn delete_prefix(&self, namespace: &str, key_prefix: &[u8]) -> StreamResult<()>;
15    async fn list_prefix(
16        &self,
17        namespace: &str,
18        key_prefix: &[u8],
19    ) -> StreamResult<Vec<(Vec<u8>, Vec<u8>)>>;
20
21    /// Persistent backends (e.g., SQLite) return the database file path; in-memory stores return `None`.
22    fn persistent_db_path(&self) -> Option<&Path> {
23        None
24    }
25}
26
27pub fn state_db_err(reason: impl Into<String>) -> crate::error::StreamError {
28    task_failed(0, reason)
29}
30
31pub fn in_memory_job_state() -> Arc<dyn StateStore> {
32    Arc::new(InMemoryStateStore::default())
33}
34
35type InMemoryStoreMap = HashMap<(String, Vec<u8>), Vec<u8>>;
36
37#[derive(Default)]
38pub struct InMemoryStateStore {
39    inner: Mutex<InMemoryStoreMap>,
40}
41
42#[async_trait]
43impl StateStore for InMemoryStateStore {
44    async fn put(&self, namespace: &str, key: &[u8], value: &[u8]) -> StreamResult<()> {
45        self.inner
46            .lock()
47            .unwrap()
48            .insert((namespace.to_string(), key.to_vec()), value.to_vec());
49        Ok(())
50    }
51
52    async fn get(&self, namespace: &str, key: &[u8]) -> StreamResult<Option<Vec<u8>>> {
53        Ok(self
54            .inner
55            .lock()
56            .unwrap()
57            .get(&(namespace.to_string(), key.to_vec()))
58            .cloned())
59    }
60
61    async fn delete(&self, namespace: &str, key: &[u8]) -> StreamResult<()> {
62        self.inner
63            .lock()
64            .unwrap()
65            .remove(&(namespace.to_string(), key.to_vec()));
66        Ok(())
67    }
68
69    async fn delete_prefix(&self, namespace: &str, key_prefix: &[u8]) -> StreamResult<()> {
70        let ns = namespace.to_string();
71        let mut guard = self.inner.lock().unwrap();
72        guard.retain(|(n, k), _| !(n == &ns && k.starts_with(key_prefix)));
73        Ok(())
74    }
75
76    async fn list_prefix(
77        &self,
78        namespace: &str,
79        key_prefix: &[u8],
80    ) -> StreamResult<Vec<(Vec<u8>, Vec<u8>)>> {
81        let ns = namespace.to_string();
82        let guard = self.inner.lock().unwrap();
83        let mut out: Vec<(Vec<u8>, Vec<u8>)> = guard
84            .iter()
85            .filter(|((n, k), _)| n == &ns && k.starts_with(key_prefix))
86            .map(|((_, k), v)| (k.clone(), v.clone()))
87            .collect();
88        out.sort_by(|a, b| a.0.cmp(&b.0));
89        Ok(out)
90    }
91}
tinyflow_api/state/mod.rs

tinyflow_api/state/
mod.rs
