floxide_core/distributed/
context_store.rs

1use crate::context::Context;
2use crate::merge::Merge;
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::Arc;
6use thiserror::Error;
7use tokio::sync::Mutex;
8#[derive(Debug, Error)]
9pub enum ContextStoreError {
10    #[error("I/O error: {0}")]
11    Io(String),
12    #[error("Other error: {0}")]
13    Other(String),
14}
15
16#[async_trait]
17pub trait ContextStore<C: Context + Merge + Default>: Send + Sync {
18    async fn get(&self, run_id: &str) -> Result<Option<C>, ContextStoreError>;
19    async fn set(&self, run_id: &str, ctx: C) -> Result<(), ContextStoreError>;
20    async fn merge(&self, run_id: &str, ctx: C) -> Result<(), ContextStoreError>;
21}
22
23/// In-memory implementation for testing and local runs.
24pub struct InMemoryContextStore<C: Context + Merge + Default> {
25    inner: Arc<Mutex<HashMap<String, C>>>,
26}
27
28impl<C: Context + Merge + Default> Default for InMemoryContextStore<C> {
29    fn default() -> Self {
30        Self {
31            inner: Arc::new(Mutex::new(HashMap::new())),
32        }
33    }
34}
35
36impl<C: Context + Merge + Default> Clone for InMemoryContextStore<C> {
37    fn clone(&self) -> Self {
38        Self {
39            inner: self.inner.clone(),
40        }
41    }
42}
43
44#[async_trait]
45impl<C: Context + Merge + Default> ContextStore<C> for InMemoryContextStore<C> {
46    async fn get(&self, run_id: &str) -> Result<Option<C>, ContextStoreError> {
47        let map = self.inner.lock().await;
48        Ok(map.get(run_id).cloned())
49    }
50    async fn set(&self, run_id: &str, ctx: C) -> Result<(), ContextStoreError> {
51        let mut map = self.inner.lock().await;
52        map.insert(run_id.to_string(), ctx);
53        Ok(())
54    }
55    async fn merge(&self, run_id: &str, ctx: C) -> Result<(), ContextStoreError> {
56        let mut map = self.inner.lock().await;
57        map.entry(run_id.to_string())
58            .and_modify(|existing| existing.merge(ctx.clone()))
59            .or_insert(ctx);
60        Ok(())
61    }
62}
floxide_core/distributed/context_store.rs

floxide_core/distributed/
context_store.rs
