Skip to main content

a3s_code_core/flow_graph/
decision.rs

1use super::decision_ledger::{
2    FlowDecisionClaimOutcome, FlowDecisionLedger, MemoryFlowDecisionLedger,
3};
4use a3s_flow::RetryPolicy;
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
9use std::sync::Arc;
10use std::time::{Duration, Instant};
11use thiserror::Error;
12use tokio::sync::Mutex;
13
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
15#[serde(tag = "type", rename_all = "snake_case")]
16pub enum FlowDecision {
17    ScheduleStep { step: FlowDecisionStep },
18    ScheduleSteps { steps: Vec<FlowDecisionStep> },
19    Complete { output: Value },
20    Fail { error: String },
21}
22
23#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
24pub struct FlowDecisionStep {
25    pub step_id: String,
26    pub step_name: String,
27    pub input: Value,
28    #[serde(default)]
29    pub retry: RetryPolicy,
30}
31
32/// A graph proposal submitted through a host-owned Flow boundary.
33#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
34pub struct FlowDecisionRequest {
35    pub decision_id: String,
36    pub run_id: String,
37    pub authority_branch_id: String,
38    pub causation_event_id: String,
39    pub decision: FlowDecision,
40}
41
42#[async_trait]
43pub trait FlowDecisionSink: Send + Sync {
44    /// Submit using `request.decision_id` as the downstream idempotency key.
45    /// Implementations must deduplicate that key because an expired lease can
46    /// be reclaimed after a process crashes between submission and receipt.
47    async fn submit(
48        &self,
49        request: &FlowDecisionRequest,
50    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
51}
52
53#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
54#[serde(rename_all = "snake_case")]
55pub enum FlowDecisionHealthStatus {
56    Healthy,
57    Degraded,
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
61pub struct FlowDecisionHealthSnapshot {
62    pub status: FlowDecisionHealthStatus,
63    pub attempted: u64,
64    pub claimed: u64,
65    pub takeovers: u64,
66    pub completed: u64,
67    pub duplicates: u64,
68    pub busy: u64,
69    pub conflicts: u64,
70    pub rejected: u64,
71    pub lease_renewals: u64,
72    pub lease_lost: u64,
73    pub sink_failures: u64,
74    pub ledger_failures: u64,
75    pub failures: u64,
76    pub cancellations: u64,
77    pub in_flight: u64,
78    pub average_dispatch_micros: u64,
79    pub max_dispatch_micros: u64,
80    pub last_success_at_ms: Option<u64>,
81    pub last_failure_at_ms: Option<u64>,
82}
83
84#[derive(Default)]
85struct FlowDecisionMetrics {
86    attempted: AtomicU64,
87    claimed: AtomicU64,
88    takeovers: AtomicU64,
89    completed: AtomicU64,
90    duplicates: AtomicU64,
91    busy: AtomicU64,
92    conflicts: AtomicU64,
93    rejected: AtomicU64,
94    lease_renewals: AtomicU64,
95    lease_lost: AtomicU64,
96    sink_failures: AtomicU64,
97    ledger_failures: AtomicU64,
98    failures: AtomicU64,
99    cancellations: AtomicU64,
100    in_flight: AtomicU64,
101    total_dispatch_micros: AtomicU64,
102    max_dispatch_micros: AtomicU64,
103    last_success_at_ms: AtomicU64,
104    last_failure_at_ms: AtomicU64,
105    degraded: AtomicBool,
106}
107
108/// Enforces production-branch authority and idempotent successful submission.
109pub struct FlowDecisionDispatcher {
110    production_branch_id: String,
111    sink: Arc<dyn FlowDecisionSink>,
112    ledger: Arc<dyn FlowDecisionLedger>,
113    owner_id: String,
114    lease_ms: u64,
115    dispatch_lock: Mutex<()>,
116    metrics: Arc<FlowDecisionMetrics>,
117}
118
119impl FlowDecisionDispatcher {
120    pub fn new(production_branch_id: impl Into<String>, sink: Arc<dyn FlowDecisionSink>) -> Self {
121        Self::with_ledger(
122            production_branch_id,
123            sink,
124            Arc::new(MemoryFlowDecisionLedger::new()),
125        )
126    }
127
128    pub fn with_ledger(
129        production_branch_id: impl Into<String>,
130        sink: Arc<dyn FlowDecisionSink>,
131        ledger: Arc<dyn FlowDecisionLedger>,
132    ) -> Self {
133        Self {
134            production_branch_id: production_branch_id.into(),
135            sink,
136            ledger,
137            owner_id: format!("decision-dispatcher-{}", uuid::Uuid::new_v4()),
138            lease_ms: 30_000,
139            dispatch_lock: Mutex::new(()),
140            metrics: Arc::new(FlowDecisionMetrics::default()),
141        }
142    }
143
144    pub fn with_lease_ms(mut self, lease_ms: u64) -> Self {
145        self.lease_ms = lease_ms.max(1);
146        self
147    }
148
149    pub fn health(&self) -> FlowDecisionHealthSnapshot {
150        let attempted = self.metrics.attempted.load(Ordering::Relaxed);
151        let total_micros = self.metrics.total_dispatch_micros.load(Ordering::Relaxed);
152        let last_success_at_ms = nonzero(self.metrics.last_success_at_ms.load(Ordering::Relaxed));
153        let last_failure_at_ms = nonzero(self.metrics.last_failure_at_ms.load(Ordering::Relaxed));
154        FlowDecisionHealthSnapshot {
155            status: if self.metrics.degraded.load(Ordering::Relaxed) {
156                FlowDecisionHealthStatus::Degraded
157            } else {
158                FlowDecisionHealthStatus::Healthy
159            },
160            attempted,
161            claimed: self.metrics.claimed.load(Ordering::Relaxed),
162            takeovers: self.metrics.takeovers.load(Ordering::Relaxed),
163            completed: self.metrics.completed.load(Ordering::Relaxed),
164            duplicates: self.metrics.duplicates.load(Ordering::Relaxed),
165            busy: self.metrics.busy.load(Ordering::Relaxed),
166            conflicts: self.metrics.conflicts.load(Ordering::Relaxed),
167            rejected: self.metrics.rejected.load(Ordering::Relaxed),
168            lease_renewals: self.metrics.lease_renewals.load(Ordering::Relaxed),
169            lease_lost: self.metrics.lease_lost.load(Ordering::Relaxed),
170            sink_failures: self.metrics.sink_failures.load(Ordering::Relaxed),
171            ledger_failures: self.metrics.ledger_failures.load(Ordering::Relaxed),
172            failures: self.metrics.failures.load(Ordering::Relaxed),
173            cancellations: self.metrics.cancellations.load(Ordering::Relaxed),
174            in_flight: self.metrics.in_flight.load(Ordering::Relaxed),
175            average_dispatch_micros: total_micros.checked_div(attempted).unwrap_or(0),
176            max_dispatch_micros: self.metrics.max_dispatch_micros.load(Ordering::Relaxed),
177            last_success_at_ms,
178            last_failure_at_ms,
179        }
180    }
181
182    /// Returns `true` only when the sink accepted a new decision.
183    pub async fn dispatch(
184        &self,
185        request: &FlowDecisionRequest,
186    ) -> Result<bool, FlowDecisionDispatchError> {
187        let mut in_flight = DecisionInFlight::new(Arc::clone(&self.metrics));
188        let result = self.dispatch_inner(request).await;
189        in_flight.finish();
190        self.record_result(&result);
191        result
192    }
193
194    async fn dispatch_inner(
195        &self,
196        request: &FlowDecisionRequest,
197    ) -> Result<bool, FlowDecisionDispatchError> {
198        if request.authority_branch_id != self.production_branch_id {
199            return Err(FlowDecisionDispatchError::UnauthorizedBranch {
200                expected: self.production_branch_id.clone(),
201                actual: request.authority_branch_id.clone(),
202            });
203        }
204        if request.decision_id.trim().is_empty() || request.causation_event_id.trim().is_empty() {
205            return Err(FlowDecisionDispatchError::InvalidIdentity);
206        }
207        let _dispatch_guard = self.dispatch_lock.lock().await;
208        let request_hash = request_hash(request)?;
209        match self
210            .ledger
211            .claim(
212                &request.decision_id,
213                &request_hash,
214                &self.owner_id,
215                now_ms(),
216                self.lease_ms,
217            )
218            .await
219            .map_err(|error| FlowDecisionDispatchError::Ledger(error.to_string()))?
220        {
221            FlowDecisionClaimOutcome::Completed => return Ok(false),
222            FlowDecisionClaimOutcome::Busy {
223                lease_expires_at_ms,
224            } => {
225                return Err(FlowDecisionDispatchError::Busy {
226                    lease_expires_at_ms,
227                })
228            }
229            FlowDecisionClaimOutcome::Conflict => {
230                return Err(FlowDecisionDispatchError::DecisionIdConflict(
231                    request.decision_id.clone(),
232                ))
233            }
234            FlowDecisionClaimOutcome::Claimed { attempt } => {
235                self.metrics.claimed.fetch_add(1, Ordering::Relaxed);
236                if attempt > 1 {
237                    self.metrics.takeovers.fetch_add(1, Ordering::Relaxed);
238                }
239            }
240        }
241        let heartbeat_period = Duration::from_millis((self.lease_ms / 3).max(1));
242        let first_heartbeat = tokio::time::Instant::now() + heartbeat_period;
243        let mut heartbeat = tokio::time::interval_at(first_heartbeat, heartbeat_period);
244        let submission = self.sink.submit(request);
245        tokio::pin!(submission);
246        let submission_result = loop {
247            tokio::select! {
248                result = &mut submission => break result,
249                _ = heartbeat.tick() => {
250                    let renewed = self.ledger
251                        .renew(
252                            &request.decision_id,
253                            &request_hash,
254                            &self.owner_id,
255                            now_ms(),
256                            self.lease_ms,
257                        )
258                        .await
259                        .map_err(|error| FlowDecisionDispatchError::Ledger(error.to_string()))?;
260                    if !renewed {
261                        return Err(FlowDecisionDispatchError::LeaseLost(
262                            request.decision_id.clone(),
263                        ));
264                    }
265                    self.metrics.lease_renewals.fetch_add(1, Ordering::Relaxed);
266                }
267            }
268        };
269        if let Err(error) = submission_result {
270            if let Err(release_error) = self
271                .ledger
272                .release(&request.decision_id, &request_hash, &self.owner_id)
273                .await
274            {
275                tracing::warn!(decision_id = request.decision_id, error = %release_error, "failed to release Flow decision claim");
276            }
277            return Err(FlowDecisionDispatchError::Sink(error.to_string()));
278        }
279        self.ledger
280            .complete(
281                &request.decision_id,
282                &request_hash,
283                &self.owner_id,
284                now_ms(),
285            )
286            .await
287            .map_err(|error| FlowDecisionDispatchError::Ledger(error.to_string()))?;
288        Ok(true)
289    }
290
291    fn record_result(&self, result: &Result<bool, FlowDecisionDispatchError>) {
292        match result {
293            Ok(true) => {
294                self.metrics.completed.fetch_add(1, Ordering::Relaxed);
295                self.metrics.degraded.store(false, Ordering::Relaxed);
296                self.metrics
297                    .last_success_at_ms
298                    .store(now_ms(), Ordering::Relaxed);
299            }
300            Ok(false) => {
301                self.metrics.duplicates.fetch_add(1, Ordering::Relaxed);
302                self.metrics.degraded.store(false, Ordering::Relaxed);
303                self.metrics
304                    .last_success_at_ms
305                    .store(now_ms(), Ordering::Relaxed);
306            }
307            Err(FlowDecisionDispatchError::Busy { .. }) => {
308                self.metrics.busy.fetch_add(1, Ordering::Relaxed);
309            }
310            Err(FlowDecisionDispatchError::DecisionIdConflict(_)) => {
311                self.metrics.conflicts.fetch_add(1, Ordering::Relaxed);
312            }
313            Err(
314                FlowDecisionDispatchError::UnauthorizedBranch { .. }
315                | FlowDecisionDispatchError::InvalidIdentity,
316            ) => {
317                self.metrics.rejected.fetch_add(1, Ordering::Relaxed);
318            }
319            Err(FlowDecisionDispatchError::LeaseLost(_)) => {
320                self.metrics.lease_lost.fetch_add(1, Ordering::Relaxed);
321                self.record_failure();
322            }
323            Err(FlowDecisionDispatchError::Ledger(_)) => {
324                self.metrics.ledger_failures.fetch_add(1, Ordering::Relaxed);
325                self.record_failure();
326            }
327            Err(FlowDecisionDispatchError::Sink(_)) => {
328                self.metrics.sink_failures.fetch_add(1, Ordering::Relaxed);
329                self.record_failure();
330            }
331        }
332    }
333
334    fn record_failure(&self) {
335        self.metrics.failures.fetch_add(1, Ordering::Relaxed);
336        self.metrics.degraded.store(true, Ordering::Relaxed);
337        self.metrics
338            .last_failure_at_ms
339            .store(now_ms(), Ordering::Relaxed);
340    }
341}
342
343struct DecisionInFlight {
344    metrics: Arc<FlowDecisionMetrics>,
345    started: Instant,
346    completed: bool,
347}
348
349impl DecisionInFlight {
350    fn new(metrics: Arc<FlowDecisionMetrics>) -> Self {
351        metrics.attempted.fetch_add(1, Ordering::Relaxed);
352        metrics.in_flight.fetch_add(1, Ordering::Relaxed);
353        Self {
354            metrics,
355            started: Instant::now(),
356            completed: false,
357        }
358    }
359
360    fn finish(&mut self) {
361        self.record_elapsed();
362        self.metrics.in_flight.fetch_sub(1, Ordering::Relaxed);
363        self.completed = true;
364    }
365
366    fn record_elapsed(&self) {
367        let elapsed = self.started.elapsed().as_micros().min(u128::from(u64::MAX)) as u64;
368        self.metrics
369            .total_dispatch_micros
370            .fetch_add(elapsed, Ordering::Relaxed);
371        self.metrics
372            .max_dispatch_micros
373            .fetch_max(elapsed, Ordering::Relaxed);
374    }
375}
376
377impl Drop for DecisionInFlight {
378    fn drop(&mut self) {
379        if self.completed {
380            return;
381        }
382        self.record_elapsed();
383        self.metrics.in_flight.fetch_sub(1, Ordering::Relaxed);
384        self.metrics.failures.fetch_add(1, Ordering::Relaxed);
385        self.metrics.cancellations.fetch_add(1, Ordering::Relaxed);
386        self.metrics.degraded.store(true, Ordering::Relaxed);
387        self.metrics
388            .last_failure_at_ms
389            .store(now_ms(), Ordering::Relaxed);
390    }
391}
392
393#[derive(Debug, Error, PartialEq, Eq)]
394pub enum FlowDecisionDispatchError {
395    #[error("graph branch is not authorized for production Flow decisions: expected `{expected}`, got `{actual}`")]
396    UnauthorizedBranch { expected: String, actual: String },
397    #[error("decision_id and causation_event_id must be non-empty")]
398    InvalidIdentity,
399    #[error("Flow decision `{0}` reuses an existing decision id with different content")]
400    DecisionIdConflict(String),
401    #[error("Flow decision is owned by another dispatcher until {lease_expires_at_ms}")]
402    Busy { lease_expires_at_ms: u64 },
403    #[error("Flow decision `{0}` lost its lease while the sink was running")]
404    LeaseLost(String),
405    #[error("Flow decision ledger failed: {0}")]
406    Ledger(String),
407    #[error("Flow decision sink failed: {0}")]
408    Sink(String),
409}
410
411fn request_hash(request: &FlowDecisionRequest) -> Result<String, FlowDecisionDispatchError> {
412    serde_json::to_vec(request)
413        .map(sha256::digest)
414        .map_err(|error| FlowDecisionDispatchError::Ledger(error.to_string()))
415}
416
417fn now_ms() -> u64 {
418    use std::time::{SystemTime, UNIX_EPOCH};
419    SystemTime::now()
420        .duration_since(UNIX_EPOCH)
421        .unwrap_or_default()
422        .as_millis()
423        .min(u128::from(u64::MAX)) as u64
424}
425
426fn nonzero(value: u64) -> Option<u64> {
427    (value != 0).then_some(value)
428}
429
430#[cfg(test)]
431mod tests {
432    use super::*;
433    use crate::{FileFlowDecisionLedger, FlowDecisionLedger};
434    use std::sync::atomic::{AtomicUsize, Ordering};
435
436    #[derive(Default)]
437    struct RecordingSink(Mutex<Vec<String>>);
438
439    #[async_trait]
440    impl FlowDecisionSink for RecordingSink {
441        async fn submit(
442            &self,
443            request: &FlowDecisionRequest,
444        ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
445            self.0.lock().await.push(request.decision_id.clone());
446            Ok(())
447        }
448    }
449
450    fn request(branch: &str) -> FlowDecisionRequest {
451        FlowDecisionRequest {
452            decision_id: "decision-1".into(),
453            run_id: "run-1".into(),
454            authority_branch_id: branch.into(),
455            causation_event_id: "event-1".into(),
456            decision: FlowDecision::Complete {
457                output: Value::Null,
458            },
459        }
460    }
461
462    #[tokio::test]
463    async fn submits_once_and_rejects_fork_branches() {
464        let sink = Arc::new(RecordingSink::default());
465        let dispatcher = FlowDecisionDispatcher::new("production", sink.clone());
466        assert!(dispatcher.dispatch(&request("production")).await.unwrap());
467        assert!(!dispatcher.dispatch(&request("production")).await.unwrap());
468        assert_eq!(sink.0.lock().await.as_slice(), ["decision-1"]);
469        assert!(matches!(
470            dispatcher.dispatch(&request("fork")).await,
471            Err(FlowDecisionDispatchError::UnauthorizedBranch { .. })
472        ));
473        let health = dispatcher.health();
474        assert_eq!(health.status, FlowDecisionHealthStatus::Healthy);
475        assert_eq!(health.attempted, 3);
476        assert_eq!(health.claimed, 1);
477        assert_eq!(health.completed, 1);
478        assert_eq!(health.duplicates, 1);
479        assert_eq!(health.rejected, 1);
480        assert_eq!(health.in_flight, 0);
481        assert!(health.last_success_at_ms.is_some());
482    }
483
484    #[tokio::test]
485    async fn completed_receipt_survives_dispatcher_restart() {
486        let sink = Arc::new(RecordingSink::default());
487        let ledger = Arc::new(MemoryFlowDecisionLedger::new());
488        let first = FlowDecisionDispatcher::with_ledger("production", sink.clone(), ledger.clone());
489        assert!(first.dispatch(&request("production")).await.unwrap());
490        let restarted = FlowDecisionDispatcher::with_ledger("production", sink.clone(), ledger);
491        assert!(!restarted.dispatch(&request("production")).await.unwrap());
492        assert_eq!(sink.0.lock().await.len(), 1);
493    }
494
495    #[tokio::test]
496    async fn decision_id_cannot_be_reused_with_different_content() {
497        let sink = Arc::new(RecordingSink::default());
498        let ledger = Arc::new(MemoryFlowDecisionLedger::new());
499        let dispatcher = FlowDecisionDispatcher::with_ledger("production", sink, ledger);
500        dispatcher.dispatch(&request("production")).await.unwrap();
501        let mut conflicting = request("production");
502        conflicting.decision = FlowDecision::Fail {
503            error: "different".into(),
504        };
505        assert!(matches!(
506            dispatcher.dispatch(&conflicting).await,
507            Err(FlowDecisionDispatchError::DecisionIdConflict(_))
508        ));
509        assert_eq!(dispatcher.health().conflicts, 1);
510    }
511
512    struct FailingOnceSink(AtomicUsize);
513
514    #[async_trait]
515    impl FlowDecisionSink for FailingOnceSink {
516        async fn submit(
517            &self,
518            _request: &FlowDecisionRequest,
519        ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
520            if self.0.fetch_add(1, Ordering::SeqCst) == 0 {
521                return Err(std::io::Error::other("transient").into());
522            }
523            Ok(())
524        }
525    }
526
527    #[tokio::test]
528    async fn sink_failure_releases_claim_for_retry() {
529        let sink = Arc::new(FailingOnceSink(AtomicUsize::new(0)));
530        let dispatcher = FlowDecisionDispatcher::with_ledger(
531            "production",
532            sink.clone(),
533            Arc::new(MemoryFlowDecisionLedger::new()),
534        );
535        assert!(matches!(
536            dispatcher.dispatch(&request("production")).await,
537            Err(FlowDecisionDispatchError::Sink(_))
538        ));
539        assert!(dispatcher.dispatch(&request("production")).await.unwrap());
540        assert_eq!(sink.0.load(Ordering::SeqCst), 2);
541        let health = dispatcher.health();
542        assert_eq!(health.status, FlowDecisionHealthStatus::Healthy);
543        assert_eq!(health.sink_failures, 1);
544        assert_eq!(health.failures, 1);
545        assert_eq!(health.completed, 1);
546    }
547
548    #[tokio::test]
549    async fn independent_file_ledgers_serialize_claim_and_allow_expired_takeover() {
550        let directory = tempfile::tempdir().unwrap();
551        let left = FileFlowDecisionLedger::new(directory.path());
552        let right = FileFlowDecisionLedger::new(directory.path());
553        let (left_claim, right_claim) = tokio::join!(
554            left.claim("decision", "hash", "left", 100, 50),
555            right.claim("decision", "hash", "right", 100, 50),
556        );
557        let claims = [left_claim.unwrap(), right_claim.unwrap()];
558        assert_eq!(
559            claims
560                .iter()
561                .filter(|claim| matches!(claim, FlowDecisionClaimOutcome::Claimed { .. }))
562                .count(),
563            1
564        );
565        assert_eq!(
566            claims
567                .iter()
568                .filter(|claim| matches!(claim, FlowDecisionClaimOutcome::Busy { .. }))
569                .count(),
570            1
571        );
572        assert_eq!(
573            right
574                .claim("decision", "hash", "takeover", 151, 50)
575                .await
576                .unwrap(),
577            FlowDecisionClaimOutcome::Claimed { attempt: 2 }
578        );
579        right
580            .complete("decision", "hash", "takeover", 160)
581            .await
582            .unwrap();
583        assert_eq!(right.prune_completed(161).await.unwrap(), 1);
584        assert_eq!(
585            left.claim("decision", "hash", "new", 162, 50)
586                .await
587                .unwrap(),
588            FlowDecisionClaimOutcome::Claimed { attempt: 1 }
589        );
590    }
591
592    #[tokio::test]
593    async fn lease_renewal_requires_current_owner_and_preserves_attempt() {
594        let ledger = MemoryFlowDecisionLedger::new();
595        assert_eq!(
596            ledger
597                .claim("decision", "hash", "owner", 100, 30)
598                .await
599                .unwrap(),
600            FlowDecisionClaimOutcome::Claimed { attempt: 1 }
601        );
602        assert!(!ledger
603            .renew("decision", "hash", "other", 120, 30)
604            .await
605            .unwrap());
606        assert!(ledger
607            .renew("decision", "hash", "owner", 120, 30)
608            .await
609            .unwrap());
610        assert_eq!(
611            ledger
612                .claim("decision", "hash", "other", 145, 30)
613                .await
614                .unwrap(),
615            FlowDecisionClaimOutcome::Busy {
616                lease_expires_at_ms: 150
617            }
618        );
619        assert_eq!(
620            ledger
621                .claim("decision", "hash", "other", 151, 30)
622                .await
623                .unwrap(),
624            FlowDecisionClaimOutcome::Claimed { attempt: 2 }
625        );
626        assert!(!ledger
627            .renew("decision", "hash", "owner", 152, 30)
628            .await
629            .unwrap());
630    }
631
632    #[tokio::test]
633    async fn file_ledger_persists_renewed_lease_across_instances() {
634        let directory = tempfile::tempdir().unwrap();
635        let owner = FileFlowDecisionLedger::new(directory.path());
636        let competitor = FileFlowDecisionLedger::new(directory.path());
637        assert_eq!(
638            owner
639                .claim("decision", "hash", "owner", 100, 20)
640                .await
641                .unwrap(),
642            FlowDecisionClaimOutcome::Claimed { attempt: 1 }
643        );
644        assert!(owner
645            .renew("decision", "hash", "owner", 115, 30)
646            .await
647            .unwrap());
648        assert_eq!(
649            competitor
650                .claim("decision", "hash", "competitor", 125, 20)
651                .await
652                .unwrap(),
653            FlowDecisionClaimOutcome::Busy {
654                lease_expires_at_ms: 145
655            }
656        );
657    }
658
659    struct SlowCountingSink(AtomicUsize);
660
661    #[async_trait]
662    impl FlowDecisionSink for SlowCountingSink {
663        async fn submit(
664            &self,
665            _request: &FlowDecisionRequest,
666        ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
667            self.0.fetch_add(1, Ordering::SeqCst);
668            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
669            Ok(())
670        }
671    }
672
673    #[tokio::test]
674    async fn competing_file_backed_dispatchers_submit_to_sink_once() {
675        let directory = tempfile::tempdir().unwrap();
676        let sink = Arc::new(SlowCountingSink(AtomicUsize::new(0)));
677        let left = FlowDecisionDispatcher::with_ledger(
678            "production",
679            sink.clone(),
680            Arc::new(FileFlowDecisionLedger::new(directory.path())),
681        );
682        let right = FlowDecisionDispatcher::with_ledger(
683            "production",
684            sink.clone(),
685            Arc::new(FileFlowDecisionLedger::new(directory.path())),
686        );
687        let request = request("production");
688        let (left_result, right_result) =
689            tokio::join!(left.dispatch(&request), right.dispatch(&request));
690        let results = [left_result, right_result];
691        assert_eq!(
692            results
693                .iter()
694                .filter(|result| matches!(result, Ok(true)))
695                .count(),
696            1
697        );
698        assert_eq!(
699            results
700                .iter()
701                .filter(|result| matches!(result, Err(FlowDecisionDispatchError::Busy { .. })))
702                .count(),
703            1
704        );
705        assert_eq!(sink.0.load(Ordering::SeqCst), 1);
706    }
707
708    struct RenewalNotifyingLedger {
709        inner: MemoryFlowDecisionLedger,
710        renewed: Arc<tokio::sync::Notify>,
711    }
712
713    #[async_trait]
714    impl FlowDecisionLedger for RenewalNotifyingLedger {
715        async fn claim(
716            &self,
717            decision_id: &str,
718            request_hash: &str,
719            owner_id: &str,
720            now_ms: u64,
721            lease_ms: u64,
722        ) -> anyhow::Result<FlowDecisionClaimOutcome> {
723            self.inner
724                .claim(decision_id, request_hash, owner_id, now_ms, lease_ms)
725                .await
726        }
727
728        async fn renew(
729            &self,
730            decision_id: &str,
731            request_hash: &str,
732            owner_id: &str,
733            now_ms: u64,
734            lease_ms: u64,
735        ) -> anyhow::Result<bool> {
736            let renewed = self
737                .inner
738                .renew(decision_id, request_hash, owner_id, now_ms, lease_ms)
739                .await?;
740            if renewed {
741                self.renewed.notify_one();
742            }
743            Ok(renewed)
744        }
745
746        async fn complete(
747            &self,
748            decision_id: &str,
749            request_hash: &str,
750            owner_id: &str,
751            completed_at_ms: u64,
752        ) -> anyhow::Result<()> {
753            self.inner
754                .complete(decision_id, request_hash, owner_id, completed_at_ms)
755                .await
756        }
757
758        async fn release(
759            &self,
760            decision_id: &str,
761            request_hash: &str,
762            owner_id: &str,
763        ) -> anyhow::Result<()> {
764            self.inner
765                .release(decision_id, request_hash, owner_id)
766                .await
767        }
768    }
769
770    struct ControlledSink {
771        calls: AtomicUsize,
772        finish: Arc<tokio::sync::Notify>,
773    }
774
775    #[async_trait]
776    impl FlowDecisionSink for ControlledSink {
777        async fn submit(
778            &self,
779            _request: &FlowDecisionRequest,
780        ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
781            self.calls.fetch_add(1, Ordering::SeqCst);
782            self.finish.notified().await;
783            Ok(())
784        }
785    }
786
787    #[tokio::test]
788    async fn heartbeat_prevents_takeover_while_sink_is_running() {
789        let renewed = Arc::new(tokio::sync::Notify::new());
790        let finish = Arc::new(tokio::sync::Notify::new());
791        let sink = Arc::new(ControlledSink {
792            calls: AtomicUsize::new(0),
793            finish: finish.clone(),
794        });
795        let ledger = Arc::new(RenewalNotifyingLedger {
796            inner: MemoryFlowDecisionLedger::new(),
797            renewed: renewed.clone(),
798        });
799        let left = Arc::new(
800            FlowDecisionDispatcher::with_ledger("production", sink.clone(), ledger.clone())
801                .with_lease_ms(30),
802        );
803        let right = FlowDecisionDispatcher::with_ledger("production", sink.clone(), ledger)
804            .with_lease_ms(30);
805        let request = request("production");
806        let left_dispatcher = left.clone();
807        let left_request = request.clone();
808        let left_task = tokio::spawn(async move { left_dispatcher.dispatch(&left_request).await });
809        tokio::time::timeout(std::time::Duration::from_secs(2), renewed.notified())
810            .await
811            .expect("dispatcher must renew its live claim");
812        let right_result = right.dispatch(&request).await;
813        assert!(matches!(
814            right_result,
815            Err(FlowDecisionDispatchError::Busy { .. })
816        ));
817        finish.notify_one();
818        assert!(left_task.await.unwrap().unwrap());
819        assert_eq!(sink.calls.load(Ordering::SeqCst), 1);
820        let left_health = left.health();
821        assert_eq!(left_health.completed, 1);
822        assert!(left_health.lease_renewals > 0);
823        assert_eq!(right.health().busy, 1);
824    }
825
826    #[tokio::test]
827    async fn dispatcher_health_counts_expired_claim_takeover() {
828        let ledger = Arc::new(MemoryFlowDecisionLedger::new());
829        let request = request("production");
830        let hash = request_hash(&request).unwrap();
831        ledger
832            .claim(
833                &request.decision_id,
834                &hash,
835                "expired-owner",
836                now_ms().saturating_sub(10),
837                1,
838            )
839            .await
840            .unwrap();
841        let dispatcher = FlowDecisionDispatcher::with_ledger(
842            "production",
843            Arc::new(RecordingSink::default()),
844            ledger,
845        );
846        assert!(dispatcher.dispatch(&request).await.unwrap());
847        let health = dispatcher.health();
848        assert_eq!(health.claimed, 1);
849        assert_eq!(health.takeovers, 1);
850        assert_eq!(health.completed, 1);
851    }
852
853    #[derive(Default)]
854    struct LeaseLosingLedger;
855
856    #[async_trait]
857    impl FlowDecisionLedger for LeaseLosingLedger {
858        async fn claim(
859            &self,
860            _decision_id: &str,
861            _request_hash: &str,
862            _owner_id: &str,
863            _now_ms: u64,
864            _lease_ms: u64,
865        ) -> anyhow::Result<FlowDecisionClaimOutcome> {
866            Ok(FlowDecisionClaimOutcome::Claimed { attempt: 1 })
867        }
868
869        async fn renew(
870            &self,
871            _decision_id: &str,
872            _request_hash: &str,
873            _owner_id: &str,
874            _now_ms: u64,
875            _lease_ms: u64,
876        ) -> anyhow::Result<bool> {
877            Ok(false)
878        }
879
880        async fn complete(
881            &self,
882            _decision_id: &str,
883            _request_hash: &str,
884            _owner_id: &str,
885            _completed_at_ms: u64,
886        ) -> anyhow::Result<()> {
887            Ok(())
888        }
889
890        async fn release(
891            &self,
892            _decision_id: &str,
893            _request_hash: &str,
894            _owner_id: &str,
895        ) -> anyhow::Result<()> {
896            Ok(())
897        }
898    }
899
900    #[derive(Default)]
901    struct ClaimFailingLedger;
902
903    #[async_trait]
904    impl FlowDecisionLedger for ClaimFailingLedger {
905        async fn claim(
906            &self,
907            _decision_id: &str,
908            _request_hash: &str,
909            _owner_id: &str,
910            _now_ms: u64,
911            _lease_ms: u64,
912        ) -> anyhow::Result<FlowDecisionClaimOutcome> {
913            anyhow::bail!("ledger unavailable")
914        }
915
916        async fn renew(
917            &self,
918            _decision_id: &str,
919            _request_hash: &str,
920            _owner_id: &str,
921            _now_ms: u64,
922            _lease_ms: u64,
923        ) -> anyhow::Result<bool> {
924            Ok(false)
925        }
926
927        async fn complete(
928            &self,
929            _decision_id: &str,
930            _request_hash: &str,
931            _owner_id: &str,
932            _completed_at_ms: u64,
933        ) -> anyhow::Result<()> {
934            Ok(())
935        }
936
937        async fn release(
938            &self,
939            _decision_id: &str,
940            _request_hash: &str,
941            _owner_id: &str,
942        ) -> anyhow::Result<()> {
943            Ok(())
944        }
945    }
946
947    #[tokio::test]
948    async fn dispatcher_health_records_ledger_failure() {
949        let dispatcher = FlowDecisionDispatcher::with_ledger(
950            "production",
951            Arc::new(RecordingSink::default()),
952            Arc::new(ClaimFailingLedger),
953        );
954        assert!(matches!(
955            dispatcher.dispatch(&request("production")).await,
956            Err(FlowDecisionDispatchError::Ledger(_))
957        ));
958        let health = dispatcher.health();
959        assert_eq!(health.status, FlowDecisionHealthStatus::Degraded);
960        assert_eq!(health.ledger_failures, 1);
961        assert_eq!(health.in_flight, 0);
962    }
963
964    struct CancellableSink {
965        started: AtomicUsize,
966        completed: AtomicUsize,
967    }
968
969    #[async_trait]
970    impl FlowDecisionSink for CancellableSink {
971        async fn submit(
972            &self,
973            _request: &FlowDecisionRequest,
974        ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
975            self.started.fetch_add(1, Ordering::SeqCst);
976            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
977            self.completed.fetch_add(1, Ordering::SeqCst);
978            Ok(())
979        }
980    }
981
982    #[tokio::test]
983    async fn lost_lease_cancels_in_flight_sink_future() {
984        let sink = Arc::new(CancellableSink {
985            started: AtomicUsize::new(0),
986            completed: AtomicUsize::new(0),
987        });
988        let dispatcher = FlowDecisionDispatcher::with_ledger(
989            "production",
990            sink.clone(),
991            Arc::new(LeaseLosingLedger),
992        )
993        .with_lease_ms(9);
994        assert!(matches!(
995            dispatcher.dispatch(&request("production")).await,
996            Err(FlowDecisionDispatchError::LeaseLost(_))
997        ));
998        tokio::time::sleep(std::time::Duration::from_millis(60)).await;
999        assert_eq!(sink.started.load(Ordering::SeqCst), 1);
1000        assert_eq!(sink.completed.load(Ordering::SeqCst), 0);
1001        let health = dispatcher.health();
1002        assert_eq!(health.status, FlowDecisionHealthStatus::Degraded);
1003        assert_eq!(health.lease_lost, 1);
1004        assert_eq!(health.in_flight, 0);
1005        assert!(health.last_failure_at_ms.is_some());
1006    }
1007
1008    struct BlockingSink(Arc<tokio::sync::Notify>);
1009
1010    #[async_trait]
1011    impl FlowDecisionSink for BlockingSink {
1012        async fn submit(
1013            &self,
1014            _request: &FlowDecisionRequest,
1015        ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
1016            self.0.notify_one();
1017            std::future::pending().await
1018        }
1019    }
1020
1021    #[tokio::test]
1022    async fn cancelled_dispatch_does_not_leak_in_flight_health() {
1023        let started = Arc::new(tokio::sync::Notify::new());
1024        let dispatcher = Arc::new(FlowDecisionDispatcher::new(
1025            "production",
1026            Arc::new(BlockingSink(started.clone())),
1027        ));
1028        let task_dispatcher = dispatcher.clone();
1029        let pending =
1030            tokio::spawn(async move { task_dispatcher.dispatch(&request("production")).await });
1031        started.notified().await;
1032        assert_eq!(dispatcher.health().in_flight, 1);
1033        pending.abort();
1034        assert!(pending.await.unwrap_err().is_cancelled());
1035        let health = dispatcher.health();
1036        assert_eq!(health.status, FlowDecisionHealthStatus::Degraded);
1037        assert_eq!(health.attempted, 1);
1038        assert_eq!(health.cancellations, 1);
1039        assert_eq!(health.in_flight, 0);
1040    }
1041}