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#[derive(Debug, Default)]
15pub struct InMemoryEventStore {
16 runs: Arc<Mutex<HashMap<String, Vec<FlowEventEnvelope>>>>,
17}
18
19impl InMemoryEventStore {
20 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}