af_workflow/
provider.rs

1//! Provider-local backpressure and circuit breaking.
2
3use std::collections::HashMap;
4use std::sync::{Arc, Mutex};
5use std::time::{Duration, Instant};
6
7use tokio::sync::{OwnedSemaphorePermit, Semaphore};
8
9/// Dispatch priority; higher priorities keep reserved capacity.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
11pub enum WorkPriority {
12    /// Lowest.
13    Notification,
14    /// Ordinary workflow work.
15    Normal,
16    /// Funds and control actions.
17    Control,
18    /// Emergency stop and revocation.
19    Emergency,
20}
21
22/// Bulkhead, circuit-breaker and backoff limits of a provider.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct ProviderLimits {
25    /// Total concurrent calls.
26    pub concurrency: usize,
27    /// Slots kept for control-priority work.
28    pub reserved_control_concurrency: usize,
29    /// Concurrent calls per tenant.
30    pub tenant_concurrency: usize,
31    /// Consecutive failures that open the circuit.
32    pub failure_threshold: u32,
33    /// Open duration before a probe.
34    pub recovery_after: Duration,
35    /// Per-call deadline.
36    pub max_elapsed: Duration,
37    /// First backoff.
38    pub initial_backoff: Duration,
39    /// Backoff cap.
40    pub max_backoff: Duration,
41}
42
43impl ProviderLimits {
44    /// Reject zero or inconsistent limits.
45    pub fn validate(&self) -> Result<(), String> {
46        if self.concurrency == 0
47            || self.tenant_concurrency == 0
48            || self.failure_threshold == 0
49            || self.recovery_after.is_zero()
50            || self.max_elapsed.is_zero()
51        {
52            return Err("provider limits must be positive".into());
53        }
54        if self.tenant_concurrency > self.concurrency {
55            return Err("tenant concurrency cannot exceed provider concurrency".into());
56        }
57        if self.reserved_control_concurrency >= self.concurrency {
58            return Err("reserved control concurrency must be below provider concurrency".into());
59        }
60        if self.initial_backoff.is_zero() || self.initial_backoff > self.max_backoff {
61            return Err("provider backoff range is invalid".into());
62        }
63        Ok(())
64    }
65
66    /// Exponential backoff with deterministic jitter for `attempt`.
67    pub fn backoff(&self, attempt: u32, jitter_seed: u64) -> Duration {
68        let factor = 1_u32.checked_shl(attempt.min(20)).unwrap_or(u32::MAX);
69        let base = self
70            .initial_backoff
71            .saturating_mul(factor)
72            .min(self.max_backoff);
73        let jitter_ceiling = (base.as_millis() / 4).max(1) as u64;
74        base.saturating_add(Duration::from_millis(jitter_seed % jitter_ceiling))
75            .min(self.max_backoff)
76    }
77}
78
79/// Why a permit was refused.
80#[derive(Debug, Clone, PartialEq, Eq)]
81pub enum ProviderBlock {
82    /// The circuit is open.
83    CircuitOpen {
84        /// Time until a probe is allowed.
85        retry_after: Duration,
86    },
87    /// Every allowed slot is busy.
88    Saturated,
89}
90
91struct Circuit {
92    consecutive_failures: u32,
93    opened_at: Option<Instant>,
94    probe_in_flight: bool,
95}
96
97/// ponytail: one semaphore per tenant is kept in a map; idle entries are swept
98/// once the map passes this size so a long tail of one-off tenants cannot grow
99/// it without bound. Upgrade path is an LRU if sweeping ever shows up in profiles.
100const TENANT_SWEEP_THRESHOLD: usize = 1_024;
101
102/// Bulkhead plus circuit breaker guarding one provider.
103pub struct ProviderGate {
104    limits: ProviderLimits,
105    global: Arc<Semaphore>,
106    normal: Arc<Semaphore>,
107    tenants: Mutex<HashMap<String, Arc<Semaphore>>>,
108    circuit: Arc<Mutex<Circuit>>,
109}
110
111/// Holding a permit keeps provider capacity reserved. Dropping it without a
112/// `record_*` call releases capacity and clears a recovery probe without
113/// touching the failure counter, so a store error or a panic can never leave
114/// the circuit stuck open.
115pub struct ProviderPermit {
116    _global: OwnedSemaphorePermit,
117    _normal: Option<OwnedSemaphorePermit>,
118    _tenant: Option<OwnedSemaphorePermit>,
119    is_probe: bool,
120    circuit: Arc<Mutex<Circuit>>,
121}
122
123impl Drop for ProviderPermit {
124    fn drop(&mut self) {
125        if self.is_probe {
126            self.circuit
127                .lock()
128                .unwrap_or_else(std::sync::PoisonError::into_inner)
129                .probe_in_flight = false;
130        }
131    }
132}
133
134impl ProviderGate {
135    /// Gate for validated `limits`.
136    pub fn new(limits: ProviderLimits) -> Result<Self, String> {
137        limits.validate()?;
138        Ok(Self {
139            global: Arc::new(Semaphore::new(limits.concurrency)),
140            normal: Arc::new(Semaphore::new(
141                limits.concurrency - limits.reserved_control_concurrency,
142            )),
143            limits,
144            tenants: Mutex::new(HashMap::new()),
145            circuit: Arc::new(Mutex::new(Circuit {
146                consecutive_failures: 0,
147                opened_at: None,
148                probe_in_flight: false,
149            })),
150        })
151    }
152
153    /// Acquire a permit for `tenant` at `priority`.
154    pub fn try_acquire(
155        &self,
156        tenant_id: &str,
157        priority: WorkPriority,
158        now: Instant,
159    ) -> Result<ProviderPermit, ProviderBlock> {
160        let is_probe = {
161            let mut circuit = self
162                .circuit
163                .lock()
164                .unwrap_or_else(std::sync::PoisonError::into_inner);
165            match circuit.opened_at {
166                None => false,
167                Some(opened_at) => {
168                    let elapsed = now.saturating_duration_since(opened_at);
169                    if elapsed < self.limits.recovery_after || circuit.probe_in_flight {
170                        return Err(ProviderBlock::CircuitOpen {
171                            retry_after: self.limits.recovery_after.saturating_sub(elapsed),
172                        });
173                    }
174                    circuit.probe_in_flight = true;
175                    true
176                }
177            }
178        };
179        let normal = if matches!(priority, WorkPriority::Notification | WorkPriority::Normal) {
180            match self.normal.clone().try_acquire_owned() {
181                Ok(permit) => Some(permit),
182                Err(_) => {
183                    self.release_probe(is_probe);
184                    return Err(ProviderBlock::Saturated);
185                }
186            }
187        } else {
188            None
189        };
190        let global = match self.global.clone().try_acquire_owned() {
191            Ok(permit) => permit,
192            Err(_) => {
193                self.release_probe(is_probe);
194                return Err(ProviderBlock::Saturated);
195            }
196        };
197        let tenant = if priority == WorkPriority::Emergency {
198            None
199        } else {
200            let tenant = {
201                let mut tenants = self
202                    .tenants
203                    .lock()
204                    .unwrap_or_else(std::sync::PoisonError::into_inner);
205                if tenants.len() >= TENANT_SWEEP_THRESHOLD {
206                    let full = self.limits.tenant_concurrency;
207                    tenants.retain(|_, semaphore| semaphore.available_permits() < full);
208                }
209                tenants
210                    .entry(tenant_id.to_owned())
211                    .or_insert_with(|| Arc::new(Semaphore::new(self.limits.tenant_concurrency)))
212                    .clone()
213            };
214            match tenant.try_acquire_owned() {
215                Ok(permit) => Some(permit),
216                Err(_) => {
217                    self.release_probe(is_probe);
218                    return Err(ProviderBlock::Saturated);
219                }
220            }
221        };
222        Ok(ProviderPermit {
223            _global: global,
224            _normal: normal,
225            _tenant: tenant,
226            is_probe,
227            circuit: self.circuit.clone(),
228        })
229    }
230
231    /// Return a permit after success; closes a half-open circuit.
232    pub fn record_success(&self, permit: ProviderPermit) {
233        let mut circuit = self
234            .circuit
235            .lock()
236            .unwrap_or_else(std::sync::PoisonError::into_inner);
237        circuit.consecutive_failures = 0;
238        circuit.opened_at = None;
239        drop(circuit);
240        drop(permit);
241    }
242
243    /// Return a permit after failure; may open the circuit.
244    pub fn record_failure(&self, permit: ProviderPermit, now: Instant) {
245        let mut circuit = self
246            .circuit
247            .lock()
248            .unwrap_or_else(std::sync::PoisonError::into_inner);
249        circuit.consecutive_failures = circuit.consecutive_failures.saturating_add(1);
250        if permit.is_probe || circuit.consecutive_failures >= self.limits.failure_threshold {
251            circuit.opened_at = Some(now);
252        }
253        drop(circuit);
254        drop(permit);
255    }
256
257    /// Give capacity back without judging the provider. Use this when the
258    /// caller failed before or after the provider ran (store conflict, stale
259    /// claim) so infrastructure errors do not trip the provider's circuit.
260    pub fn release(&self, permit: ProviderPermit) {
261        drop(permit);
262    }
263
264    fn release_probe(&self, is_probe: bool) {
265        if is_probe {
266            self.circuit
267                .lock()
268                .unwrap_or_else(std::sync::PoisonError::into_inner)
269                .probe_in_flight = false;
270        }
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277
278    fn limits() -> ProviderLimits {
279        ProviderLimits {
280            concurrency: 2,
281            reserved_control_concurrency: 1,
282            tenant_concurrency: 1,
283            failure_threshold: 2,
284            recovery_after: Duration::from_secs(10),
285            max_elapsed: Duration::from_secs(30),
286            initial_backoff: Duration::from_millis(100),
287            max_backoff: Duration::from_secs(1),
288        }
289    }
290
291    #[test]
292    fn tenant_bulkhead_and_circuit_recovery_are_bounded() {
293        let gate = ProviderGate::new(limits()).unwrap();
294        let now = Instant::now();
295        let first = gate.try_acquire("one", WorkPriority::Normal, now).unwrap();
296        let emergency = gate
297            .try_acquire("one", WorkPriority::Emergency, now)
298            .expect("reserved control capacity must remain available");
299        gate.record_success(emergency);
300        assert_eq!(
301            gate.try_acquire("two", WorkPriority::Normal, now).err(),
302            Some(ProviderBlock::Saturated)
303        );
304        gate.record_failure(first, now);
305        let second = gate.try_acquire("two", WorkPriority::Normal, now).unwrap();
306        gate.record_failure(second, now);
307        assert!(matches!(
308            gate.try_acquire("one", WorkPriority::Normal, now),
309            Err(ProviderBlock::CircuitOpen { .. })
310        ));
311        let probe = gate
312            .try_acquire("one", WorkPriority::Normal, now + Duration::from_secs(10))
313            .unwrap();
314        gate.record_success(probe);
315        assert!(gate
316            .try_acquire("one", WorkPriority::Normal, now + Duration::from_secs(10))
317            .is_ok());
318    }
319
320    #[test]
321    fn dropped_probe_permit_does_not_wedge_the_circuit() {
322        let gate = ProviderGate::new(limits()).unwrap();
323        let now = Instant::now();
324        for tenant in ["one", "two"] {
325            let permit = gate.try_acquire(tenant, WorkPriority::Normal, now).unwrap();
326            gate.record_failure(permit, now);
327        }
328        let later = now + Duration::from_secs(10);
329        let probe = gate
330            .try_acquire("one", WorkPriority::Normal, later)
331            .unwrap();
332        assert!(matches!(
333            gate.try_acquire("two", WorkPriority::Normal, later),
334            Err(ProviderBlock::CircuitOpen { .. })
335        ));
336        gate.release(probe);
337        let probe = gate
338            .try_acquire("two", WorkPriority::Normal, later)
339            .expect("released probe must allow the next probe");
340        gate.record_success(probe);
341        assert!(gate.try_acquire("one", WorkPriority::Normal, later).is_ok());
342    }
343
344    #[test]
345    fn tenant_map_is_swept_once_it_grows() {
346        let gate = ProviderGate::new(ProviderLimits {
347            concurrency: 4,
348            tenant_concurrency: 1,
349            ..limits()
350        })
351        .unwrap();
352        let now = Instant::now();
353        for index in 0..TENANT_SWEEP_THRESHOLD + 5 {
354            let permit = gate
355                .try_acquire(&format!("tenant-{index}"), WorkPriority::Normal, now)
356                .unwrap();
357            gate.record_success(permit);
358        }
359        assert!(gate.tenants.lock().unwrap().len() <= TENANT_SWEEP_THRESHOLD + 1);
360    }
361
362    #[test]
363    fn backoff_is_capped() {
364        let limits = limits();
365        assert_eq!(limits.backoff(30, 99), Duration::from_secs(1));
366    }
367}
af_workflow/provider.rs

af_workflow/
provider.rs
