load-balancer 0.5.0

use tokio::{spawn, sync::Mutex, sync::RwLock, task::JoinHandle, task::yield_now, time::sleep};

use crate::LoadBalancer;
use std::{
    future::Future,
    sync::Arc,
    time::{Duration, Instant},
};

/// A single entry in the cooldown load balancer.
///
/// Each entry tracks the minimum reuse interval, the last allocation
/// timestamp, and the associated value.
pub struct Entry<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Minimum duration before this entry can be reused.
    pub interval: Duration,
    /// Timestamp of the last allocation (`None` if never used).
    pub last: Mutex<Option<Instant>>,
    /// The underlying resource value.
    pub value: T,
}

impl<T> Clone for Entry<T>
where
    T: Send + Sync + Clone + 'static,
{
    fn clone(&self) -> Self {
        Self {
            interval: self.interval.clone(),
            last: self.last.try_lock().unwrap().clone().into(),
            value: self.value.clone(),
        }
    }
}

/// A load balancer that allocates items based on a cooldown interval.
/// Each entry can only be reused after its interval has elapsed since the last allocation.
#[derive(Clone)]
pub struct Cooldown<T>
where
    T: Send + Sync + Clone + 'static,
{
    inner: Arc<RwLock<Vec<Entry<T>>>>,
}

impl<T> Cooldown<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Create a new `Cooldown` with a list of `(interval, value)` pairs.
    /// Each value will only be available after its interval has passed since the last allocation.
    pub fn new(entries: Vec<(Duration, T)>) -> Self {
        Self {
            inner: Arc::new(RwLock::new(
                entries
                    .into_iter()
                    .map(|(interval, value)| Entry {
                        interval,
                        last: None.into(),
                        value,
                    })
                    .collect(),
            )),
        }
    }

    /// Update the internal entries using an async callback.
    /// This allows dynamic reconfiguration of the load balancer.
    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 Cooldown<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Allocate a value asynchronously.
    /// This will loop until a value becomes available, yielding in between attempts.
    fn alloc(&self) -> impl Future<Output = T> + Send {
        async move {
            loop {
                if let Some(v) = LoadBalancer::try_alloc(self) {
                    return v;
                }

                let min_remaining = {
                    let entries = self.inner.read().await;
                    let mut min = None;

                    for entry in entries.iter() {
                        if entry.interval == Duration::ZERO {
                            continue;
                        }

                        if let Some(last_time) = *entry.last.lock().await {
                            let now = Instant::now();
                            let elapsed = now.duration_since(last_time);

                            if elapsed < entry.interval {
                                let remaining = entry.interval - elapsed;

                                if min.is_none() || remaining < min.unwrap() {
                                    min = Some(remaining);
                                }
                            }
                        }
                    }

                    min
                };

                if let Some(duration) = min_remaining {
                    tokio::time::sleep(duration).await;
                } else {
                    yield_now().await;
                }
            }
        }
    }

    /// Try to allocate a value immediately without waiting.
    /// Returns `Some(value)` if an entry is available (interval elapsed),
    /// otherwise returns `None`.
    fn try_alloc(&self) -> Option<T> {
        let entries = self.inner.try_read().ok()?;

        for entry in entries.iter() {
            if entry.interval == Duration::ZERO {
                return Some(entry.value.clone());
            }

            if let Ok(mut last) = entry.last.try_lock() {
                match *last {
                    Some(v) => {
                        let now = Instant::now();

                        if now.duration_since(v) >= entry.interval {
                            *last = Some(now);
                            return Some(entry.value.clone());
                        }
                    }
                    None => {
                        *last = Some(Instant::now());
                        return Some(entry.value.clone());
                    }
                }
            }
        }

        None
    }
}