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