use std::ops::Deref;
use crate::{
clock::{Clock, InstantClock},
GcraError, GcraState, RateLimit,
};
pub struct RateLimitGuard<C: Clock = InstantClock> {
clock: C,
rate_limit: RateLimit,
state: GcraState,
}
impl RateLimitGuard {
pub fn new_state(rate_limit: RateLimit) -> Self {
RateLimitGuard {
clock: InstantClock,
rate_limit,
state: GcraState::default(),
}
}
}
impl<C: Clock> RateLimitGuard<C> {
pub fn new(clock: C, rate_limit: RateLimit, state: GcraState) -> Self {
RateLimitGuard {
clock,
rate_limit,
state,
}
}
pub fn check_and_modify(&mut self, cost: u32) -> Result<(), GcraError> {
let RateLimitGuard {
clock,
rate_limit,
state,
} = self;
let arrived_at = clock.now();
state.check_and_modify_at(rate_limit, arrived_at, cost)
}
pub fn remaining_resources(&self) -> u32 {
self.state
.remaining_resources(&self.rate_limit, self.clock.now())
}
pub fn revert(&mut self, cost: u32) -> Result<(), GcraError> {
let RateLimitGuard {
clock,
rate_limit,
state,
} = self;
let arrived_at = clock.now();
state.revert_at(rate_limit, arrived_at, cost)
}
}
impl<C: Clock> Deref for RateLimitGuard<C> {
type Target = GcraState;
fn deref(&self) -> &Self::Target {
&self.state
}
}