async-throttle 0.3.2

#[cfg(test)]
mod tests {
    use anyhow::Result;
    use async_throttle::MultiRateLimiter;
    use futures::future::join_all;
    use std::ops::Mul;
    use std::sync::atomic::AtomicUsize;
    use std::sync::atomic::Ordering::SeqCst;
    use std::sync::Arc;
    use std::time::{Duration, Instant};

    #[tokio::test]
    async fn sync_throttle_same_key() -> Result<()> {
        let rate_limiter = MultiRateLimiter::new(Duration::from_millis(10));
        let start = Instant::now();
        static COUNT: AtomicUsize = AtomicUsize::new(0);

        for _ in 0..10 {
            rate_limiter
                .throttle("key", || async {
                    COUNT.fetch_add(1, SeqCst);
                })
                .await;
        }

        assert_eq!(COUNT.load(SeqCst), 10);
        assert!(start.elapsed().as_millis() > 89);
        Ok(())
    }

    #[tokio::test]
    async fn sync_throttle_multi_key() -> Result<()> {
        static ONE_THOUSAND_SECONDS: Duration = Duration::from_secs(1000);
        let rate_limiter = Arc::new(MultiRateLimiter::new(ONE_THOUSAND_SECONDS));
        let start = Instant::now();
        static COUNT: AtomicUsize = AtomicUsize::new(0);

        for k in 0..10 {
            rate_limiter
                .throttle(k, || async {
                    COUNT.fetch_add(1, SeqCst);
                })
                .await;
        }

        assert_eq!(COUNT.load(SeqCst), 10);
        assert!(start.elapsed() < ONE_THOUSAND_SECONDS);
        Ok(())
    }

    #[tokio::test]
    async fn async_throttle_same_key() -> Result<()> {
        let rate_limiter = Arc::new(MultiRateLimiter::new(Duration::from_millis(1)));
        let start = Instant::now();
        static COUNT: AtomicUsize = AtomicUsize::new(0);

        join_all((0..100).map(|_| {
            let rate_limiter = rate_limiter.clone();
            tokio::spawn(async move {
                rate_limiter
                    .throttle("key", || async {
                        COUNT.fetch_add(1, SeqCst);
                    })
                    .await;
                Ok::<(), anyhow::Error>(())
            })
        }))
        .await;

        assert_eq!(COUNT.load(SeqCst), 100);
        assert!(start.elapsed().as_millis() > 99);
        Ok(())
    }

    #[tokio::test]
    async fn async_throttle_multi_key_get_once() -> Result<()> {
        static ONE_THOUSAND_SECONDS: Duration = Duration::from_secs(1000);
        let rate_limiter = Arc::new(MultiRateLimiter::new(ONE_THOUSAND_SECONDS));
        let start = Instant::now();
        static COUNT: AtomicUsize = AtomicUsize::new(0);

        join_all((0..1000).map(|x| {
            let rate_limiter = rate_limiter.clone();
            tokio::spawn(async move {
                rate_limiter
                    .throttle(x, || async {
                        COUNT.fetch_add(1, SeqCst);
                    })
                    .await;
                Ok::<(), anyhow::Error>(())
            })
        }))
        .await;

        assert_eq!(COUNT.load(SeqCst), 1000);
        assert!(start.elapsed() < ONE_THOUSAND_SECONDS);
        Ok(())
    }

    #[tokio::test]
    async fn async_throttle_multi_key_get_many_times() -> Result<()> {
        let period = Duration::from_nanos(100);
        let rate_limiter = Arc::new(MultiRateLimiter::new(period));
        let start = Instant::now();
        let (max, radix): (u32, u32) = (1000, 100);
        let min_wait_time = period.mul(max / radix);
        static COUNT: AtomicUsize = AtomicUsize::new(0);

        join_all((0..max).map(|x| {
            let rate_limiter = rate_limiter.clone();
            tokio::spawn(async move {
                let target = x % radix;
                rate_limiter
                    .throttle(target, || async {
                        COUNT.fetch_add(1, SeqCst);
                    })
                    .await;
                Ok::<(), anyhow::Error>(())
            })
        }))
        .await;

        assert_eq!(COUNT.load(SeqCst), max as usize);
        assert!(start.elapsed() > min_wait_time);
        Ok(())
    }
}