Skip to main content

vantage_api_pool/resilient/
builder.rs

1//! Assembling a [`ResilientClient`]: every policy is optional, and a client
2//! built with no options at all is a plain bounded-retry HTTP client.
3
4use std::sync::atomic::AtomicUsize;
5use std::sync::Arc;
6use std::time::Duration;
7
8use tokio::sync::{RwLock, Semaphore};
9
10use super::auth::{AuthRefresher, AuthState};
11use super::breaker::CircuitBreaker;
12use super::rate::TokenBucket;
13use super::{ResilientClient, RetryPolicy, TransportObserver};
14
15/// Builder for [`ResilientClient`].
16pub struct ResilientClientBuilder {
17    http: Option<reqwest::Client>,
18    max_parallel: usize,
19    policy: RetryPolicy,
20    breaker: Option<(usize, Duration, Duration)>,
21    auth: Option<(AuthRefresher, String, String)>,
22    observer: Option<(Arc<str>, Arc<dyn TransportObserver>)>,
23    rate: Option<(f64, usize)>,
24}
25
26impl Default for ResilientClientBuilder {
27    fn default() -> Self {
28        Self {
29            http: None,
30            max_parallel: 8,
31            policy: RetryPolicy::default(),
32            breaker: None,
33            auth: None,
34            observer: None,
35            rate: None,
36        }
37    }
38}
39
40impl ResilientClientBuilder {
41    /// Cap concurrent in-flight requests for this client. Default 8.
42    pub fn max_parallel(mut self, n: usize) -> Self {
43        self.max_parallel = n.max(1);
44        self
45    }
46
47    pub fn retry(mut self, policy: RetryPolicy) -> Self {
48        self.policy = policy;
49        self
50    }
51
52    /// Open the breaker after `threshold` consecutive failures; stay open for
53    /// `cooldown` (fixed), then allow one half-open probe.
54    pub fn circuit_breaker(mut self, threshold: usize, cooldown: Duration) -> Self {
55        self.breaker = Some((threshold, cooldown, cooldown));
56        self
57    }
58
59    /// Like `circuit_breaker`, but each failed probe doubles the cooldown up
60    /// to `max_cooldown`; a success resets it to `base_cooldown`.
61    pub fn circuit_breaker_growing(
62        mut self,
63        threshold: usize,
64        base_cooldown: Duration,
65        max_cooldown: Duration,
66    ) -> Self {
67        self.breaker = Some((threshold, base_cooldown, max_cooldown));
68        self
69    }
70
71    /// The spec's default: 5 failures, 5 s doubling to 60 s.
72    pub fn default_breaker(self) -> Self {
73        self.circuit_breaker_growing(5, Duration::from_secs(5), Duration::from_secs(60))
74    }
75
76    /// Bearer auth, re-acquired lazily and on `401`.
77    pub fn bearer_auth(self, refresher: AuthRefresher) -> Self {
78        self.auth(refresher, "Authorization", "Bearer ")
79    }
80
81    /// Auth with a custom header name and scheme prefix.
82    pub fn auth(
83        mut self,
84        refresher: AuthRefresher,
85        header: impl Into<String>,
86        scheme: impl Into<String>,
87    ) -> Self {
88        self.auth = Some((refresher, header.into(), scheme.into()));
89        self
90    }
91
92    pub fn http_client(mut self, client: reqwest::Client) -> Self {
93        self.http = Some(client);
94        self
95    }
96
97    /// At most `per_second` attempts per second, with a burst of
98    /// `per_second.ceil()`.
99    pub fn rate_limit(self, per_second: f64) -> Self {
100        let burst = per_second.ceil().max(1.0) as usize;
101        self.rate_limit_with_burst(per_second, burst)
102    }
103
104    pub fn rate_limit_with_burst(mut self, per_second: f64, burst: usize) -> Self {
105        self.rate = Some((per_second, burst));
106        self
107    }
108
109    /// Report every attempt, retry and breaker transition under `key`. See
110    /// [`TransportObserver::on_event`] for what the callback may do.
111    pub fn observer(
112        mut self,
113        key: impl Into<Arc<str>>,
114        observer: Arc<dyn TransportObserver>,
115    ) -> Self {
116        self.observer = Some((key.into(), observer));
117        self
118    }
119
120    pub fn build(self) -> ResilientClient {
121        ResilientClient {
122            http: self.http.unwrap_or_default(),
123            semaphore: Arc::new(Semaphore::new(self.max_parallel)),
124            policy: self.policy,
125            breaker: self
126                .breaker
127                .map(|(t, base, max)| Arc::new(CircuitBreaker::new(t, base, max))),
128            auth: self.auth.map(|(refresh, header, scheme)| {
129                Arc::new(AuthState {
130                    token: RwLock::new(None),
131                    refresh,
132                    header,
133                    scheme,
134                })
135            }),
136            in_flight: Arc::new(AtomicUsize::new(0)),
137            peak_in_flight: Arc::new(AtomicUsize::new(0)),
138            observer: self.observer,
139            rate: self.rate.map(|(r, b)| Arc::new(TokenBucket::new(r, b))),
140        }
141    }
142}