Riven/riven/src/req/token_bucket.test.rs

#![cfg(test)]

use fake_instant::FakeInstant as Instant;

/// This is a hack to test token bucket, substituting `FakeInstant` in place of `Instant`.
mod token_bucket {
    include!("token_bucket.rs");

    mod tests {
        use super::*;

        pub const ZERO: Duration = Duration::new(0, 0);

        #[test]
        fn test_basic() {
            Instant::set_time(50_000);
            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 0.95, 1.0);
            assert!(bucket.get_tokens(50), "Should have not violated limit.");
            assert_eq!(None, bucket.get_delay(), "Can get stuff.");
            assert!(!bucket.get_tokens(51), "Should have violated limit.");
        }

        #[test]
        fn test_internal_constructor() {
            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 1.0, 1.0);
            assert_eq!(100, bucket.burst_limit);

            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 1e-6, 1.0);
            assert_eq!(1, bucket.burst_limit);

            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 1.0, 1e-6);
            assert_eq!(1, bucket.total_limit);
            assert_eq!(1, bucket.burst_limit);
        }

        #[test]
        fn test_saturated_100_burst() {
            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 1.00, 1.0);

            Instant::set_time(50_000);
            assert!(
                bucket.get_tokens(100),
                "All tokens should be immediately available."
            );
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(1001); // Extra buffer for Duration(0).
            assert!(
                bucket.get_tokens(100),
                "All tokens should be available after a bucket duration."
            );
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
        }

        #[test]
        fn test_saturated_95_burst() {
            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 0.95, 1.0);

            Instant::set_time(50_000);
            assert!(
                bucket.get_tokens(95),
                "95 tokens should be immediately available."
            );
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(475);
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
            Instant::advance_time(476); // Total 951. Extra buffer for Duration(0).

            assert!(bucket.get_tokens(5), "Last 5 tokens should be available.");
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(51); // Total 1002.
            assert!(bucket.get_tokens(90), "90 tokens should be available.");
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(951);
            assert!(bucket.get_tokens(10), "Last 10 tokens should be available.");
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
        }

        #[test]
        fn test_violated_50_burst() {
            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 0.50, 1.0);

            Instant::set_time(50_000);
            assert!(!bucket.get_tokens(90), "Burst should be violated.");
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
        }

        #[test]
        fn test_saturated_50_burst() {
            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 0.50, 1.0);

            Instant::set_time(50_000);
            assert!(
                bucket.get_tokens(50),
                "Half the tokens should be immediately available."
            );
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(501); // Extra buffer for Duration(0).
            assert!(
                bucket.get_tokens(50),
                "Half the tokens should be available after a half bucket duration."
            );
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(501);
            assert!(
                bucket.get_tokens(50),
                "Half the tokens should be available after a full bucket duration."
            );
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(501);
            assert!(bucket.get_tokens(50));
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
        }

        #[test]
        fn test_saturated_90_burst_rate_usage_factor_50() {
            let rate_usage_factor = 0.5;
            let bucket = VectorTokenBucket::new(
                Duration::from_millis(1000),
                100,
                ZERO,
                0.90,
                rate_usage_factor,
            );

            Instant::set_time(50_000);
            assert!(
                bucket.get_tokens(45),
                "45 tokens should be immediately available."
            );
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(475);
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
            Instant::advance_time(476); // Total 951. Extra buffer for Duration(0).

            assert!(bucket.get_tokens(5), "Last 5 tokens should be available.");
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(51); // Total 1002.
            assert!(bucket.get_tokens(40), "45 tokens should be available.");
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");

            Instant::advance_time(951);
            assert!(bucket.get_tokens(10), "Last 10 tokens should be available.");
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay.");
        }

        #[test]
        fn test_many() {
            Instant::set_time(50_000);
            let bucket = VectorTokenBucket::new(Duration::from_millis(1000), 100, ZERO, 0.5, 1.0);
            assert!(
                bucket.get_tokens(50),
                "Should have not violated limit. i=-1."
            );
            assert_ne!(None, bucket.get_delay(), "Bucket should have delay. i=-1.");
            for i in 0..20_000 {
                Instant::advance_time(501);
                assert!(
                    bucket.get_tokens(50),
                    "Should have not violated limit. i={}.",
                    i
                );
                assert_ne!(
                    None,
                    bucket.get_delay(),
                    "Bucket should have delay. i={}.",
                    i
                );
                Instant::advance_time(501);
                assert!(
                    bucket.get_tokens(50),
                    "Should have not violated limit. i={}.",
                    i
                );
                assert_ne!(
                    None,
                    bucket.get_delay(),
                    "Bucket should have delay. i={}.",
                    i
                );
            }
            assert!(
                bucket.timestamps.lock().len() < 110,
                "Should not memory leak."
            );
        }
    }
}