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::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    pub fn new() -> Self {
21        Self::default()
22    }
23}
24
25#[async_trait]
26impl FlowEventStore for InMemoryEventStore {
27    async fn append(&self, run_id: &str, event: FlowEvent) -> Result<FlowEventEnvelope> {
28        let mut runs = self.runs.lock().await;
29        ensure_linked_flow_run_exists(&runs, &event)?;
30        append_in_memory(&mut runs, run_id, event)
31    }
32
33    async fn append_if_sequence(
34        &self,
35        run_id: &str,
36        expected_sequence: u64,
37        event: FlowEvent,
38    ) -> Result<FlowEventEnvelope> {
39        let mut runs = self.runs.lock().await;
40        ensure_linked_flow_run_exists(&runs, &event)?;
41        let actual_sequence = runs
42            .get(run_id)
43            .and_then(|events| events.last())
44            .map_or(0, |event| event.sequence);
45        if actual_sequence != expected_sequence {
46            return Err(FlowError::EventConflict {
47                run_id: run_id.to_string(),
48                expected_sequence,
49                actual_sequence,
50            });
51        }
52        append_in_memory(&mut runs, run_id, event)
53    }
54
55    async fn list(&self, run_id: &str) -> Result<Vec<FlowEventEnvelope>> {
56        let runs = self.runs.lock().await;
57        match runs.get(run_id) {
58            Some(events) => Ok(events.clone()),
59            None => Err(FlowError::RunNotFound(run_id.to_string())),
60        }
61    }
62
63    async fn list_run_ids(&self) -> Result<Vec<String>> {
64        let runs = self.runs.lock().await;
65        let mut ids: Vec<String> = runs.keys().cloned().collect();
66        ids.sort();
67        Ok(ids)
68    }
69}
70
71fn ensure_linked_flow_run_exists(
72    runs: &HashMap<String, Vec<FlowEventEnvelope>>,
73    event: &FlowEvent,
74) -> Result<()> {
75    let Some(linked_run_id) = linked_flow_run_id(event) else {
76        return Ok(());
77    };
78    if runs.get(linked_run_id).is_none_or(Vec::is_empty) {
79        return Err(FlowError::RunNotFound(linked_run_id.to_string()));
80    }
81    Ok(())
82}
83
84fn append_in_memory(
85    runs: &mut HashMap<String, Vec<FlowEventEnvelope>>,
86    run_id: &str,
87    event: FlowEvent,
88) -> Result<FlowEventEnvelope> {
89    let events = runs.entry(run_id.to_string()).or_default();
90    let envelope = FlowEventEnvelope {
91        run_id: run_id.to_string(),
92        sequence: events.last().map_or(1, |event| event.sequence + 1),
93        event_id: Uuid::new_v4(),
94        timestamp: Utc::now(),
95        event,
96    };
97    events.push(envelope.clone());
98    Ok(envelope)
99}