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#[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}