load_balancer/
random.rs

1use crate::LoadBalancer;
2use rand::seq::IteratorRandom;
3use std::future::Future;
4use std::sync::Arc;
5use std::time::Duration;
6use tokio::sync::RwLock;
7use tokio::task::JoinHandle;
8use tokio::{spawn, time::sleep};
9
10/// A single entry in the Random load balancer.
11#[derive(Clone)]
12pub struct Entry<T>
13where
14    T: Send + Sync + Clone + 'static,
15{
16    /// The underlying value of type `T`.
17    pub value: T,
18}
19
20/// A load balancer that randomly selects an entry.
21#[derive(Clone)]
22pub struct Random<T>
23where
24    T: Send + Sync + Clone + 'static,
25{
26    /// The inner list of entries.
27    inner: Arc<RwLock<Vec<Entry<T>>>>,
28}
29
30impl<T> Random<T>
31where
32    T: Send + Sync + Clone + 'static,
33{
34    /// Create a new random load balancer from a vector of values.
35    pub fn new(inner: Vec<T>) -> Self {
36        Self {
37            inner: Arc::new(RwLock::new(
38                inner.into_iter().map(|value| Entry { value }).collect(),
39            )),
40        }
41    }
42
43    /// Update the inner entries using an async callback.
44    pub async fn update<F, R, N>(&self, handler: F) -> anyhow::Result<N>
45    where
46        F: Fn(Arc<RwLock<Vec<Entry<T>>>>) -> R,
47        R: Future<Output = anyhow::Result<N>>,
48    {
49        handler(self.inner.clone()).await
50    }
51
52    /// Spawn a background task that calls `handler` every `interval`.
53    pub async fn update_timer<F, R>(
54        &self,
55        handler: F,
56        interval: Duration,
57    ) -> anyhow::Result<JoinHandle<()>>
58    where
59        F: Fn(Arc<RwLock<Vec<Entry<T>>>>) -> R + Send + Sync + 'static,
60        R: Future<Output = anyhow::Result<()>> + Send,
61    {
62        handler(self.inner.clone()).await?;
63
64        let inner = self.inner.clone();
65
66        Ok(spawn(async move {
67            loop {
68                sleep(interval).await;
69                let _ = handler(inner.clone()).await;
70            }
71        }))
72    }
73}
74
75impl<T> LoadBalancer<T> for Random<T>
76where
77    T: Send + Sync + Clone + 'static,
78{
79    /// Asynchronously allocate a random entry.
80    async fn alloc(&self) -> T {
81        self.inner
82            .read()
83            .await
84            .iter()
85            .choose(&mut rand::rng())
86            .map(|v| v.value.clone())
87            .unwrap()
88    }
89
90    /// Try to allocate a random entry synchronously.
91    fn try_alloc(&self) -> Option<T> {
92        self.inner
93            .try_read()
94            .ok()?
95            .iter()
96            .choose(&mut rand::rng())
97            .map(|v| v.value.clone())
98    }
99}
load_balancer/random.rs

load_balancer/
random.rs
