plecto_control/upstream/
circuit_breaker.rs1use std::sync::Arc;
6use std::sync::atomic::Ordering;
7
8use super::UpstreamGroup;
9
10#[derive(Debug)]
14pub struct RequestPermit {
15 group: Option<Arc<UpstreamGroup>>,
17}
18
19impl Drop for RequestPermit {
20 fn drop(&mut self) {
21 if let Some(group) = &self.group {
22 group.in_flight.fetch_sub(1, Ordering::Relaxed);
23 }
24 }
25}
26
27impl UpstreamGroup {
28 pub fn try_acquire(self: &Arc<Self>) -> Option<RequestPermit> {
34 if self.max_requests == 0 {
35 return Some(RequestPermit { group: None });
36 }
37 let prev = self.in_flight.fetch_add(1, Ordering::Relaxed);
41 if prev >= self.max_requests {
42 self.in_flight.fetch_sub(1, Ordering::Relaxed);
43 return None;
44 }
45 Some(RequestPermit {
46 group: Some(self.clone()),
47 })
48 }
49}
50
51#[cfg(test)]
52mod tests {
53 use crate::manifest::{
54 AddressSpec, CircuitBreaker, HealthConfig, LbAlgorithm, OutlierDetection, Upstream,
55 };
56 use crate::upstream::UpstreamRegistry;
57
58 fn health(healthy_threshold: u32, unhealthy_threshold: u32) -> HealthConfig {
59 HealthConfig {
60 path: "/healthz".to_string(),
61 interval_ms: 100,
62 timeout_ms: 50,
63 healthy_threshold,
64 unhealthy_threshold,
65 port: None,
66 }
67 }
68
69 fn upstream(name: &str, addrs: &[&str], h: HealthConfig) -> Upstream {
70 Upstream {
71 name: name.to_string(),
72 addresses: addrs
73 .iter()
74 .map(|s| AddressSpec::Bare(s.to_string()))
75 .collect(),
76 lb_algorithm: LbAlgorithm::RoundRobin,
77 hash: None,
78 tls: None,
79 resolve_interval_ms: 0,
80 health: h,
81 request_timeout_ms: 30_000,
82 max_retries: 1,
83 overall_timeout_ms: 0,
84 circuit_breaker: CircuitBreaker::default(),
85 outlier_detection: OutlierDetection::default(),
86 }
87 }
88
89 #[test]
90 fn circuit_breaker_caps_concurrent_in_flight_and_releases_on_drop() {
91 let reg = UpstreamRegistry::new();
94 reg.reconcile(
95 &[Upstream {
96 name: "u".to_string(),
97 addresses: vec![AddressSpec::Bare("a:1".to_string())],
98 lb_algorithm: LbAlgorithm::RoundRobin,
99 hash: None,
100 tls: None,
101 resolve_interval_ms: 0,
102 health: health(1, 1),
103 request_timeout_ms: 30_000,
104 max_retries: 0,
105 overall_timeout_ms: 0,
106 circuit_breaker: CircuitBreaker { max_requests: 2 },
107 outlier_detection: OutlierDetection::default(),
108 }],
109 std::path::Path::new("."),
110 )
111 .unwrap();
112 let g = reg.group("u").unwrap();
113
114 let p1 = g.try_acquire().expect("1st slot is under the cap");
115 let _p2 = g.try_acquire().expect("2nd slot is under the cap");
116 assert!(
117 g.try_acquire().is_none(),
118 "the 3rd concurrent request is over the cap → rejected"
119 );
120
121 drop(p1);
122 let _p3 = g.try_acquire().expect("a freed slot is reusable");
123 assert!(
124 g.try_acquire().is_none(),
125 "still capped at 2 after reusing the freed slot"
126 );
127 }
128
129 #[test]
130 fn circuit_breaker_zero_is_unlimited() {
131 let reg = UpstreamRegistry::new();
134 reg.reconcile(
135 &[upstream("u", &["a:1"], health(1, 1))],
136 std::path::Path::new("."),
137 )
138 .unwrap();
139 let g = reg.group("u").unwrap();
140 let permits: Vec<_> = (0..1000)
141 .map(|_| g.try_acquire().expect("an unlimited breaker never rejects"))
142 .collect();
143 assert_eq!(permits.len(), 1000);
144 }
145}