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, N>(&self, handle: F) -> anyhow::Result<N>
43    where
44        F: Fn(Arc<RwLock<Vec<Entry<T>>>>) -> R,
45        R: Future<Output = anyhow::Result<N>>,
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) -> T {
57        self.inner
58            .read()
59            .await
60            .iter()
61            .choose(&mut rand::rng())
62            .map(|v| v.value.clone())
63            .unwrap()
64    }
65
66    /// Try to allocate a random entry synchronously.
67    fn try_alloc(&self) -> Option<T> {
68        self.inner
69            .try_read()
70            .ok()?
71            .iter()
72            .choose(&mut rand::rng())
73            .map(|v| v.value.clone())
74    }
75}
76
77#[async_trait]
78impl<T> BoxLoadBalancer<T> for RandomLoadBalancer<T>
79where
80    T: Send + Sync + Clone + 'static,
81{
82    /// Asynchronously allocate a random entry.
83    async fn alloc(&self) -> T {
84        self.inner
85            .read()
86            .await
87            .iter()
88            .choose(&mut rand::rng())
89            .map(|v| v.value.clone())
90            .unwrap()
91    }
92
93    /// Try to allocate a random entry synchronously.
94    fn try_alloc(&self) -> Option<T> {
95        self.inner
96            .try_read()
97            .ok()?
98            .iter()
99            .choose(&mut rand::rng())
100            .map(|v| v.value.clone())
101    }
102}
load_balancer/random.rs

load_balancer/
random.rs
