a3s_code_core/flow_graph/
decision_ledger.rs

1use anyhow::{Context, Result};
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4use std::collections::BTreeMap;
5use std::path::{Path, PathBuf};
6use tokio::io::AsyncWriteExt;
7use tokio::sync::Mutex;
8
9const MAX_DECISION_LEDGER_BYTES: usize = 16 * 1024 * 1024;
10const DECISION_LEDGER_SCHEMA_VERSION: u32 = 1;
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub enum FlowDecisionClaimOutcome {
14    Claimed { attempt: u32 },
15    Completed,
16    Busy { lease_expires_at_ms: u64 },
17    Conflict,
18}
19
20#[async_trait]
21pub trait FlowDecisionLedger: Send + Sync {
22    async fn claim(
23        &self,
24        decision_id: &str,
25        request_hash: &str,
26        owner_id: &str,
27        now_ms: u64,
28        lease_ms: u64,
29    ) -> Result<FlowDecisionClaimOutcome>;
30
31    /// Extend a pending claim only while it is still owned by `owner_id`.
32    /// Returns `false` after completion, takeover, release, or identity conflict.
33    async fn renew(
34        &self,
35        decision_id: &str,
36        request_hash: &str,
37        owner_id: &str,
38        now_ms: u64,
39        lease_ms: u64,
40    ) -> Result<bool>;
41
42    async fn complete(
43        &self,
44        decision_id: &str,
45        request_hash: &str,
46        owner_id: &str,
47        completed_at_ms: u64,
48    ) -> Result<()>;
49
50    async fn release(&self, decision_id: &str, request_hash: &str, owner_id: &str) -> Result<()>;
51
52    /// Remove completed receipts older than the host's retention cutoff.
53    async fn prune_completed(&self, _before_ms: u64) -> Result<usize> {
54        anyhow::bail!("Flow decision ledger does not support receipt pruning")
55    }
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
59#[serde(rename_all = "snake_case")]
60enum ClaimStatus {
61    Pending,
62    Completed,
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize)]
66struct ClaimRecord {
67    request_hash: String,
68    status: ClaimStatus,
69    owner_id: String,
70    lease_expires_at_ms: u64,
71    attempts: u32,
72    #[serde(default, skip_serializing_if = "Option::is_none")]
73    completed_at_ms: Option<u64>,
74}
75
76#[derive(Debug, Default, Serialize, Deserialize)]
77struct DecisionLedgerFile {
78    schema_version: u32,
79    records: BTreeMap<String, ClaimRecord>,
80}
81
82#[derive(Serialize)]
83struct DecisionLedgerFileRef<'a> {
84    schema_version: u32,
85    records: &'a BTreeMap<String, ClaimRecord>,
86}
87
88#[derive(Debug, Default)]
89pub struct MemoryFlowDecisionLedger {
90    records: Mutex<BTreeMap<String, ClaimRecord>>,
91}
92
93impl MemoryFlowDecisionLedger {
94    pub fn new() -> Self {
95        Self::default()
96    }
97}
98
99#[async_trait]
100impl FlowDecisionLedger for MemoryFlowDecisionLedger {
101    async fn claim(
102        &self,
103        decision_id: &str,
104        request_hash: &str,
105        owner_id: &str,
106        now_ms: u64,
107        lease_ms: u64,
108    ) -> Result<FlowDecisionClaimOutcome> {
109        let mut records = self.records.lock().await;
110        claim_record(
111            &mut records,
112            decision_id,
113            request_hash,
114            owner_id,
115            now_ms,
116            lease_ms,
117        )
118    }
119
120    async fn renew(
121        &self,
122        decision_id: &str,
123        request_hash: &str,
124        owner_id: &str,
125        now_ms: u64,
126        lease_ms: u64,
127    ) -> Result<bool> {
128        let mut records = self.records.lock().await;
129        Ok(renew_record(
130            &mut records,
131            decision_id,
132            request_hash,
133            owner_id,
134            now_ms,
135            lease_ms,
136        ))
137    }
138
139    async fn complete(
140        &self,
141        decision_id: &str,
142        request_hash: &str,
143        owner_id: &str,
144        completed_at_ms: u64,
145    ) -> Result<()> {
146        let mut records = self.records.lock().await;
147        complete_record(
148            &mut records,
149            decision_id,
150            request_hash,
151            owner_id,
152            completed_at_ms,
153        )
154    }
155
156    async fn release(&self, decision_id: &str, request_hash: &str, owner_id: &str) -> Result<()> {
157        let mut records = self.records.lock().await;
158        release_record(&mut records, decision_id, request_hash, owner_id)
159    }
160
161    async fn prune_completed(&self, before_ms: u64) -> Result<usize> {
162        let mut records = self.records.lock().await;
163        Ok(prune_records(&mut records, before_ms))
164    }
165}
166
167#[derive(Debug)]
168pub struct FileFlowDecisionLedger {
169    root: PathBuf,
170    process_lock: Mutex<()>,
171}
172
173impl FileFlowDecisionLedger {
174    pub fn new(root: impl Into<PathBuf>) -> Self {
175        Self {
176            root: root.into(),
177            process_lock: Mutex::new(()),
178        }
179    }
180
181    fn data_path(&self) -> PathBuf {
182        self.root.join("flow-decisions.json")
183    }
184
185    async fn acquire_file_lock(&self) -> Result<std::fs::File> {
186        tokio::fs::create_dir_all(&self.root)
187            .await
188            .with_context(|| format!("create decision ledger `{}`", self.root.display()))?;
189        let path = self.root.join(".flow-decisions.lock");
190        tokio::task::spawn_blocking(move || {
191            use fs2::FileExt;
192            let file = std::fs::OpenOptions::new()
193                .create(true)
194                .truncate(false)
195                .read(true)
196                .write(true)
197                .open(&path)
198                .with_context(|| format!("open decision ledger lock `{}`", path.display()))?;
199            file.lock_exclusive()
200                .with_context(|| format!("lock decision ledger `{}`", path.display()))?;
201            Ok(file)
202        })
203        .await
204        .context("decision ledger lock task failed")?
205    }
206
207    async fn mutate<T>(
208        &self,
209        mutation: impl FnOnce(&mut BTreeMap<String, ClaimRecord>) -> Result<T>,
210    ) -> Result<T> {
211        let _process_guard = self.process_lock.lock().await;
212        let _file_guard = self.acquire_file_lock().await?;
213        let mut records = read_records(&self.data_path()).await?;
214        let result = mutation(&mut records)?;
215        write_records(&self.data_path(), &records).await?;
216        Ok(result)
217    }
218}
219
220#[async_trait]
221impl FlowDecisionLedger for FileFlowDecisionLedger {
222    async fn claim(
223        &self,
224        decision_id: &str,
225        request_hash: &str,
226        owner_id: &str,
227        now_ms: u64,
228        lease_ms: u64,
229    ) -> Result<FlowDecisionClaimOutcome> {
230        self.mutate(|records| {
231            claim_record(
232                records,
233                decision_id,
234                request_hash,
235                owner_id,
236                now_ms,
237                lease_ms,
238            )
239        })
240        .await
241    }
242
243    async fn renew(
244        &self,
245        decision_id: &str,
246        request_hash: &str,
247        owner_id: &str,
248        now_ms: u64,
249        lease_ms: u64,
250    ) -> Result<bool> {
251        self.mutate(|records| {
252            Ok(renew_record(
253                records,
254                decision_id,
255                request_hash,
256                owner_id,
257                now_ms,
258                lease_ms,
259            ))
260        })
261        .await
262    }
263
264    async fn complete(
265        &self,
266        decision_id: &str,
267        request_hash: &str,
268        owner_id: &str,
269        completed_at_ms: u64,
270    ) -> Result<()> {
271        self.mutate(|records| {
272            complete_record(
273                records,
274                decision_id,
275                request_hash,
276                owner_id,
277                completed_at_ms,
278            )
279        })
280        .await
281    }
282
283    async fn release(&self, decision_id: &str, request_hash: &str, owner_id: &str) -> Result<()> {
284        self.mutate(|records| release_record(records, decision_id, request_hash, owner_id))
285            .await
286    }
287
288    async fn prune_completed(&self, before_ms: u64) -> Result<usize> {
289        self.mutate(|records| Ok(prune_records(records, before_ms)))
290            .await
291    }
292}
293
294fn claim_record(
295    records: &mut BTreeMap<String, ClaimRecord>,
296    decision_id: &str,
297    request_hash: &str,
298    owner_id: &str,
299    now_ms: u64,
300    lease_ms: u64,
301) -> Result<FlowDecisionClaimOutcome> {
302    let lease_expires_at_ms = now_ms.saturating_add(lease_ms.max(1));
303    match records.get_mut(decision_id) {
304        None => {
305            records.insert(
306                decision_id.to_string(),
307                ClaimRecord {
308                    request_hash: request_hash.to_string(),
309                    status: ClaimStatus::Pending,
310                    owner_id: owner_id.to_string(),
311                    lease_expires_at_ms,
312                    attempts: 1,
313                    completed_at_ms: None,
314                },
315            );
316            Ok(FlowDecisionClaimOutcome::Claimed { attempt: 1 })
317        }
318        Some(record) if record.request_hash != request_hash => {
319            Ok(FlowDecisionClaimOutcome::Conflict)
320        }
321        Some(record) if record.status == ClaimStatus::Completed => {
322            Ok(FlowDecisionClaimOutcome::Completed)
323        }
324        Some(record) if record.lease_expires_at_ms > now_ms && record.owner_id != owner_id => {
325            Ok(FlowDecisionClaimOutcome::Busy {
326                lease_expires_at_ms: record.lease_expires_at_ms,
327            })
328        }
329        Some(record) => {
330            record.owner_id = owner_id.to_string();
331            record.lease_expires_at_ms = lease_expires_at_ms;
332            record.attempts = record.attempts.saturating_add(1);
333            Ok(FlowDecisionClaimOutcome::Claimed {
334                attempt: record.attempts,
335            })
336        }
337    }
338}
339
340fn complete_record(
341    records: &mut BTreeMap<String, ClaimRecord>,
342    decision_id: &str,
343    request_hash: &str,
344    owner_id: &str,
345    completed_at_ms: u64,
346) -> Result<()> {
347    let record = records
348        .get_mut(decision_id)
349        .with_context(|| format!("decision claim `{decision_id}` does not exist"))?;
350    if record.request_hash != request_hash {
351        anyhow::bail!("decision `{decision_id}` request hash conflicts with its claim");
352    }
353    if record.status == ClaimStatus::Completed {
354        return Ok(());
355    }
356    if record.owner_id != owner_id {
357        anyhow::bail!("decision `{decision_id}` is owned by another dispatcher");
358    }
359    record.status = ClaimStatus::Completed;
360    record.lease_expires_at_ms = 0;
361    record.completed_at_ms = Some(completed_at_ms);
362    Ok(())
363}
364
365fn renew_record(
366    records: &mut BTreeMap<String, ClaimRecord>,
367    decision_id: &str,
368    request_hash: &str,
369    owner_id: &str,
370    now_ms: u64,
371    lease_ms: u64,
372) -> bool {
373    let Some(record) = records.get_mut(decision_id) else {
374        return false;
375    };
376    if record.request_hash != request_hash
377        || record.status != ClaimStatus::Pending
378        || record.owner_id != owner_id
379    {
380        return false;
381    }
382    record.lease_expires_at_ms = now_ms.saturating_add(lease_ms.max(1));
383    true
384}
385
386fn release_record(
387    records: &mut BTreeMap<String, ClaimRecord>,
388    decision_id: &str,
389    request_hash: &str,
390    owner_id: &str,
391) -> Result<()> {
392    let Some(record) = records.get_mut(decision_id) else {
393        return Ok(());
394    };
395    if record.request_hash != request_hash || record.status == ClaimStatus::Completed {
396        return Ok(());
397    }
398    if record.owner_id == owner_id {
399        record.owner_id.clear();
400        record.lease_expires_at_ms = 0;
401    }
402    Ok(())
403}
404
405fn prune_records(records: &mut BTreeMap<String, ClaimRecord>, before_ms: u64) -> usize {
406    let before = records.len();
407    records.retain(|_, record| {
408        record.status != ClaimStatus::Completed
409            || record.completed_at_ms.unwrap_or(u64::MAX) >= before_ms
410    });
411    before - records.len()
412}
413
414async fn read_records(path: &Path) -> Result<BTreeMap<String, ClaimRecord>> {
415    let bytes = match tokio::fs::read(path).await {
416        Ok(bytes) => bytes,
417        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(BTreeMap::new()),
418        Err(error) => return Err(error).context("read decision ledger"),
419    };
420    if bytes.len() > MAX_DECISION_LEDGER_BYTES {
421        anyhow::bail!("decision ledger exceeds {MAX_DECISION_LEDGER_BYTES} bytes");
422    }
423    let ledger: DecisionLedgerFile =
424        serde_json::from_slice(&bytes).context("decode decision ledger")?;
425    if ledger.schema_version > DECISION_LEDGER_SCHEMA_VERSION {
426        anyhow::bail!(
427            "decision ledger schema {} is newer than supported schema {}",
428            ledger.schema_version,
429            DECISION_LEDGER_SCHEMA_VERSION
430        );
431    }
432    Ok(ledger.records)
433}
434
435async fn write_records(path: &Path, records: &BTreeMap<String, ClaimRecord>) -> Result<()> {
436    let bytes = serde_json::to_vec(&DecisionLedgerFileRef {
437        schema_version: DECISION_LEDGER_SCHEMA_VERSION,
438        records,
439    })
440    .context("encode decision ledger")?;
441    if bytes.len() > MAX_DECISION_LEDGER_BYTES {
442        anyhow::bail!("decision ledger exceeds {MAX_DECISION_LEDGER_BYTES} bytes");
443    }
444    let temp = path.with_extension(format!("tmp-{}", uuid::Uuid::new_v4()));
445    let result = async {
446        let mut file = tokio::fs::File::create(&temp)
447            .await
448            .context("create decision ledger generation")?;
449        file.write_all(&bytes)
450            .await
451            .context("write decision ledger")?;
452        file.sync_all().await.context("sync decision ledger")?;
453        drop(file);
454        let temp_copy = temp.clone();
455        let path = path.to_path_buf();
456        tokio::task::spawn_blocking(move || {
457            tempfile::TempPath::try_from_path(temp_copy)?
458                .persist(path)
459                .map_err(|error| error.error)
460        })
461        .await
462        .context("publish decision ledger task failed")??;
463        Ok::<_, anyhow::Error>(())
464    }
465    .await;
466    if result.is_err() {
467        let _ = tokio::fs::remove_file(temp).await;
468    }
469    result
470}
a3s_code_core/flow_graph/decision_ledger.rs

a3s_code_core/flow_graph/
decision_ledger.rs
