load_balancer/
simple.rs

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

load_balancer/
simple.rs
