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

load_balancer/
round_robin.rs
