Skip to main content

a3s_flow/store/
memory.rs

1use async_trait::async_trait;
2use chrono::Utc;
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::Mutex;
6use uuid::Uuid;
7
8use crate::error::{FlowError, Result};
9use crate::model::{FlowEvent, FlowEventEnvelope};
10
11use super::{retention::required_linked_flow_run_id, FlowEventStore};
12
13/// In-memory event store for tests, local development, and embedded hosts.
14#[derive(Debug, Default)]
15pub struct InMemoryEventStore {
16    runs: Arc<Mutex<HashMap<String, Vec<FlowEventEnvelope>>>>,
17}
18
19impl InMemoryEventStore {
20    /// Create an empty in-memory event store.
21    pub fn new() -> Self {
22        Self::default()
23    }
24}
25
26#[async_trait]
27impl FlowEventStore for InMemoryEventStore {
28    async fn append(&self, run_id: &str, event: FlowEvent) -> Result<FlowEventEnvelope> {
29        let mut runs = self.runs.lock().await;
30        ensure_linked_flow_run_exists(&runs, &event)?;
31        append_in_memory(&mut runs, run_id, event)
32    }
33
34    async fn append_if_sequence(
35        &self,
36        run_id: &str,
37        expected_sequence: u64,
38        event: FlowEvent,
39    ) -> Result<FlowEventEnvelope> {
40        let mut runs = self.runs.lock().await;
41        ensure_linked_flow_run_exists(&runs, &event)?;
42        let actual_sequence = runs
43            .get(run_id)
44            .and_then(|events| events.last())
45            .map_or(0, |event| event.sequence);
46        if actual_sequence != expected_sequence {
47            return Err(FlowError::EventConflict {
48                run_id: run_id.to_string(),
49                expected_sequence,
50                actual_sequence,
51            });
52        }
53        append_in_memory(&mut runs, run_id, event)
54    }
55
56    async fn list(&self, run_id: &str) -> Result<Vec<FlowEventEnvelope>> {
57        let runs = self.runs.lock().await;
58        match runs.get(run_id) {
59            Some(events) => Ok(events.clone()),
60            None => Err(FlowError::RunNotFound(run_id.to_string())),
61        }
62    }
63
64    async fn list_run_ids(&self) -> Result<Vec<String>> {
65        let runs = self.runs.lock().await;
66        let mut ids: Vec<String> = runs.keys().cloned().collect();
67        ids.sort();
68        Ok(ids)
69    }
70}
71
72fn ensure_linked_flow_run_exists(
73    runs: &HashMap<String, Vec<FlowEventEnvelope>>,
74    event: &FlowEvent,
75) -> Result<()> {
76    let Some(linked_run_id) = required_linked_flow_run_id(event) else {
77        return Ok(());
78    };
79    if runs.get(linked_run_id).is_none_or(Vec::is_empty) {
80        return Err(FlowError::RunNotFound(linked_run_id.to_string()));
81    }
82    Ok(())
83}
84
85fn append_in_memory(
86    runs: &mut HashMap<String, Vec<FlowEventEnvelope>>,
87    run_id: &str,
88    event: FlowEvent,
89) -> Result<FlowEventEnvelope> {
90    let events = runs.entry(run_id.to_string()).or_default();
91    let envelope = FlowEventEnvelope {
92        run_id: run_id.to_string(),
93        sequence: events.last().map_or(1, |event| event.sequence + 1),
94        event_id: Uuid::new_v4(),
95        timestamp: Utc::now(),
96        event,
97    };
98    events.push(envelope.clone());
99    Ok(envelope)
100}