load-balancer 0.6.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};

#[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>>>,
}

/// A load balancer that randomly selects an entry.
#[derive(Clone)]
pub struct Random<T>
where
    T: Send + Sync + Clone + 'static,
{
    inner: Arc<Inner<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(entries: Vec<T>) -> Self {
        Self {
            inner: Arc::new(Inner {
                entries: RwLock::new(
                    entries.into_iter().map(|value| Entry { value }).collect(),
                ),
            }),
        }
    }

    /// Run an async callback with access to the internal state.
    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 Random<T>
where
    T: Send + Sync + Clone + 'static,
{
    async fn alloc(&self) -> T {
        self.inner
            .entries
            .read()
            .await
            .iter()
            .choose(&mut rand::rng())
            .map(|v| v.value.clone())
            .unwrap()
    }

    fn try_alloc(&self) -> Option<T> {
        self.inner
            .entries
            .try_read()
            .ok()?
            .iter()
            .choose(&mut rand::rng())
            .map(|v| v.value.clone())
    }
}