systemprompt_database/resilience/
breaker.rs

1//! A circuit breaker that fast-fails calls to an unhealthy dependency.
2//!
3//! Admission is a [`Probe`] token: while it lives it occupies one of the
4//! half-open probe slots, and it must be settled with `success`/`failure`.
5//! A probe dropped unsettled — a cancelled future — frees its slot without
6//! changing the breaker's mode, so a client disconnect can never exhaust the
7//! probe budget and leave the breaker open forever.
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12use std::sync::{Mutex, MutexGuard, PoisonError};
13use std::time::Instant;
14
15use super::config::BreakerConfig;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18enum Mode {
19    Closed,
20    Open,
21    HalfOpen,
22}
23
24#[derive(Debug)]
25struct State {
26    mode: Mode,
27    consecutive_failures: u32,
28    open_until: Option<Instant>,
29    probes_in_flight: u32,
30}
31
32#[derive(Debug, Clone, Copy)]
33pub struct Tripped;
34
35/// One admitted call. Settle it with [`Probe::success`] or
36/// [`Probe::failure`]; dropping it unsettled releases the probe slot only.
37#[derive(Debug)]
38#[must_use = "an unsettled probe neither closes nor reopens the breaker"]
39pub struct Probe<'a> {
40    breaker: &'a CircuitBreaker,
41    counted: bool,
42    settled: bool,
43}
44
45impl Probe<'_> {
46    pub fn success(mut self) {
47        self.settled = true;
48        self.breaker.settle(self.counted, true);
49    }
50
51    pub fn failure(mut self) {
52        self.settled = true;
53        self.breaker.settle(self.counted, false);
54    }
55}
56
57impl Drop for Probe<'_> {
58    fn drop(&mut self) {
59        if !self.settled && self.counted {
60            let mut state = self.breaker.lock();
61            state.probes_in_flight = state.probes_in_flight.saturating_sub(1);
62        }
63    }
64}
65
66#[derive(Debug)]
67pub struct CircuitBreaker {
68    key: String,
69    cfg: BreakerConfig,
70    state: Mutex<State>,
71}
72
73impl CircuitBreaker {
74    pub fn new(key: impl Into<String>, cfg: BreakerConfig) -> Self {
75        Self {
76            key: key.into(),
77            cfg,
78            state: Mutex::new(State {
79                mode: Mode::Closed,
80                consecutive_failures: 0,
81                open_until: None,
82                probes_in_flight: 0,
83            }),
84        }
85    }
86
87    pub fn acquire(&self) -> Result<Probe<'_>, Tripped> {
88        let mut state = self.lock();
89        let counted = match state.mode {
90            Mode::Closed => false,
91            Mode::Open => {
92                let cooled_down = state
93                    .open_until
94                    .is_some_and(|until| Instant::now() >= until);
95                if !cooled_down {
96                    return Err(Tripped);
97                }
98                self.transition(&mut state, Mode::HalfOpen);
99                state.probes_in_flight = 1;
100                true
101            },
102            Mode::HalfOpen => {
103                if state.probes_in_flight >= self.cfg.half_open_max_probes {
104                    return Err(Tripped);
105                }
106                state.probes_in_flight += 1;
107                true
108            },
109        };
110        drop(state);
111        Ok(Probe {
112            breaker: self,
113            counted,
114            settled: false,
115        })
116    }
117
118    pub fn record_success(&self) {
119        self.settle(false, true);
120    }
121
122    pub fn record_failure(&self) {
123        self.settle(false, false);
124    }
125
126    fn settle(&self, counted: bool, success: bool) {
127        let mut state = self.lock();
128        if counted {
129            state.probes_in_flight = state.probes_in_flight.saturating_sub(1);
130        }
131        if success {
132            state.consecutive_failures = 0;
133            if state.mode != Mode::Closed {
134                self.transition(&mut state, Mode::Closed);
135                state.open_until = None;
136            }
137            return;
138        }
139        state.consecutive_failures = state.consecutive_failures.saturating_add(1);
140
141        let should_open = state.mode == Mode::HalfOpen
142            || state.consecutive_failures >= self.cfg.failure_threshold;
143        if should_open && state.mode != Mode::Open {
144            self.transition(&mut state, Mode::Open);
145            state.open_until = Some(Instant::now() + self.cfg.open_cooldown);
146        }
147    }
148
149    #[must_use]
150    pub fn is_open(&self) -> bool {
151        self.lock().mode == Mode::Open
152    }
153
154    fn transition(&self, state: &mut State, to: Mode) {
155        let from = state.mode;
156        if from != to {
157            state.mode = to;
158            tracing::warn!(key = %self.key, ?from, ?to, "circuit breaker state transition");
159        }
160    }
161
162    fn lock(&self) -> MutexGuard<'_, State> {
163        self.state.lock().unwrap_or_else(PoisonError::into_inner)
164    }
165}
systemprompt_database/resilience/breaker.rs

systemprompt_database/resilience/
breaker.rs
