load_balancer/
fault_tolerant.rs

1use crate::FilteredLoadBalancer;
2use crate::LoadBalancer;
3use std::sync::atomic::Ordering::{Acquire, Release};
4use std::{
5    future::Future,
6    sync::{Arc, atomic::AtomicU64},
7    time::{Duration, Instant},
8};
9use tokio::{
10    spawn,
11    sync::{Mutex, RwLock},
12    task::{JoinHandle, yield_now},
13    time::sleep,
14};
15
16/// A single entry in the fault-tolerant load balancer.
17///
18/// `max_count == 0` = unlimited.  `max_error_count == 0` = never auto-disable.
19pub struct Entry<T>
20where
21    T: Send + Sync + Clone + 'static,
22{
23    pub max_count: u64,
24    pub max_error_count: u64,
25    pub count: AtomicU64,
26    pub error_count: AtomicU64,
27    pub value: T,
28}
29
30impl<T> Entry<T>
31where
32    T: Send + Sync + Clone + 'static,
33{
34    /// Clear both counts to zero.
35    pub fn reset(&self) {
36        self.count.store(0, Release);
37        self.error_count.store(0, Release);
38    }
39
40    /// Mark this entry as disabled (sets `error_count` to `max_error_count`).
41    pub fn disable(&self) {
42        self.error_count.store(self.max_error_count, Release);
43    }
44}
45
46impl<T> Clone for Entry<T>
47where
48    T: Send + Sync + Clone + 'static,
49{
50    fn clone(&self) -> Self {
51        Self {
52            max_count: self.max_count.clone(),
53            max_error_count: self.max_error_count.clone(),
54            count: self.count.load(Acquire).into(),
55            error_count: self.error_count.load(Acquire).into(),
56            value: self.value.clone(),
57        }
58    }
59}
60
61pub struct Inner<T>
62where
63    T: Send + Sync + Clone + 'static,
64{
65    pub entries: RwLock<Vec<Entry<T>>>,
66    pub timer: Mutex<Option<JoinHandle<()>>>,
67    pub interval: RwLock<Duration>,
68    pub next_reset: RwLock<Instant>,
69}
70
71/// Fault-tolerant load balancer: entries auto-disable after too many errors.
72///
73/// Each entry has per-interval allocation and error caps.  Call [`FaultTolerant::success`]
74/// / [`FaultTolerant::failure`] with the index returned by [`FilteredLoadBalancer::alloc_filter`]
75/// to adjust error counts.  Disabled entries revive on the next reset.
76#[derive(Clone)]
77pub struct FaultTolerant<T>
78where
79    T: Send + Sync + Clone + 'static,
80{
81    inner: Arc<Inner<T>>,
82}
83
84impl<T> FaultTolerant<T>
85where
86    T: Send + Sync + Clone + 'static,
87{
88    /// Create a `FaultTolerant` with a 1-second reset interval.
89    ///
90    /// Each tuple is `(max_count, max_error_count, value)`.
91    /// `max_error_count == 0` means errors never disable the entry.
92    pub fn new(entries: Vec<(u64, u64, T)>) -> Self {
93        Self::new_interval(entries, Duration::from_secs(1))
94    }
95
96    /// Create a `FaultTolerant` with a custom reset `interval`.
97    ///
98    /// Counts (both allocation and error) reset every `interval`.
99    pub fn new_interval(entries: Vec<(u64, u64, T)>, interval: Duration) -> Self {
100        Self {
101            inner: Arc::new(Inner {
102                entries: entries
103                    .into_iter()
104                    .map(|(max_count, max_error_count, value)| Entry {
105                        max_count,
106                        max_error_count,
107                        value,
108                        count: 0.into(),
109                        error_count: 0.into(),
110                    })
111                    .collect::<Vec<_>>()
112                    .into(),
113                timer: Mutex::new(None),
114                interval: interval.into(),
115                next_reset: RwLock::new(Instant::now() + interval),
116            }),
117        }
118    }
119
120    /// Run an async callback with access to the internal state.
121    pub async fn update<F, R, N>(&self, handler: F) -> anyhow::Result<N>
122    where
123        F: Fn(Arc<Inner<T>>) -> R,
124        R: Future<Output = anyhow::Result<N>>,
125    {
126        handler(self.inner.clone()).await
127    }
128
129    /// Spawn a background task that calls `handler` every `interval`.
130    ///
131    /// The handler is called once immediately; if that initial call fails
132    /// the error is returned and no background task is spawned.
133    /// Returns a [`JoinHandle`] that can be aborted to stop the timer.
134    pub async fn update_timer<F, R>(
135        &self,
136        handler: F,
137        interval: Duration,
138    ) -> anyhow::Result<JoinHandle<()>>
139    where
140        F: Fn(Arc<Inner<T>>) -> R + Send + Sync + 'static,
141        R: Future<Output = anyhow::Result<()>> + Send,
142    {
143        handler(self.inner.clone()).await?;
144
145        let inner = self.inner.clone();
146
147        Ok(spawn(async move {
148            loop {
149                sleep(interval).await;
150                let _ = handler(inner.clone()).await;
151            }
152        }))
153    }
154
155    /// Decrement the error count of the entry at `index` (floor of zero).
156    pub fn success(&self, index: usize) {
157        if let Ok(entries) = self.inner.entries.try_read() {
158            if let Some(entry) = entries.get(index) {
159                let _ = entry.error_count.fetch_update(Acquire, Acquire, |current| {
160                    if current == 0 {
161                        None
162                    } else {
163                        Some(current - 1)
164                    }
165                });
166            }
167        }
168    }
169
170    /// Increment the error count of the entry at `index`.
171    /// When it reaches `max_error_count` the entry is disabled until the next reset.
172    pub fn failure(&self, index: usize) {
173        if let Ok(entries) = self.inner.entries.try_read() {
174            if let Some(entry) = entries.get(index) {
175                entry.error_count.fetch_add(1, Release);
176            }
177        }
178    }
179}
180
181impl<T> FilteredLoadBalancer<T> for FaultTolerant<T>
182where
183    T: Clone + Send + Sync + 'static,
184{
185    async fn alloc_filter(&self, mut predicate: impl FnMut((usize, &T)) -> bool) -> (usize, T) {
186        loop {
187            if let Some(result) = self.try_alloc_filter(&mut predicate) {
188                return result;
189            }
190
191            let now = Instant::now();
192
193            let next = *self.inner.next_reset.read().await;
194
195            let remaining = if now < next {
196                next - now
197            } else {
198                Duration::ZERO
199            };
200
201            if remaining > Duration::ZERO {
202                sleep(remaining).await;
203            } else {
204                yield_now().await;
205            }
206        }
207    }
208
209    fn try_alloc_filter(
210        &self,
211        mut predicate: impl FnMut((usize, &T)) -> bool,
212    ) -> Option<(usize, T)> {
213        if let Ok(mut timer_guard) = self.inner.timer.try_lock() {
214            if timer_guard.is_none() {
215                let this = self.inner.clone();
216
217                *timer_guard = Some(spawn(async move {
218                    let mut interval = *this.interval.read().await;
219
220                    *this.next_reset.write().await = Instant::now() + interval;
221
222                    loop {
223                        sleep(match this.interval.try_read() {
224                            Ok(v) => {
225                                interval = *v;
226                                interval
227                            }
228                            Err(_) => interval,
229                        })
230                        .await;
231
232                        let now = Instant::now();
233
234                        let entries = this.entries.read().await;
235
236                        for entry in entries.iter() {
237                            entry.count.store(0, Release);
238                        }
239
240                        *this.next_reset.write().await = now + interval;
241                    }
242                }));
243            }
244        }
245
246        if let Ok(entries) = self.inner.entries.try_read() {
247            for (i, entry) in entries.iter().enumerate() {
248                if entry.max_error_count != 0
249                    && entry.error_count.load(Acquire) >= entry.max_error_count
250                {
251                    continue;
252                }
253
254                if !predicate((i, &entry.value)) {
255                    continue;
256                }
257
258                let count = entry.count.load(Acquire);
259
260                if entry.max_count == 0
261                    || (count < entry.max_count
262                        && entry
263                            .count
264                            .compare_exchange(count, count + 1, Release, Acquire)
265                            .is_ok())
266                {
267                    return Some((i, entry.value.clone()));
268                }
269            }
270        }
271
272        None
273    }
274}
275
276impl<T> LoadBalancer<T> for FaultTolerant<T>
277where
278    T: Clone + Send + Sync + 'static,
279{
280    fn alloc(&self) -> impl Future<Output = T> + Send {
281        async move { self.alloc_filter(|_| true).await.1 }
282    }
283
284    fn try_alloc(&self) -> Option<T> {
285        self.try_alloc_filter(|_| true).map(|v| v.1)
286    }
287}
load_balancer/fault_tolerant.rs

load_balancer/
fault_tolerant.rs
