use std::fmt;
use std::sync::Arc;
use chrono::{DateTime, Datelike, Duration, Local, Timelike, Utc, Weekday};
use crate::{Error, LimitExceeded, Result, TimeUnit, TimeWindow};
use super::inner::EventStoreInner;
use super::query::Query;
use super::EventId;
pub struct Limiter {
store: Arc<EventStoreInner>,
constraints: Vec<Constraint>,
}
#[derive(Debug, Clone)]
pub enum Constraint {
AtMost {
event_id: String,
limit: u32,
window_count: usize,
time_unit: TimeUnit,
},
AtLeast {
event_id: String,
count: u32,
window_count: usize,
time_unit: TimeUnit,
},
Cooldown {
event_id: String,
duration: Duration,
},
Within {
prerequisite_event: String,
duration: Duration,
},
During { schedule: Schedule },
OutsideOf { schedule: Schedule },
}
#[derive(Clone)]
pub enum Schedule {
Hours {
start_hour: u8,
end_hour: u8,
use_local_tz: bool,
},
Weekdays { use_local_tz: bool },
Weekends { use_local_tz: bool },
Custom(Arc<dyn Fn(DateTime<Utc>) -> bool + Send + Sync>),
}
impl fmt::Debug for Schedule {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Schedule::Hours {
start_hour,
end_hour,
use_local_tz,
} => {
let tz_str = if *use_local_tz { "Local" } else { "UTC" };
write!(
f,
"Hours {{ start_hour: {}, end_hour: {}, timezone: {} }}",
start_hour, end_hour, tz_str
)
}
Schedule::Weekdays { use_local_tz } => {
let tz_str = if *use_local_tz { "Local" } else { "UTC" };
write!(f, "Weekdays({})", tz_str)
}
Schedule::Weekends { use_local_tz } => {
let tz_str = if *use_local_tz { "Local" } else { "UTC" };
write!(f, "Weekends({})", tz_str)
}
Schedule::Custom(_) => write!(f, "Custom(<closure>)"),
}
}
}
impl Schedule {
pub fn hours(start_hour: u8, end_hour: u8) -> Result<Self> {
Self::validate_hours(start_hour, end_hour)?;
Ok(Schedule::Hours {
start_hour,
end_hour,
use_local_tz: false,
})
}
pub fn hours_local_tz(start_hour: u8, end_hour: u8) -> Result<Self> {
Self::validate_hours(start_hour, end_hour)?;
Ok(Schedule::Hours {
start_hour,
end_hour,
use_local_tz: true,
})
}
pub fn weekdays() -> Self {
Schedule::Weekdays {
use_local_tz: false,
}
}
pub fn weekdays_local_tz() -> Self {
Schedule::Weekdays { use_local_tz: true }
}
pub fn weekends() -> Self {
Schedule::Weekends {
use_local_tz: false,
}
}
pub fn weekends_local_tz() -> Self {
Schedule::Weekends { use_local_tz: true }
}
fn validate_hours(start_hour: u8, end_hour: u8) -> Result<()> {
if start_hour > 23 {
Err(Error::InvalidHour(start_hour))
} else if end_hour > 23 {
Err(Error::InvalidHour(end_hour))
} else if start_hour >= end_hour {
Err(Error::InvalidRange(start_hour, end_hour))
} else {
Ok(())
}
}
}
pub struct LimitUsage {
pub count: u32,
pub limit: u32,
pub remaining: u32,
}
pub struct Reservation {
store: Arc<EventStoreInner>,
event_id: String,
committed: bool,
}
impl Reservation {
pub fn commit(mut self) {
self.release_and_record();
self.committed = true;
}
pub fn cancel(mut self) {
self.release_pending();
self.committed = true;
}
fn release_and_record(&self) {
if let Some(counter_arc) = self.store.events.get(&self.event_id) {
let mut counter = counter_arc.lock().unwrap();
counter.decrement_pending();
let now = self.store.clock_now();
counter.advance_if_needed(now);
counter.record(1);
counter.mark_dirty();
}
}
fn release_pending(&self) {
if let Some(counter_arc) = self.store.events.get(&self.event_id) {
let mut counter = counter_arc.lock().unwrap();
counter.decrement_pending();
counter.mark_dirty();
}
}
}
impl Drop for Reservation {
fn drop(&mut self) {
if !self.committed {
self.release_pending();
}
}
}
impl Limiter {
pub(crate) fn new(store: Arc<EventStoreInner>) -> Self {
Self {
store,
constraints: Vec::new(),
}
}
pub fn at_most(
mut self,
event_id: impl EventId,
limit: u32,
window: impl Into<TimeWindow>,
) -> Self {
let window = window.into();
self.constraints.push(Constraint::AtMost {
event_id: event_id.as_ref().to_string(),
limit,
window_count: window.count,
time_unit: window.time_unit,
});
self
}
pub fn at_least(
mut self,
event_id: impl EventId,
count: u32,
window: impl Into<TimeWindow>,
) -> Self {
let window = window.into();
self.constraints.push(Constraint::AtLeast {
event_id: event_id.as_ref().to_string(),
count,
window_count: window.count,
time_unit: window.time_unit,
});
self
}
pub fn cooldown(mut self, event_id: impl EventId, duration: Duration) -> Self {
self.constraints.push(Constraint::Cooldown {
event_id: event_id.as_ref().to_string(),
duration,
});
self
}
pub fn within(mut self, prerequisite_event: impl EventId, duration: Duration) -> Self {
self.constraints.push(Constraint::Within {
prerequisite_event: prerequisite_event.as_ref().to_string(),
duration,
});
self
}
pub fn during(mut self, schedule: Schedule) -> Self {
self.constraints.push(Constraint::During { schedule });
self
}
pub fn outside_of(mut self, schedule: Schedule) -> Self {
self.constraints.push(Constraint::OutsideOf { schedule });
self
}
pub fn check(self, event_id: impl EventId) -> Result<()> {
let event_id_str = event_id.as_ref();
self.evaluate_constraints(event_id_str)
}
pub fn check_and_record(self, event_id: impl EventId) -> Result<()> {
let event_id_str = event_id.as_ref();
self.evaluate_constraints(event_id_str)?;
let now = self.store.clock.now();
let counter = self.store.get_counter_for_record(event_id_str);
let mut counter = counter.lock().unwrap();
counter.advance_if_needed(now);
counter.record(1);
counter.mark_dirty();
Ok(())
}
pub fn reserve(self, event_id: impl EventId) -> Result<Reservation> {
let clock_now = self.store.clock_now();
let event_id_str = event_id.as_ref();
let counter_arc = self.store.get_counter_for_record(event_id_str);
let mut counter = counter_arc.lock().unwrap();
counter.advance_if_needed(clock_now);
for constraint in &self.constraints {
match constraint {
Constraint::AtMost {
event_id: _constraint_event,
limit,
window_count,
time_unit,
} => {
let total = counter.total_with_pending(*time_unit, 0..*window_count);
if total >= *limit {
return Err(Error::LimitExceeded(LimitExceeded {
event_id: event_id_str.to_string(),
constraint: constraint.clone(),
retry_after: Some(time_unit.duration()),
}));
}
}
_ => {
drop(counter);
self.evaluate_constraint(event_id_str, constraint)?;
counter = counter_arc.lock().unwrap();
}
}
}
counter.increment_pending();
counter.mark_dirty();
Ok(Reservation {
store: self.store.clone(),
event_id: event_id_str.to_string(),
committed: false,
})
}
pub fn allowed(self, event_id: impl EventId) -> bool {
self.check(event_id).is_ok()
}
pub fn usage(self, event_id: impl EventId) -> Result<LimitUsage> {
let event_id_str = event_id.as_ref();
for constraint in &self.constraints {
if let Constraint::AtMost {
event_id: constraint_event,
limit,
window_count,
time_unit,
} = constraint
{
if constraint_event == event_id_str {
let current = self
.query_sum_window(constraint_event, *window_count, *time_unit)
.unwrap_or(0);
let remaining = limit.saturating_sub(current);
return Ok(LimitUsage {
count: current,
limit: *limit,
remaining,
});
}
}
}
Ok(LimitUsage {
count: 0,
limit: 0,
remaining: 0,
})
}
fn evaluate_constraints(&self, event_id: &str) -> Result<()> {
for constraint in &self.constraints {
self.evaluate_constraint(event_id, constraint)?;
}
Ok(())
}
fn evaluate_constraint(&self, event_id: &str, constraint: &Constraint) -> Result<()> {
match constraint {
Constraint::AtMost {
event_id: constraint_event,
limit,
window_count,
time_unit,
} => {
let current = self
.query_sum_window(constraint_event, *window_count, *time_unit)
.unwrap_or(0);
if current >= *limit {
return Err(Error::LimitExceeded(LimitExceeded {
event_id: event_id.to_string(),
constraint: constraint.clone(),
retry_after: Some(time_unit.duration()),
}));
}
}
Constraint::AtLeast {
event_id: constraint_event,
count,
window_count,
time_unit,
} => {
let current = self
.query_sum_window(constraint_event, *window_count, *time_unit)
.unwrap_or(0);
if current < *count {
return Err(Error::LimitExceeded(LimitExceeded {
event_id: event_id.to_string(),
constraint: constraint.clone(),
retry_after: None,
}));
}
}
Constraint::Cooldown {
event_id: constraint_event,
duration,
} => {
if let Some(elapsed) = self.query_last_seen(constraint_event) {
if elapsed < *duration {
let remaining = *duration - elapsed;
return Err(Error::LimitExceeded(LimitExceeded {
event_id: event_id.to_string(),
constraint: constraint.clone(),
retry_after: Some(remaining),
}));
}
}
}
Constraint::Within {
prerequisite_event,
duration,
} => {
match self.query_last_seen(prerequisite_event) {
None => {
return Err(Error::LimitExceeded(LimitExceeded {
event_id: event_id.to_string(),
constraint: constraint.clone(),
retry_after: None,
}));
}
Some(elapsed) => {
if elapsed > *duration {
return Err(Error::LimitExceeded(LimitExceeded {
event_id: event_id.to_string(),
constraint: constraint.clone(),
retry_after: None,
}));
}
}
}
}
Constraint::During { schedule } => {
if !schedule.allows(&self.store) {
return Err(Error::LimitExceeded(LimitExceeded {
event_id: event_id.to_string(),
constraint: constraint.clone(),
retry_after: schedule.next_allowed(&self.store),
}));
}
}
Constraint::OutsideOf { schedule } => {
if schedule.allows(&self.store) {
return Err(Error::LimitExceeded(LimitExceeded {
event_id: event_id.to_string(),
constraint: constraint.clone(),
retry_after: schedule.block_end(&self.store),
}));
}
}
}
Ok(())
}
fn query_sum_window(
&self,
event_id: &str,
window_count: usize,
time_unit: TimeUnit,
) -> Option<u32> {
let query = Query::new(self.store.clone(), event_id.to_string());
query.last(window_count, time_unit).sum()
}
fn query_last_seen(&self, event_id: &str) -> Option<Duration> {
Query::new(self.store.clone(), event_id.to_string()).last_seen()
}
}
impl Schedule {
fn allows(&self, store: &EventStoreInner) -> bool {
let now = store.clock_now();
match self {
Schedule::Hours {
start_hour,
end_hour,
use_local_tz,
} => {
let hour = if *use_local_tz {
now.with_timezone(&Local).hour() as u8
} else {
now.hour() as u8
};
hour >= *start_hour && hour < *end_hour
}
Schedule::Weekdays { use_local_tz } => {
let weekday = if *use_local_tz {
now.with_timezone(&Local).weekday()
} else {
now.weekday()
};
!matches!(weekday, Weekday::Sat | Weekday::Sun)
}
Schedule::Weekends { use_local_tz } => {
let weekday = if *use_local_tz {
now.with_timezone(&Local).weekday()
} else {
now.weekday()
};
matches!(weekday, Weekday::Sat | Weekday::Sun)
}
Schedule::Custom(f) => f(now),
}
}
fn next_allowed(&self, store: &EventStoreInner) -> Option<Duration> {
let now = store.clock_now();
match self {
Schedule::Hours {
start_hour,
end_hour: _,
use_local_tz,
} => {
let current_hour = if *use_local_tz {
now.with_timezone(&Local).hour() as u8
} else {
now.hour() as u8
};
if current_hour < *start_hour {
let hours_until = (*start_hour - current_hour) as i64;
Some(Duration::hours(hours_until))
} else {
let hours_until = 24 - current_hour + *start_hour;
Some(Duration::hours(hours_until as i64))
}
}
Schedule::Weekdays { use_local_tz } => {
let weekday = if *use_local_tz {
now.with_timezone(&Local).weekday()
} else {
now.weekday()
};
match weekday {
Weekday::Sat => Some(Duration::days(2)),
Weekday::Sun => Some(Duration::days(1)),
_ => None,
}
}
Schedule::Weekends { use_local_tz } => {
let weekday = if *use_local_tz {
now.with_timezone(&Local).weekday()
} else {
now.weekday()
};
let days_until_sat = match weekday {
Weekday::Mon => 5,
Weekday::Tue => 4,
Weekday::Wed => 3,
Weekday::Thu => 2,
Weekday::Fri => 1,
_ => 0,
};
if days_until_sat > 0 {
Some(Duration::days(days_until_sat))
} else {
None
}
}
Schedule::Custom(_) => None, }
}
fn block_end(&self, store: &EventStoreInner) -> Option<Duration> {
let now = store.clock_now();
match self {
Schedule::Hours {
start_hour: _,
end_hour,
use_local_tz,
} => {
let current_hour = if *use_local_tz {
now.with_timezone(&Local).hour() as u8
} else {
now.hour() as u8
};
if current_hour < *end_hour {
let hours_until = (*end_hour - current_hour) as i64;
Some(Duration::hours(hours_until))
} else {
None
}
}
Schedule::Weekdays { use_local_tz } => {
let weekday = if *use_local_tz {
now.with_timezone(&Local).weekday()
} else {
now.weekday()
};
match weekday {
Weekday::Mon => Some(Duration::days(5)),
Weekday::Tue => Some(Duration::days(4)),
Weekday::Wed => Some(Duration::days(3)),
Weekday::Thu => Some(Duration::days(2)),
Weekday::Fri => Some(Duration::days(1)),
_ => None,
}
}
Schedule::Weekends { use_local_tz } => {
let weekday = if *use_local_tz {
now.with_timezone(&Local).weekday()
} else {
now.weekday()
};
match weekday {
Weekday::Sat => Some(Duration::days(2)),
Weekday::Sun => Some(Duration::days(1)),
_ => None,
}
}
Schedule::Custom(_) => None, }
}
}
#[cfg(test)]
mod tests {
use std::thread;
use chrono::{Duration, TimeZone, Weekday};
use crate::{EventStore, TestClock, TimeUnit};
use super::*;
#[test]
fn test_at_most_constraint_passes() {
let store = EventStore::new();
store.record_count("test_event", 5);
let result = store
.limit()
.at_most("test_event", 10, TimeUnit::Days)
.check("action");
assert!(result.is_ok());
}
#[test]
fn test_at_most_constraint_fails() {
let store = EventStore::new();
store.record_count("test_event", 10);
let result = store
.limit()
.at_most("test_event", 5, TimeUnit::Days)
.check("action");
assert!(result.is_err());
let Error::LimitExceeded(err) = result.unwrap_err() else {
panic!("Expected LimitExceeded error");
};
assert_eq!(err.event_id, "action");
assert!(err.retry_after.is_some());
}
#[test]
fn test_at_least_constraint_passes() {
let store = EventStore::new();
store.record_count("prerequisite", 5);
let result = store
.limit()
.at_least("prerequisite", 3, TimeUnit::Days)
.check("action");
assert!(result.is_ok());
}
#[test]
fn test_at_least_constraint_fails() {
let store = EventStore::new();
store.record_count("prerequisite", 2);
let result = store
.limit()
.at_least("prerequisite", 5, TimeUnit::Days)
.check("action");
assert!(result.is_err());
let Error::LimitExceeded(err) = result.unwrap_err() else {
panic!("Expected LimitExceeded error");
};
assert!(err.retry_after.is_none());
}
#[test]
fn test_cooldown_constraint_passes() {
let store = EventStore::new();
store.record_ago("test_event", Duration::hours(2));
let result = store
.limit()
.cooldown("test_event", Duration::hours(1))
.check("action");
assert!(result.is_ok());
}
#[test]
fn test_cooldown_constraint_fails() {
let store = EventStore::new();
store.record("test_event");
let result = store
.limit()
.cooldown("test_event", Duration::hours(1))
.check("action");
assert!(result.is_err());
let Error::LimitExceeded(err) = result.unwrap_err() else {
panic!("Expected LimitExceeded error");
};
assert!(err.retry_after.is_some());
}
#[test]
fn test_within_constraint_passes() {
let store = EventStore::new();
store.record("prerequisite");
let result = store
.limit()
.within("prerequisite", Duration::hours(1))
.check("action");
assert!(result.is_ok());
}
#[test]
fn test_within_constraint_fails_never_seen() {
let store = EventStore::new();
let result = store
.limit()
.within("prerequisite", Duration::hours(1))
.check("action");
assert!(result.is_err());
}
#[test]
fn test_within_constraint_fails_too_long_ago() {
let store = EventStore::new();
store.record_ago("prerequisite", Duration::hours(2));
let result = store
.limit()
.within("prerequisite", Duration::hours(1))
.check("action");
assert!(result.is_err());
}
#[test]
fn test_multiple_constraints_all_pass() {
let store = EventStore::new();
store.record_count("prerequisite", 5);
store.record_count("test_event", 2);
let result = store
.limit()
.at_least("prerequisite", 3, TimeUnit::Days)
.at_most("test_event", 10, TimeUnit::Days)
.check("action");
assert!(result.is_ok());
}
#[test]
fn test_multiple_constraints_first_fails() {
let store = EventStore::new();
store.record_count("prerequisite", 1);
store.record_count("test_event", 2);
let result = store
.limit()
.at_least("prerequisite", 3, TimeUnit::Days)
.at_most("test_event", 10, TimeUnit::Days)
.check("action");
assert!(result.is_err());
}
#[test]
fn test_check_and_record_success() {
let store = EventStore::new();
let result = store
.limit()
.at_most("test_event", 5, TimeUnit::Days)
.check_and_record("test_event");
assert!(result.is_ok());
let count = store.query("test_event").last_days(1).sum();
assert_eq!(count, Some(1));
}
#[test]
fn test_check_and_record_failure_does_not_record() {
let store = EventStore::new();
store.record_count("test_event", 10);
let result = store
.limit()
.at_most("test_event", 5, TimeUnit::Days)
.check_and_record("test_event");
assert!(result.is_err());
let count = store.query("test_event").last_days(1).sum();
assert_eq!(count, Some(10));
}
#[test]
fn test_allowed_convenience_method() {
let store = EventStore::new();
store.record_count("test_event", 3);
let allowed = store
.limit()
.at_most("test_event", 5, TimeUnit::Days)
.allowed("action");
assert!(allowed);
store.record_count("test_event", 2);
let not_allowed = store
.limit()
.at_most("test_event", 5, TimeUnit::Days)
.allowed("action");
assert!(!not_allowed);
}
#[test]
fn test_usage_returns_correct_counts() {
let store = EventStore::new();
store.record_count("test_event", 3);
let usage = store
.limit()
.at_most("test_event", 10, TimeUnit::Days)
.usage("test_event")
.unwrap();
assert_eq!(usage.count, 3);
assert_eq!(usage.limit, 10);
assert_eq!(usage.remaining, 7);
}
#[test]
fn test_business_hours_schedule_during_hours() {
let tuesday_10am = Utc.with_ymd_and_hms(2025, 1, 7, 10, 0, 0).unwrap();
let clock = TestClock::build_for_testing_at(tuesday_10am);
let store = EventStore::builder()
.with_clock(Arc::new(clock.clone()))
.build()
.unwrap();
let result = store
.limit()
.during(Schedule::hours(9, 17).unwrap())
.check_and_record("action");
assert!(result.is_ok());
let tuesday_7pm = Utc.with_ymd_and_hms(2025, 1, 7, 19, 0, 0).unwrap();
clock.set(tuesday_7pm);
let result = store
.limit()
.during(Schedule::hours(9, 17).unwrap())
.check("action");
assert!(result.is_err());
}
#[test]
fn test_weekdays_schedule() {
let tuesday = Utc.with_ymd_and_hms(2025, 1, 7, 12, 0, 0).unwrap();
assert_eq!(tuesday.weekday(), Weekday::Tue);
let clock = TestClock::build_for_testing_at(tuesday);
let store = EventStore::builder()
.with_clock(Arc::new(clock.clone()))
.build()
.unwrap();
let result = store
.limit()
.during(Schedule::weekdays())
.check_and_record("action");
assert!(result.is_ok());
let saturday = Utc.with_ymd_and_hms(2025, 1, 11, 12, 0, 0).unwrap();
assert_eq!(saturday.weekday(), Weekday::Sat);
clock.set(saturday);
let result = store.limit().during(Schedule::weekdays()).check("action");
assert!(result.is_err());
}
#[test]
fn test_custom_schedule() {
let store = EventStore::new();
let schedule = Schedule::Custom(Arc::new(|_time| true));
let result = store.limit().during(schedule).check("action");
assert!(result.is_ok());
}
#[test]
fn test_schedule_hours_constructor() {
let schedule = Schedule::hours(9, 17).unwrap();
match schedule {
Schedule::Hours {
start_hour,
end_hour,
..
} => {
assert_eq!(start_hour, 9);
assert_eq!(end_hour, 17);
}
_ => panic!("Expected Hours variant"),
}
}
#[test]
fn test_schedule_hours_validates_start_hour() {
let result = Schedule::hours(24, 17);
assert!(result.is_err());
match result.unwrap_err() {
Error::InvalidHour(hour) => assert_eq!(hour, 24),
_ => panic!("Expected InvalidHour error"),
}
}
#[test]
fn test_schedule_hours_validates_end_hour() {
let result = Schedule::hours(9, 25);
assert!(result.is_err());
match result.unwrap_err() {
Error::InvalidHour(hour) => assert_eq!(hour, 25),
_ => panic!("Expected InvalidHour error"),
}
}
#[test]
fn test_schedule_hours_validates_range() {
let result = Schedule::hours(17, 9);
assert!(result.is_err());
match result.unwrap_err() {
Error::InvalidRange(start, end) => {
assert_eq!(start, 17);
assert_eq!(end, 9);
}
_ => panic!("Expected InvalidRange error"),
}
}
#[test]
fn test_schedule_hours_rejects_equal_hours() {
let result = Schedule::hours(9, 9);
assert!(result.is_err());
match result.unwrap_err() {
Error::InvalidRange(start, end) => {
assert_eq!(start, 9);
assert_eq!(end, 9);
}
_ => panic!("Expected InvalidRange error"),
}
}
#[test]
fn test_schedule_hours_accepts_valid_boundary_cases() {
let schedule = Schedule::hours(0, 23).unwrap();
match schedule {
Schedule::Hours {
start_hour,
end_hour,
..
} => {
assert_eq!(start_hour, 0);
assert_eq!(end_hour, 23);
}
_ => panic!("Expected Hours variant"),
}
let schedule = Schedule::hours(0, 1).unwrap();
match schedule {
Schedule::Hours {
start_hour,
end_hour,
..
} => {
assert_eq!(start_hour, 0);
assert_eq!(end_hour, 1);
}
_ => panic!("Expected Hours variant"),
}
}
#[test]
fn test_retry_after_calculation_for_at_most() {
let store = EventStore::new();
store.record_count("test_event", 10);
let result = store
.limit()
.at_most("test_event", 5, TimeUnit::Hours)
.check("action");
assert!(result.is_err());
let Error::LimitExceeded(err) = result.unwrap_err() else {
panic!("Expected LimitExceeded error");
};
assert_eq!(err.retry_after, Some(Duration::hours(1)));
}
#[test]
fn test_cooldown_never_seen_passes() {
let store = EventStore::new();
let result = store
.limit()
.cooldown("never_seen", Duration::hours(1))
.check("action");
assert!(result.is_ok());
}
#[test]
fn test_flexible_time_windows_with_tuple() {
let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
let clock = TestClock::build_for_testing_at(start_time);
let store = EventStore::builder()
.with_clock(Arc::new(clock.clone()))
.build()
.unwrap();
for _i in 0..10 {
store.record("api");
clock.advance(Duration::hours(7));
}
let result = store
.limit()
.at_most("api", 11, (7, TimeUnit::Days))
.check("api");
assert!(result.is_ok());
let result = store
.limit()
.at_most("api", 10, (7, TimeUnit::Days))
.check("api");
assert!(result.is_err());
}
#[test]
fn test_flexible_time_windows_with_duration() {
let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
let clock = TestClock::build_for_testing_at(start_time);
let store = EventStore::builder()
.with_clock(Arc::new(clock.clone()))
.build()
.unwrap();
for _i in 0..10 {
store.record("api");
clock.advance(Duration::hours(7));
}
let result = store
.limit()
.at_most("api", 11, Duration::days(7))
.check("api");
assert!(result.is_ok());
let result = store
.limit()
.at_most("api", 10, Duration::days(7))
.check("api");
assert!(result.is_err());
}
#[test]
fn test_flexible_time_windows_backward_compatible() {
let store = EventStore::new();
store.record_count("api", 5);
let result = store
.limit()
.at_most("api", 10, TimeUnit::Days)
.check("api");
assert!(result.is_ok());
}
#[test]
fn test_flexible_time_windows_30_minutes() {
let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
let clock = TestClock::build_for_testing_at(start_time);
let store = EventStore::builder()
.with_clock(Arc::new(clock.clone()))
.build()
.unwrap();
for _i in 0..10 {
store.record("api");
clock.advance(Duration::minutes(2));
}
let result = store
.limit()
.at_most("api", 11, (30, TimeUnit::Minutes))
.check("api");
assert!(result.is_ok());
let result = store
.limit()
.at_most("api", 10, (30, TimeUnit::Minutes))
.check("api");
assert!(result.is_err());
}
#[test]
fn test_at_least_with_flexible_windows() {
let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
let clock = TestClock::build_for_testing_at(start_time);
let store = EventStore::builder()
.with_clock(Arc::new(clock.clone()))
.build()
.unwrap();
for _i in 0..10 {
store.record("prerequisite");
clock.advance(Duration::hours(7));
}
let result = store
.limit()
.at_least("prerequisite", 5, (7, TimeUnit::Days))
.check("action");
assert!(result.is_ok());
let result = store
.limit()
.at_least("prerequisite", 15, (7, TimeUnit::Days))
.check("action");
assert!(result.is_err());
}
#[test]
fn test_duration_conversion_to_time_window() {
let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
let clock = TestClock::build_for_testing_at(start_time);
let store = EventStore::builder()
.with_clock(Arc::new(clock.clone()))
.build()
.unwrap();
for _i in 0..5 {
store.record("api");
clock.advance(Duration::hours(1));
}
let result = store
.limit()
.at_most("api", 10, Duration::hours(12))
.check("api");
assert!(result.is_ok());
store.record_count("fast", 20);
let result = store
.limit()
.at_most("fast", 25, Duration::minutes(30))
.check("fast");
assert!(result.is_ok());
}
#[test]
fn test_reservation_prevents_race_condition() {
let store = EventStore::new();
let res1 = store
.limit()
.at_most("api", 1, TimeUnit::Hours)
.reserve("api")
.unwrap();
let res2 = store
.limit()
.at_most("api", 1, TimeUnit::Hours)
.reserve("api");
assert!(res2.is_err());
res1.commit();
assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
}
#[test]
fn test_reservation_auto_cancel_on_drop() {
let store = EventStore::new();
{
let _res = store
.limit()
.at_most("api", 10, TimeUnit::Hours)
.reserve("api")
.unwrap();
}
let sum = store.query("api").last_hours(1).sum();
assert!(sum.is_none() || sum == Some(0));
let res = store
.limit()
.at_most("api", 1, TimeUnit::Hours)
.reserve("api");
assert!(res.is_ok());
}
#[test]
fn test_reservation_explicit_cancel() {
let store = EventStore::new();
let res = store
.limit()
.at_most("api", 10, TimeUnit::Hours)
.reserve("api")
.unwrap();
res.cancel();
let sum = store.query("api").last_hours(1).sum();
assert!(sum.is_none() || sum == Some(0));
let res = store
.limit()
.at_most("api", 1, TimeUnit::Hours)
.reserve("api");
assert!(res.is_ok());
}
#[test]
fn test_reservation_commit_records_event() {
let store = EventStore::new();
let res = store
.limit()
.at_most("api", 10, TimeUnit::Hours)
.reserve("api")
.unwrap();
res.commit();
assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
}
#[test]
fn test_reservation_concurrent_limits() {
let store = Arc::new(EventStore::new());
let handles: Vec<_> = (0..20)
.map(|_| {
let store = store.clone();
thread::spawn(move || {
store
.limit()
.at_most("api", 10, TimeUnit::Hours)
.reserve("api")
})
})
.collect();
let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();
let successes: Vec<_> = results.into_iter().filter(|r| r.is_ok()).collect();
assert_eq!(successes.len(), 10);
for result in successes {
result.unwrap().commit();
}
assert_eq!(store.query("api").last_hours(1).sum(), Some(10));
}
#[test]
fn test_reservation_transactional_pattern() {
fn simulated_work(should_succeed: bool) -> std::result::Result<(), &'static str> {
if should_succeed {
Ok(())
} else {
Err("Work failed")
}
}
let store = EventStore::new();
let res = store
.limit()
.at_most("api", 10, TimeUnit::Hours)
.reserve("api")
.unwrap();
match simulated_work(true) {
Ok(_) => res.commit(),
Err(_) => res.cancel(),
}
assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
let res = store
.limit()
.at_most("api", 10, TimeUnit::Hours)
.reserve("api")
.unwrap();
match simulated_work(false) {
Ok(_) => res.commit(),
Err(_) => res.cancel(),
}
assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
}
#[test]
fn test_reservation_with_multiple_constraints() {
let store = EventStore::new();
store.record_count("prerequisite", 5);
let res = store
.limit()
.at_most("api", 10, TimeUnit::Hours)
.at_least("prerequisite", 3, TimeUnit::Days)
.reserve("api")
.unwrap();
res.commit();
assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
}
}