github_mcp/core/
circuit_breaker.rs

1// GitHub v3 REST API MCP server — generated by mcpify. Do not hand-edit.
2
3use std::fmt;
4use std::future::Future;
5use std::sync::Mutex;
6use std::time::{Duration, Instant};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9enum State {
10    Closed,
11    Open,
12    HalfOpen,
13}
14
15struct Inner {
16    state: State,
17    failure_count: u32,
18    opened_at: Option<Instant>,
19}
20
21/// `CircuitBreaker::execute`'s error type: either the breaker itself was
22/// open (call never ran) or the wrapped call's own error `E` propagated
23/// through. Generic over `E` rather than folding into `McpifyError` the way
24/// `targets::typescript`'s `CircuitBreaker.execute` folds every rejection
25/// into a re-thrown `unknown`, since Rust's `Result` makes losing the
26/// original error type a real, avoidable cost here.
27#[derive(Debug)]
28pub enum CircuitBreakerError<E> {
29    Open,
30    Inner(E),
31}
32
33impl<E: fmt::Display> fmt::Display for CircuitBreakerError<E> {
34    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
35        match self {
36            Self::Open => write!(f, "circuit breaker is open"),
37            Self::Inner(err) => write!(f, "{err}"),
38        }
39    }
40}
41
42impl<E: fmt::Debug + fmt::Display> std::error::Error for CircuitBreakerError<E> {}
43
44/// CLOSED -> OPEN -> HALF_OPEN state machine around outbound calls to the
45/// target API (REQ-2.3.3). Opens after `failure_threshold` consecutive
46/// failures; after `reset_timeout`, lets one trial call through
47/// (HALF_OPEN) to decide whether to close again or re-open.
48pub struct CircuitBreaker {
49    failure_threshold: u32,
50    reset_timeout: Duration,
51    inner: Mutex<Inner>,
52}
53
54impl CircuitBreaker {
55    pub fn new(failure_threshold: u32, reset_timeout: Duration) -> Self {
56        Self {
57            failure_threshold,
58            reset_timeout,
59            inner: Mutex::new(Inner {
60                state: State::Closed,
61                failure_count: 0,
62                opened_at: None,
63            }),
64        }
65    }
66
67    pub async fn execute<F, Fut, T, E>(&self, f: F) -> Result<T, CircuitBreakerError<E>>
68    where
69        F: FnOnce() -> Fut,
70        Fut: Future<Output = Result<T, E>>,
71    {
72        {
73            let mut inner = self.inner.lock().unwrap();
74            if inner.state == State::Open {
75                let elapsed = inner.opened_at.map(|at| at.elapsed()).unwrap_or_default();
76                if elapsed < self.reset_timeout {
77                    return Err(CircuitBreakerError::Open);
78                }
79                inner.state = State::HalfOpen;
80            }
81        }
82
83        match f().await {
84            Ok(value) => {
85                self.on_success();
86                Ok(value)
87            }
88            Err(err) => {
89                self.on_failure();
90                Err(CircuitBreakerError::Inner(err))
91            }
92        }
93    }
94
95    fn on_success(&self) {
96        let mut inner = self.inner.lock().unwrap();
97        inner.failure_count = 0;
98        inner.state = State::Closed;
99    }
100
101    fn on_failure(&self) {
102        let mut inner = self.inner.lock().unwrap();
103        inner.failure_count += 1;
104        if inner.state == State::HalfOpen || inner.failure_count >= self.failure_threshold {
105            inner.state = State::Open;
106            inner.opened_at = Some(Instant::now());
107        }
108    }
109}
110
111impl Default for CircuitBreaker {
112    fn default() -> Self {
113        Self::new(5, Duration::from_secs(30))
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[tokio::test]
122    async fn opens_after_the_failure_threshold_is_reached() {
123        let breaker = CircuitBreaker::new(2, Duration::from_secs(30));
124
125        for _ in 0..2 {
126            let _ = breaker
127                .execute(|| async { Err::<(), _>(anyhow::anyhow!("boom")) })
128                .await;
129        }
130
131        let result = breaker
132            .execute(|| async { Ok::<_, anyhow::Error>(()) })
133            .await;
134        assert!(matches!(result, Err(CircuitBreakerError::Open)));
135    }
136
137    #[tokio::test]
138    async fn a_success_resets_the_failure_count() {
139        let breaker = CircuitBreaker::new(2, Duration::from_secs(30));
140
141        let _ = breaker
142            .execute(|| async { Err::<(), _>(anyhow::anyhow!("boom")) })
143            .await;
144        breaker
145            .execute(|| async { Ok::<_, anyhow::Error>(()) })
146            .await
147            .unwrap();
148
149        // A single further failure shouldn't open a breaker with
150        // threshold 2, since the prior success reset the counter — the
151        // error should be the wrapped call's own error, not `Open`.
152        let result = breaker
153            .execute(|| async { Err::<(), _>(anyhow::anyhow!("boom")) })
154            .await;
155        assert!(matches!(result, Err(CircuitBreakerError::Inner(_))));
156
157        let ok = breaker
158            .execute(|| async { Ok::<_, anyhow::Error>(()) })
159            .await;
160        assert!(ok.is_ok());
161    }
162}
github_mcp/core/circuit_breaker.rs

github_mcp/core/
circuit_breaker.rs
