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

load_balancer/
cooldown.rs
