load-balancer 0.6.0

use crate::LoadBalancer;
use std::future::Future;
use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};
use std::time::Duration;
use tokio::sync::RwLock;
use tokio::task::{yield_now, JoinHandle};
use tokio::{spawn, time::sleep};

#[derive(Clone)]
pub struct Entry<T>
where
    T: Send + Sync + Clone + 'static,
{
    pub value: T,
}

pub struct Inner<T>
where
    T: Send + Sync + Clone + 'static,
{
    pub entries: RwLock<Vec<Entry<T>>>,
    pub cursor: AtomicUsize,
}

/// A round-robin load balancer that selects entries in sequential order.
#[derive(Clone)]
pub struct RoundRobin<T>
where
    T: Send + Sync + Clone + 'static,
{
    inner: Arc<Inner<T>>,
}

impl<T> RoundRobin<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Create a new `RoundRobin` from a list of values.
    pub fn new(entries: Vec<T>) -> Self {
        Self {
            inner: Arc::new(Inner {
                entries: RwLock::new(entries.into_iter().map(|v| Entry { value: v }).collect()),
                cursor: AtomicUsize::new(0),
            }),
        }
    }

    /// Update the inner state using an async callback.
    pub async fn update<F, R, N>(&self, handler: F) -> anyhow::Result<N>
    where
        F: Fn(Arc<Inner<T>>) -> R,
        R: Future<Output = anyhow::Result<N>>,
    {
        handler(self.inner.clone()).await
    }

    /// Spawn a background task that calls `handler` every `interval`.
    ///
    /// The handler is called once immediately; if that initial call fails
    /// the error is returned and no background task is spawned.
    pub async fn update_timer<F, R>(
        &self,
        handler: F,
        interval: Duration,
    ) -> anyhow::Result<JoinHandle<()>>
    where
        F: Fn(Arc<Inner<T>>) -> R + Send + Sync + 'static,
        R: Future<Output = anyhow::Result<()>> + Send,
    {
        handler(self.inner.clone()).await?;

        let inner = self.inner.clone();

        Ok(spawn(async move {
            loop {
                sleep(interval).await;
                let _ = handler(inner.clone()).await;
            }
        }))
    }
}

impl<T> LoadBalancer<T> for RoundRobin<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Asynchronously allocate the next entry in sequence.
    async fn alloc(&self) -> T {
        loop {
            match LoadBalancer::try_alloc(self) {
                Some(v) => return v,
                None => yield_now().await,
            }
        }
    }

    /// Try to allocate the next entry in sequence without awaiting.
    fn try_alloc(&self) -> Option<T> {
        let entries = self.inner.entries.try_read().ok()?;

        if entries.is_empty() {
            return None;
        }

        Some(
            entries[self.inner.cursor.fetch_add(1, Ordering::Relaxed) % entries.len()]
                .value
                .clone(),
        )
    }
}