load_balancer/
round_robin.rs

1use crate::LoadBalancer;
2use std::future::Future;
3use std::sync::{
4    Arc,
5    atomic::{AtomicUsize, Ordering},
6};
7use std::time::Duration;
8use tokio::sync::RwLock;
9use tokio::task::{yield_now, JoinHandle};
10use tokio::{spawn, time::sleep};
11
12#[derive(Clone)]
13pub struct Entry<T>
14where
15    T: Send + Sync + Clone + 'static,
16{
17    pub value: T,
18}
19
20pub struct Inner<T>
21where
22    T: Send + Sync + Clone + 'static,
23{
24    pub entries: RwLock<Vec<Entry<T>>>,
25    pub cursor: AtomicUsize,
26}
27
28/// A round-robin load balancer that selects entries in sequential order.
29#[derive(Clone)]
30pub struct RoundRobin<T>
31where
32    T: Send + Sync + Clone + 'static,
33{
34    inner: Arc<Inner<T>>,
35}
36
37impl<T> RoundRobin<T>
38where
39    T: Send + Sync + Clone + 'static,
40{
41    /// Create a new `RoundRobin` from a list of values.
42    pub fn new(entries: Vec<T>) -> Self {
43        Self {
44            inner: Arc::new(Inner {
45                entries: RwLock::new(entries.into_iter().map(|v| Entry { value: v }).collect()),
46                cursor: AtomicUsize::new(0),
47            }),
48        }
49    }
50
51    /// Update the inner state using an async callback.
52    pub async fn update<F, R, N>(&self, handler: F) -> anyhow::Result<N>
53    where
54        F: Fn(Arc<Inner<T>>) -> R,
55        R: Future<Output = anyhow::Result<N>>,
56    {
57        handler(self.inner.clone()).await
58    }
59
60    /// Spawn a background task that calls `handler` every `interval`.
61    ///
62    /// The handler is called once immediately; if that initial call fails
63    /// the error is returned and no background task is spawned.
64    pub async fn update_timer<F, R>(
65        &self,
66        handler: F,
67        interval: Duration,
68    ) -> anyhow::Result<JoinHandle<()>>
69    where
70        F: Fn(Arc<Inner<T>>) -> R + Send + Sync + 'static,
71        R: Future<Output = anyhow::Result<()>> + Send,
72    {
73        handler(self.inner.clone()).await?;
74
75        let inner = self.inner.clone();
76
77        Ok(spawn(async move {
78            loop {
79                sleep(interval).await;
80                let _ = handler(inner.clone()).await;
81            }
82        }))
83    }
84}
85
86impl<T> LoadBalancer<T> for RoundRobin<T>
87where
88    T: Send + Sync + Clone + 'static,
89{
90    /// Asynchronously allocate the next entry in sequence.
91    async fn alloc(&self) -> T {
92        loop {
93            match LoadBalancer::try_alloc(self) {
94                Some(v) => return v,
95                None => yield_now().await,
96            }
97        }
98    }
99
100    /// Try to allocate the next entry in sequence without awaiting.
101    fn try_alloc(&self) -> Option<T> {
102        let entries = self.inner.entries.try_read().ok()?;
103
104        if entries.is_empty() {
105            return None;
106        }
107
108        Some(
109            entries[self.inner.cursor.fetch_add(1, Ordering::Relaxed) % entries.len()]
110                .value
111                .clone(),
112        )
113    }
114}
load_balancer/round_robin.rs

load_balancer/
round_robin.rs
