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) -> T {
71        let entries = self.inner.entries.read().await;
72
73        entries[self.inner.cursor.fetch_add(1, Ordering::Relaxed) % entries.len()]
74            .value
75            .clone()
76    }
77
78    /// Try to allocate the next entry in sequence without awaiting.
79    fn try_alloc(&self) -> Option<T> {
80        let entries = self.inner.entries.try_read().ok()?;
81
82        if entries.is_empty() {
83            return None;
84        }
85
86        Some(
87            entries[self.inner.cursor.fetch_add(1, Ordering::Relaxed) % entries.len()]
88                .value
89                .clone(),
90        )
91    }
92}
93
94#[async_trait]
95impl<T> BoxLoadBalancer<T> for SimpleLoadBalancer<T>
96where
97    T: Send + Sync + Clone + 'static,
98{
99    /// Asynchronously allocate the next entry (BoxLoadBalancer version).
100    async fn alloc(&self) -> T {
101        LoadBalancer::alloc(self).await
102    }
103
104    /// Try to allocate the next entry (BoxLoadBalancer version).
105    fn try_alloc(&self) -> Option<T> {
106        LoadBalancer::try_alloc(self)
107    }
108}
load_balancer/simple.rs

load_balancer/
simple.rs
