foyer-bench 0.1.4

//  Copyright 2024 Foyer Project Authors
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

// Copyright 2023 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::time::{Duration, Instant};

pub struct RateLimiter {
    capacity: f64,
    quota: f64,

    last: Instant,
}

impl RateLimiter {
    pub fn new(capacity: f64) -> Self {
        Self {
            capacity,
            quota: 0.0,
            last: Instant::now(),
        }
    }

    pub fn consume(&mut self, weight: f64) -> Option<Duration> {
        let now = Instant::now();
        let refill = now.duration_since(self.last).as_secs_f64() * self.capacity;
        self.last = now;
        self.quota = f64::min(self.quota + refill, self.capacity);
        self.quota -= weight;
        if self.quota >= 0.0 {
            return None;
        }
        let wait = Duration::from_secs_f64((-self.quota) / self.capacity);
        Some(wait)
    }
}