1. 1//! Cross-run circuit breaker for provider health. Only retryable failures
2. 2//! (429/529/network) trip it; auth/config errors are the caller's problem.
3. 3//! While open, steps fail fast with the remaining cooldown instead of
4. 4//! burning their retry budget against a dead provider.
5. 5
6. 6use std::collections::HashMap;
7. 7use std::sync::Mutex;
8. 8use std::time::{Duration, Instant};
9. 9
10. 10#[derive(Debug, Clone)]
11. 11pub struct CircuitBreakerConfig {
12. 12 /// Consecutive retryable failures before the circuit opens.
13. 13 pub threshold: u32,
14. 14 /// How long an open circuit stays open before allowing a probe.
15. 15 pub cooldown: Duration,
16. 16}
17. 17
18. 18impl Default for CircuitBreakerConfig {
19. 19 fn default() -> Self {
20. 20 CircuitBreakerConfig {
21. 21 threshold: 5,
22. 22 cooldown: Duration::from_secs(60),
23. 23 }
24. 24 }
25. 25}
26. 26
27. 27#[derive(Debug, Default)]
28. 28struct State {
29. 29 consecutive_failures: u32,
30. 30 opened_at: Option<Instant>,
31. 31 probe_in_flight: bool,
32. 32}
33. 33
34. 34#[derive(Debug)]
35. 35pub struct CircuitBreaker {
36. 36 config: CircuitBreakerConfig,
37. 37 state: Mutex<HashMap<String, State>>,
38. 38}
39. 39
40. 40#[derive(Debug, thiserror::Error)]
41. 41#[error(
42. 42 "circuit open for provider (cooling down {remaining_secs}s after {failures} consecutive failures)"
43. 43)]
44. 44pub struct CircuitOpen {
45. 45 pub remaining_secs: u64,
46. 46 pub failures: u32,
47. 47}
48. 48
49. 49impl CircuitBreaker {
50. 50 pub fn new(config: CircuitBreakerConfig) -> Self {
51. 51 CircuitBreaker {
52. 52 config,
53. 53 state: Mutex::new(HashMap::new()),
54. 54 }
55. 55 }
56. 56
57. 57 /// Err when the circuit is open and the cooldown has not elapsed.
58. 58 pub fn check(&self) -> Result<(), CircuitOpen> {
59. 59 self.check_key("")
60. 60 }
61. 61
62. 62 /// Check one provider/key circuit. Provider names are used by the agent
63. 63 /// as the stable health domain; an empty key means the process-global
64. 64 /// helper semantics for callers that do not have a route identity.
65. 65 pub fn check_key(&self, key: &str) -> Result<(), CircuitOpen> {
66. 66 let mut states = self.lock();
67. 67 let st = states.entry(key.to_string()).or_default();
68. 68 if st.probe_in_flight {
69. 69 return Err(CircuitOpen {
70. 70 remaining_secs: 1,
71. 71 failures: st.consecutive_failures,
72. 72 });
73. 73 }
74. 74 if let Some(opened_at) = st.opened_at {
75. 75 let elapsed = opened_at.elapsed();
76. 76 if elapsed < self.config.cooldown {
77. 77 return Err(CircuitOpen {
78. 78 remaining_secs: (self.config.cooldown - elapsed).as_secs().max(1),
79. 79 failures: st.consecutive_failures,
80. 80 });
81. 81 }
82. 82 // Cooldown elapsed: reserve exactly one half-open probe. Other
83. 83 // callers remain fail-closed until that probe settles.
84. 84 st.opened_at = None;
85. 85 st.probe_in_flight = true;
86. 86 }
87. 87 Ok(())
88. 88 }
89. 89
90. 90 pub fn record_success(&self) {
91. 91 self.record_success_key("");
92. 92 }
93. 93
94. 94 pub fn record_success_key(&self, key: &str) {
95. 95 let mut states = self.lock();
96. 96 let st = states.entry(key.to_string()).or_default();
97. 97 st.consecutive_failures = 0;
98. 98 st.opened_at = None;
99. 99 st.probe_in_flight = false;
100. 100 }
101. 101
102. 102 /// Only retryable failures should call this.
103. 103 pub fn record_failure(&self) {
104. 104 self.record_failure_key("");
105. 105 }
106. 106
107. 107 pub fn record_failure_key(&self, key: &str) {
108. 108 let mut states = self.lock();
109. 109 let st = states.entry(key.to_string()).or_default();
110. 110 st.consecutive_failures = st.consecutive_failures.saturating_add(1);
111. 111 if st.consecutive_failures >= self.config.threshold && st.opened_at.is_none() {
112. 112 st.opened_at = Some(Instant::now());
113. 113 }
114. 114 st.probe_in_flight = false;
115. 115 }
116. 116
117. 117 fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<String, State>> {
118. 118 self.state
119. 119 .lock()
120. 120 .unwrap_or_else(std::sync::PoisonError::into_inner)
121. 121 }
122. 122}
123. 123