rostrum 14.0.1

use anyhow::Result;
use futures_util::Future;
use std::{
    ops::{Add, Sub},
    time::{Duration, Instant},
};

use crate::{errors::rpc_timeout, signal::Waiter, thread::run_with_timeout};

pub struct TimeoutTrigger {
    start: Instant,
    timeout: Duration,
}

impl TimeoutTrigger {
    pub fn new(timeout: Duration) -> TimeoutTrigger {
        TimeoutTrigger {
            start: Instant::now(),
            timeout,
        }
    }

    #[inline]
    pub fn check(&self) -> Result<()> {
        Waiter::shutdown_check()?;
        if self.start.elapsed() >= self.timeout {
            return Err(rpc_timeout(&self.timeout));
        }
        Ok(())
    }

    pub async fn timeout<F>(&self, future: F) -> Result<F::Output>
    where
        F: Future,
    {
        let res = run_with_timeout(self.timeout, future).await;
        match res {
            Ok(r) => Ok(r),
            Err(_) => Err(rpc_timeout(&self.timeout)),
        }
    }
}

pub struct LeakyBucket {
    /// Duration since last time we filled the level
    last_checked: Instant,
    /// How much to fill level per second
    fill_rate: f64,
    /// Maximum allowed burst (in seconds)
    max_capacity: u32,
    /// Current maximum allowed burst.
    current_level: f64,
}

impl LeakyBucket {
    pub fn new(fill_rate: f64, capacity: u32) -> LeakyBucket {
        LeakyBucket {
            last_checked: Instant::now(),
            fill_rate,
            max_capacity: capacity,
            current_level: capacity as f64,
        }
    }

    fn refill(&mut self) {
        let now = Instant::now();
        let elapsed = now.duration_since(self.last_checked);
        let fills = elapsed.as_secs_f64() / self.fill_rate;

        self.current_level = f64::min(self.current_level.add(fills), self.max_capacity as f64);

        self.last_checked = now;
    }

    pub fn check_out(&mut self) -> TimeoutTrigger {
        self.refill();

        let timeout = if self.current_level <= 0.25 {
            // If user has spent all his time quota, only allow RPC calls up to 250ms.
            Duration::from_millis(250)
        } else {
            Duration::from_secs(self.current_level as u64)
        };

        TimeoutTrigger {
            start: Instant::now(),
            timeout,
        }
    }

    pub fn check_in(&mut self, spent: &TimeoutTrigger) {
        let time_spent = spent.start.elapsed();
        self.current_level = self.current_level.sub(time_spent.as_secs_f64());
    }
}

#[cfg(test)]
mod tests {

    use super::*;
    use std::thread::sleep;
    #[cfg(feature = "with-benchmarks")]
    use test::Bencher;

    #[test]
    fn test_timeout() {
        let timeout = TimeoutTrigger::new(Duration::from_millis(50));
        assert!(!timeout.check().is_err());
        sleep(Duration::from_millis(100));
        assert!(timeout.check().is_err());
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_timeout_check(b: &mut Bencher) {
        let timeout = TimeoutTrigger::new(Duration::from_secs(60));
        b.iter(|| {
            let x: Vec<Result<()>> = (0..100_000).map(|_| timeout.check()).collect();
            x
        })
    }

    #[test]
    fn test_leaky_bucket_refill() {
        let mut bucket = LeakyBucket::new(1.0, 10);
        bucket.current_level = 5.0;
        sleep(Duration::from_millis(100));
        bucket.refill();
        // Check if bucket refilled correctly
        assert!((bucket.current_level - 5.1).abs() < 0.05);
    }

    #[test]
    fn test_leaky_bucket_checkout_and_checkin() {
        let mut bucket = LeakyBucket::new(1.0, 10); // Adjusted fill rate for faster testing

        // Perform checkout
        let timeout_trigger = bucket.check_out();

        // Simulate some work that takes 100 ms
        sleep(Duration::from_millis(100));

        // Check-in the timeout_trigger
        bucket.check_in(&timeout_trigger);

        println!("level {}", bucket.current_level);
        assert!((bucket.current_level - 9.9).abs() < 0.15);
    }
}