load-balancer 0.5.0

use crate::LoadBalancer;
use rand::seq::IteratorRandom;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
use tokio::{spawn, time::sleep};

/// A single entry in the Random load balancer.
#[derive(Clone)]
pub struct Entry<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// The underlying value of type `T`.
    pub value: T,
}

/// A load balancer that randomly selects an entry.
#[derive(Clone)]
pub struct Random<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// The inner list of entries.
    inner: Arc<RwLock<Vec<Entry<T>>>>,
}

impl<T> Random<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Create a new random load balancer from a vector of values.
    pub fn new(inner: Vec<T>) -> Self {
        Self {
            inner: Arc::new(RwLock::new(
                inner.into_iter().map(|value| Entry { value }).collect(),
            )),
        }
    }

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

    /// Spawn a background task that calls `handler` every `interval`.
    pub async fn update_timer<F, R>(
        &self,
        handler: F,
        interval: Duration,
    ) -> anyhow::Result<JoinHandle<()>>
    where
        F: Fn(Arc<RwLock<Vec<Entry<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 Random<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Asynchronously allocate a random entry.
    async fn alloc(&self) -> T {
        self.inner
            .read()
            .await
            .iter()
            .choose(&mut rand::rng())
            .map(|v| v.value.clone())
            .unwrap()
    }

    /// Try to allocate a random entry synchronously.
    fn try_alloc(&self) -> Option<T> {
        self.inner
            .try_read()
            .ok()?
            .iter()
            .choose(&mut rand::rng())
            .map(|v| v.value.clone())
    }
}