load_balancer/
random.rs

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

load_balancer/
random.rs
