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#[derive(Clone)]
11pub struct Entry<T>
12where
13    T: Send + Sync + Clone + 'static,
14{
15    pub value: T,
16}
17
18pub struct Inner<T>
19where
20    T: Send + Sync + Clone + 'static,
21{
22    pub entries: RwLock<Vec<Entry<T>>>,
23}
24
25/// A load balancer that randomly selects an entry.
26#[derive(Clone)]
27pub struct Random<T>
28where
29    T: Send + Sync + Clone + 'static,
30{
31    inner: Arc<Inner<T>>,
32}
33
34impl<T> Random<T>
35where
36    T: Send + Sync + Clone + 'static,
37{
38    /// Create a new random load balancer from a vector of values.
39    pub fn new(entries: Vec<T>) -> Self {
40        Self {
41            inner: Arc::new(Inner {
42                entries: RwLock::new(
43                    entries.into_iter().map(|value| Entry { value }).collect(),
44                ),
45            }),
46        }
47    }
48
49    /// Run an async callback with access to the internal state.
50    pub async fn update<F, R, N>(&self, handler: F) -> anyhow::Result<N>
51    where
52        F: Fn(Arc<Inner<T>>) -> R,
53        R: Future<Output = anyhow::Result<N>>,
54    {
55        handler(self.inner.clone()).await
56    }
57
58    /// Spawn a background task that calls `handler` every `interval`.
59    ///
60    /// The handler is called once immediately; if that initial call fails
61    /// the error is returned and no background task is spawned.
62    pub async fn update_timer<F, R>(
63        &self,
64        handler: F,
65        interval: Duration,
66    ) -> anyhow::Result<JoinHandle<()>>
67    where
68        F: Fn(Arc<Inner<T>>) -> R + Send + Sync + 'static,
69        R: Future<Output = anyhow::Result<()>> + Send,
70    {
71        handler(self.inner.clone()).await?;
72
73        let inner = self.inner.clone();
74
75        Ok(spawn(async move {
76            loop {
77                sleep(interval).await;
78                let _ = handler(inner.clone()).await;
79            }
80        }))
81    }
82}
83
84impl<T> LoadBalancer<T> for Random<T>
85where
86    T: Send + Sync + Clone + 'static,
87{
88    async fn alloc(&self) -> T {
89        self.inner
90            .entries
91            .read()
92            .await
93            .iter()
94            .choose(&mut rand::rng())
95            .map(|v| v.value.clone())
96            .unwrap()
97    }
98
99    fn try_alloc(&self) -> Option<T> {
100        self.inner
101            .entries
102            .try_read()
103            .ok()?
104            .iter()
105            .choose(&mut rand::rng())
106            .map(|v| v.value.clone())
107    }
108}
load_balancer/random.rs

load_balancer/
random.rs
