Skip to main content

beam_core/
workflow.rs

1use std::collections::HashMap;
2use std::fs::{self, OpenOptions};
3use std::io::{Read, Write};
4use std::path::{Path, PathBuf};
5use std::sync::{Arc, Mutex, OnceLock};
6use std::thread;
7use std::time::Duration;
8
9use anyhow::{Context, Result};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12
13#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
14#[serde(rename_all = "snake_case")]
15pub enum WorkflowActor {
16    Scheduler,
17    Human,
18    System,
19    Worker,
20    HostExecutor,
21    Supervisor,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
25pub struct WorkflowEventEnvelope {
26    pub event_id: String,
27    pub run_id: String,
28    pub timestamp: u64,
29    pub schema_version: u32,
30    pub actor: WorkflowActor,
31    #[serde(rename = "type")]
32    pub event_type: String,
33    pub payload: Value,
34    #[serde(default, skip_serializing_if = "Option::is_none")]
35    pub payload_hash: Option<String>,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
39pub struct EventDraft {
40    pub event_type: String,
41    pub actor: WorkflowActor,
42    pub payload: Value,
43    #[serde(default)]
44    pub timestamp: Option<u64>,
45    #[serde(default)]
46    pub payload_hash: Option<String>,
47}
48
49static RUN_MUTEXES: OnceLock<Mutex<HashMap<String, Arc<Mutex<()>>>>> = OnceLock::new();
50
51fn run_mutex(run_id: &str) -> Arc<Mutex<()>> {
52    let map = RUN_MUTEXES.get_or_init(|| Mutex::new(HashMap::new()));
53    let mut guard = map.lock().expect("workflow mutex map poisoned");
54    guard
55        .entry(run_id.to_string())
56        .or_insert_with(|| Arc::new(Mutex::new(())))
57        .clone()
58}
59
60#[derive(Debug, Clone)]
61pub struct EventLog {
62    pub run_id: String,
63    pub run_dir: PathBuf,
64    pub events_file: PathBuf,
65    pub blob_dir: PathBuf,
66    seq: u64,
67    seq_loaded: bool,
68    cached_len: u64,
69}
70
71impl EventLog {
72    pub fn new(run_id: impl Into<String>, base_dir: impl AsRef<Path>) -> Result<Self> {
73        let run_id = run_id.into();
74        if run_id.trim().is_empty() {
75            anyhow::bail!("EventLog: runId required");
76        }
77        let base_dir = base_dir.as_ref();
78        if base_dir.as_os_str().is_empty() {
79            anyhow::bail!("EventLog: baseDir required");
80        }
81        let run_dir = base_dir.join(&run_id);
82        let events_file = run_dir.join("events.ndjson");
83        let blob_dir = run_dir.join("blobs");
84        fs::create_dir_all(&blob_dir).context("failed to create workflow run directories")?;
85        Ok(Self {
86            run_id,
87            run_dir,
88            events_file,
89            blob_dir,
90            seq: 0,
91            seq_loaded: false,
92            cached_len: 0,
93        })
94    }
95
96    pub fn append(&mut self, draft: EventDraft) -> Result<WorkflowEventEnvelope> {
97        let mutex = run_mutex(&self.run_id);
98        let _guard = mutex.lock().expect("workflow run mutex poisoned");
99        let _lock = FileLock::acquire(&self.events_file)?;
100        self.refresh_seq_if_stale()?;
101
102        let next_seq = self.seq + 1;
103        let timestamp = draft.timestamp.unwrap_or_else(now_ms);
104        let candidate = WorkflowEventEnvelope {
105            event_id: format!("{}-{}", self.run_id, next_seq),
106            run_id: self.run_id.clone(),
107            timestamp,
108            schema_version: 1,
109            actor: draft.actor,
110            event_type: draft.event_type,
111            payload: draft.payload,
112            payload_hash: draft.payload_hash,
113        };
114
115        let line = serde_json::to_string(&candidate)? + "\n";
116        let mut file = OpenOptions::new()
117            .create(true)
118            .append(true)
119            .open(&self.events_file)
120            .with_context(|| format!("failed to open {}", self.events_file.display()))?;
121        file.write_all(line.as_bytes())?;
122        file.flush()?;
123
124        self.seq = next_seq;
125        self.seq_loaded = true;
126        self.cached_len = fs::metadata(&self.events_file)?.len();
127        Ok(candidate)
128    }
129
130    pub fn read_all(&self) -> Result<Vec<WorkflowEventEnvelope>> {
131        if !self.events_file.exists() {
132            return Ok(Vec::new());
133        }
134        let mut raw = String::new();
135        fs::File::open(&self.events_file)?.read_to_string(&mut raw)?;
136        let mut events = Vec::new();
137        for (idx, line) in raw.lines().enumerate() {
138            if line.trim().is_empty() {
139                continue;
140            }
141            let event = serde_json::from_str::<WorkflowEventEnvelope>(line)
142                .with_context(|| format!("corrupt event at line {}", idx + 1))?;
143            events.push(event);
144        }
145        Ok(events)
146    }
147
148    pub fn read_blob(&self, reference: &str) -> Result<Vec<u8>> {
149        Ok(fs::read(reference)?)
150    }
151
152    pub fn current_seq(&mut self) -> Result<u64> {
153        self.refresh_seq_if_stale()?;
154        Ok(self.seq)
155    }
156
157    fn refresh_seq_if_stale(&mut self) -> Result<()> {
158        if !self.events_file.exists() {
159            self.seq = 0;
160            self.cached_len = 0;
161            self.seq_loaded = true;
162            return Ok(());
163        }
164        let len = fs::metadata(&self.events_file)?.len();
165        if self.seq_loaded && len == self.cached_len {
166            return Ok(());
167        }
168        let events = self.read_all()?;
169        self.seq = events
170            .iter()
171            .filter_map(|event| event.event_id.rsplit_once('-'))
172            .filter_map(|(_, seq)| seq.parse::<u64>().ok())
173            .max()
174            .unwrap_or(0);
175        self.cached_len = len;
176        self.seq_loaded = true;
177        Ok(())
178    }
179}
180
181struct FileLock {
182    path: PathBuf,
183}
184
185impl FileLock {
186    fn acquire(target: &Path) -> Result<Self> {
187        let path = target.with_extension("lock");
188        for _ in 0..300 {
189            match OpenOptions::new().write(true).create_new(true).open(&path) {
190                Ok(_) => return Ok(Self { path }),
191                Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
192                    thread::sleep(Duration::from_millis(10));
193                }
194                Err(err) => return Err(err.into()),
195            }
196        }
197        anyhow::bail!("timed out acquiring workflow file lock: {}", path.display());
198    }
199}
200
201impl Drop for FileLock {
202    fn drop(&mut self) {
203        let _ = fs::remove_file(&self.path);
204    }
205}
206
207fn now_ms() -> u64 {
208    let now = std::time::SystemTime::now();
209    now.duration_since(std::time::UNIX_EPOCH)
210        .unwrap_or_default()
211        .as_millis() as u64
212}