load-balancer 0.6.0

use crate::FilteredLoadBalancer;
use crate::LoadBalancer;
use std::sync::atomic::Ordering::{Acquire, Release};
use std::{
    future::Future,
    sync::{Arc, atomic::AtomicU64},
    time::{Duration, Instant},
};
use tokio::{
    spawn,
    sync::{Mutex, RwLock},
    task::{JoinHandle, yield_now},
    time::sleep,
};

/// A single entry in the fault-tolerant load balancer.
///
/// `max_count == 0` = unlimited.  `max_error_count == 0` = never auto-disable.
pub struct Entry<T>
where
    T: Send + Sync + Clone + 'static,
{
    pub max_count: u64,
    pub max_error_count: u64,
    pub count: AtomicU64,
    pub error_count: AtomicU64,
    pub value: T,
}

impl<T> Entry<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Clear both counts to zero.
    pub fn reset(&self) {
        self.count.store(0, Release);
        self.error_count.store(0, Release);
    }

    /// Mark this entry as disabled (sets `error_count` to `max_error_count`).
    pub fn disable(&self) {
        self.error_count.store(self.max_error_count, Release);
    }
}

impl<T> Clone for Entry<T>
where
    T: Send + Sync + Clone + 'static,
{
    fn clone(&self) -> Self {
        Self {
            max_count: self.max_count.clone(),
            max_error_count: self.max_error_count.clone(),
            count: self.count.load(Acquire).into(),
            error_count: self.error_count.load(Acquire).into(),
            value: self.value.clone(),
        }
    }
}

pub struct Inner<T>
where
    T: Send + Sync + Clone + 'static,
{
    pub entries: RwLock<Vec<Entry<T>>>,
    pub timer: Mutex<Option<JoinHandle<()>>>,
    pub interval: RwLock<Duration>,
    pub next_reset: RwLock<Instant>,
}

/// Fault-tolerant load balancer: entries auto-disable after too many errors.
///
/// Each entry has per-interval allocation and error caps.  Call [`FaultTolerant::success`]
/// / [`FaultTolerant::failure`] with the index returned by [`FilteredLoadBalancer::alloc_filter`]
/// to adjust error counts.  Disabled entries revive on the next reset.
#[derive(Clone)]
pub struct FaultTolerant<T>
where
    T: Send + Sync + Clone + 'static,
{
    inner: Arc<Inner<T>>,
}

impl<T> FaultTolerant<T>
where
    T: Send + Sync + Clone + 'static,
{
    /// Create a `FaultTolerant` with a 1-second reset interval.
    ///
    /// Each tuple is `(max_count, max_error_count, value)`.
    /// `max_error_count == 0` means errors never disable the entry.
    pub fn new(entries: Vec<(u64, u64, T)>) -> Self {
        Self::new_interval(entries, Duration::from_secs(1))
    }

    /// Create a `FaultTolerant` with a custom reset `interval`.
    ///
    /// Counts (both allocation and error) reset every `interval`.
    pub fn new_interval(entries: Vec<(u64, u64, T)>, interval: Duration) -> Self {
        Self {
            inner: Arc::new(Inner {
                entries: entries
                    .into_iter()
                    .map(|(max_count, max_error_count, value)| Entry {
                        max_count,
                        max_error_count,
                        value,
                        count: 0.into(),
                        error_count: 0.into(),
                    })
                    .collect::<Vec<_>>()
                    .into(),
                timer: Mutex::new(None),
                interval: interval.into(),
                next_reset: RwLock::new(Instant::now() + interval),
            }),
        }
    }

    /// 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.
    /// Returns a [`JoinHandle`] that can be aborted to stop the timer.
    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;
            }
        }))
    }

    /// Decrement the error count of the entry at `index` (floor of zero).
    pub fn success(&self, index: usize) {
        if let Ok(entries) = self.inner.entries.try_read() {
            if let Some(entry) = entries.get(index) {
                let _ = entry.error_count.fetch_update(Acquire, Acquire, |current| {
                    if current == 0 {
                        None
                    } else {
                        Some(current - 1)
                    }
                });
            }
        }
    }

    /// Increment the error count of the entry at `index`.
    /// When it reaches `max_error_count` the entry is disabled until the next reset.
    pub fn failure(&self, index: usize) {
        if let Ok(entries) = self.inner.entries.try_read() {
            if let Some(entry) = entries.get(index) {
                entry.error_count.fetch_add(1, Release);
            }
        }
    }
}

impl<T> FilteredLoadBalancer<T> for FaultTolerant<T>
where
    T: Clone + Send + Sync + 'static,
{
    async fn alloc_filter(&self, mut predicate: impl FnMut((usize, &T)) -> bool) -> (usize, T) {
        loop {
            if let Some(result) = self.try_alloc_filter(&mut predicate) {
                return result;
            }

            let now = Instant::now();

            let next = *self.inner.next_reset.read().await;

            let remaining = if now < next {
                next - now
            } else {
                Duration::ZERO
            };

            if remaining > Duration::ZERO {
                sleep(remaining).await;
            } else {
                yield_now().await;
            }
        }
    }

    fn try_alloc_filter(
        &self,
        mut predicate: impl FnMut((usize, &T)) -> bool,
    ) -> Option<(usize, T)> {
        if let Ok(mut timer_guard) = self.inner.timer.try_lock() {
            if timer_guard.is_none() {
                let this = self.inner.clone();

                *timer_guard = Some(spawn(async move {
                    let mut interval = *this.interval.read().await;

                    *this.next_reset.write().await = Instant::now() + interval;

                    loop {
                        sleep(match this.interval.try_read() {
                            Ok(v) => {
                                interval = *v;
                                interval
                            }
                            Err(_) => interval,
                        })
                        .await;

                        let now = Instant::now();

                        let entries = this.entries.read().await;

                        for entry in entries.iter() {
                            entry.count.store(0, Release);
                        }

                        *this.next_reset.write().await = now + interval;
                    }
                }));
            }
        }

        if let Ok(entries) = self.inner.entries.try_read() {
            for (i, entry) in entries.iter().enumerate() {
                if entry.max_error_count != 0
                    && entry.error_count.load(Acquire) >= entry.max_error_count
                {
                    continue;
                }

                if !predicate((i, &entry.value)) {
                    continue;
                }

                let count = entry.count.load(Acquire);

                if entry.max_count == 0
                    || (count < entry.max_count
                        && entry
                            .count
                            .compare_exchange(count, count + 1, Release, Acquire)
                            .is_ok())
                {
                    return Some((i, entry.value.clone()));
                }
            }
        }

        None
    }
}

impl<T> LoadBalancer<T> for FaultTolerant<T>
where
    T: Clone + Send + Sync + 'static,
{
    fn alloc(&self) -> impl Future<Output = T> + Send {
        async move { self.alloc_filter(|_| true).await.1 }
    }

    fn try_alloc(&self) -> Option<T> {
        self.try_alloc_filter(|_| true).map(|v| v.1)
    }
}