load_balancer/
cooldown.rs

1use tokio::{spawn, sync::Mutex, sync::RwLock, task::JoinHandle, task::yield_now, time::sleep};
2
3use crate::LoadBalancer;
4use std::{
5    future::Future,
6    sync::Arc,
7    time::{Duration, Instant},
8};
9
10/// A single entry in the cooldown load balancer.
11///
12/// Each entry tracks the minimum reuse interval, the last allocation
13/// timestamp, and the associated value.
14pub struct Entry<T>
15where
16    T: Send + Sync + Clone + 'static,
17{
18    /// Minimum duration before this entry can be reused.
19    pub interval: Duration,
20    /// Timestamp of the last allocation (`None` if never used).
21    pub last: Mutex<Option<Instant>>,
22    /// The underlying resource value.
23    pub value: T,
24}
25
26impl<T> Clone for Entry<T>
27where
28    T: Send + Sync + Clone + 'static,
29{
30    fn clone(&self) -> Self {
31        Self {
32            interval: self.interval.clone(),
33            last: self.last.try_lock().unwrap().clone().into(),
34            value: self.value.clone(),
35        }
36    }
37}
38
39/// A load balancer that allocates items based on a cooldown interval.
40/// Each entry can only be reused after its interval has elapsed since the last allocation.
41#[derive(Clone)]
42pub struct Cooldown<T>
43where
44    T: Send + Sync + Clone + 'static,
45{
46    inner: Arc<RwLock<Vec<Entry<T>>>>,
47}
48
49impl<T> Cooldown<T>
50where
51    T: Send + Sync + Clone + 'static,
52{
53    /// Create a new `Cooldown` with a list of `(interval, value)` pairs.
54    /// Each value will only be available after its interval has passed since the last allocation.
55    pub fn new(entries: Vec<(Duration, T)>) -> Self {
56        Self {
57            inner: Arc::new(RwLock::new(
58                entries
59                    .into_iter()
60                    .map(|(interval, value)| Entry {
61                        interval,
62                        last: None.into(),
63                        value,
64                    })
65                    .collect(),
66            )),
67        }
68    }
69
70    /// Update the internal entries using an async callback.
71    /// This allows dynamic reconfiguration of the load balancer.
72    pub async fn update<F, R, N>(&self, handler: F) -> anyhow::Result<N>
73    where
74        F: Fn(Arc<RwLock<Vec<Entry<T>>>>) -> R,
75        R: Future<Output = anyhow::Result<N>>,
76    {
77        handler(self.inner.clone()).await
78    }
79
80    /// Spawn a background task that calls `handler` every `interval`.
81    pub async fn update_timer<F, R>(
82        &self,
83        handler: F,
84        interval: Duration,
85    ) -> anyhow::Result<JoinHandle<()>>
86    where
87        F: Fn(Arc<RwLock<Vec<Entry<T>>>>) -> R + Send + Sync + 'static,
88        R: Future<Output = anyhow::Result<()>> + Send,
89    {
90        handler(self.inner.clone()).await?;
91
92        let inner = self.inner.clone();
93
94        Ok(spawn(async move {
95            loop {
96                sleep(interval).await;
97                let _ = handler(inner.clone()).await;
98            }
99        }))
100    }
101}
102
103impl<T> LoadBalancer<T> for Cooldown<T>
104where
105    T: Send + Sync + Clone + 'static,
106{
107    /// Allocate a value asynchronously.
108    /// This will loop until a value becomes available, yielding in between attempts.
109    fn alloc(&self) -> impl Future<Output = T> + Send {
110        async move {
111            loop {
112                if let Some(v) = LoadBalancer::try_alloc(self) {
113                    return v;
114                }
115
116                let min_remaining = {
117                    let entries = self.inner.read().await;
118                    let mut min = None;
119
120                    for entry in entries.iter() {
121                        if entry.interval == Duration::ZERO {
122                            continue;
123                        }
124
125                        if let Some(last_time) = *entry.last.lock().await {
126                            let now = Instant::now();
127                            let elapsed = now.duration_since(last_time);
128
129                            if elapsed < entry.interval {
130                                let remaining = entry.interval - elapsed;
131
132                                if min.is_none() || remaining < min.unwrap() {
133                                    min = Some(remaining);
134                                }
135                            }
136                        }
137                    }
138
139                    min
140                };
141
142                if let Some(duration) = min_remaining {
143                    tokio::time::sleep(duration).await;
144                } else {
145                    yield_now().await;
146                }
147            }
148        }
149    }
150
151    /// Try to allocate a value immediately without waiting.
152    /// Returns `Some(value)` if an entry is available (interval elapsed),
153    /// otherwise returns `None`.
154    fn try_alloc(&self) -> Option<T> {
155        let entries = self.inner.try_read().ok()?;
156
157        for entry in entries.iter() {
158            if entry.interval == Duration::ZERO {
159                return Some(entry.value.clone());
160            }
161
162            if let Ok(mut last) = entry.last.try_lock() {
163                match *last {
164                    Some(v) => {
165                        let now = Instant::now();
166
167                        if now.duration_since(v) >= entry.interval {
168                            *last = Some(now);
169                            return Some(entry.value.clone());
170                        }
171                    }
172                    None => {
173                        *last = Some(Instant::now());
174                        return Some(entry.value.clone());
175                    }
176                }
177            }
178        }
179
180        None
181    }
182}
load_balancer/cooldown.rs

load_balancer/
cooldown.rs
