use chrono::{Datelike, Local, Timelike};
use parking_lot::RwLock;
pub use crate::protocol::{BandwidthLimits, ScheduleRule};
pub struct BandwidthScheduler {
rules: Vec<ScheduleRule>,
default_limits: BandwidthLimits,
current_limits: RwLock<BandwidthLimits>,
}
impl BandwidthScheduler {
pub fn new(rules: Vec<ScheduleRule>, default_limits: BandwidthLimits) -> Self {
let current = Self::evaluate_rules(&rules, &default_limits);
Self {
rules,
default_limits,
current_limits: RwLock::new(current),
}
}
pub fn with_defaults(download: Option<u64>, upload: Option<u64>) -> Self {
Self::new(Vec::new(), BandwidthLimits { download, upload })
}
pub fn get_limits(&self) -> BandwidthLimits {
*self.current_limits.read()
}
pub fn update(&self) -> bool {
let new_limits = Self::evaluate_rules(&self.rules, &self.default_limits);
let mut current = self.current_limits.write();
if *current != new_limits {
tracing::info!(
"Bandwidth limits changed: download={:?} upload={:?}",
new_limits.download,
new_limits.upload
);
*current = new_limits;
true
} else {
false
}
}
fn evaluate_rules(rules: &[ScheduleRule], default: &BandwidthLimits) -> BandwidthLimits {
let now = Local::now();
let hour = now.hour() as u8;
let weekday = now.weekday();
for rule in rules {
if rule.matches(hour, weekday) {
return BandwidthLimits {
download: rule.download_limit,
upload: rule.upload_limit,
};
}
}
*default
}
pub fn rules(&self) -> &[ScheduleRule] {
&self.rules
}
pub fn add_rule(&mut self, rule: ScheduleRule) {
self.rules.push(rule);
self.update();
}
pub fn clear_rules(&mut self) {
self.rules.clear();
self.update();
}
pub fn set_rules(&mut self, rules: Vec<ScheduleRule>) {
self.rules = rules;
self.update();
}
pub fn set_defaults(&mut self, limits: BandwidthLimits) {
self.default_limits = limits;
self.update();
}
}
impl Default for BandwidthScheduler {
fn default() -> Self {
Self::with_defaults(None, None)
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Weekday;
#[test]
fn test_schedule_rule_match_simple() {
let rule = ScheduleRule::all_days(9, 17, Some(1_000_000), None);
assert!(rule.matches(9, Weekday::Mon));
assert!(rule.matches(12, Weekday::Wed));
assert!(rule.matches(17, Weekday::Fri));
assert!(!rule.matches(8, Weekday::Mon));
assert!(!rule.matches(18, Weekday::Wed));
assert!(!rule.matches(0, Weekday::Fri));
}
#[test]
fn test_schedule_rule_match_overnight() {
let rule = ScheduleRule::all_days(22, 6, Some(10_000_000), None);
assert!(rule.matches(22, Weekday::Mon));
assert!(rule.matches(23, Weekday::Mon));
assert!(rule.matches(0, Weekday::Tue));
assert!(rule.matches(3, Weekday::Tue));
assert!(rule.matches(6, Weekday::Tue));
assert!(!rule.matches(7, Weekday::Mon));
assert!(!rule.matches(12, Weekday::Mon));
assert!(!rule.matches(21, Weekday::Mon));
}
#[test]
fn test_schedule_rule_weekdays() {
let rule = ScheduleRule::weekdays(9, 17, Some(500_000), None);
assert!(rule.matches(12, Weekday::Mon));
assert!(rule.matches(12, Weekday::Tue));
assert!(rule.matches(12, Weekday::Wed));
assert!(rule.matches(12, Weekday::Thu));
assert!(rule.matches(12, Weekday::Fri));
assert!(!rule.matches(12, Weekday::Sat));
assert!(!rule.matches(12, Weekday::Sun));
}
#[test]
fn test_schedule_rule_weekends() {
let rule = ScheduleRule::weekends(0, 23, None, None);
assert!(rule.matches(12, Weekday::Sat));
assert!(rule.matches(12, Weekday::Sun));
assert!(!rule.matches(12, Weekday::Mon));
assert!(!rule.matches(12, Weekday::Fri));
}
#[test]
fn test_scheduler_no_rules() {
let scheduler = BandwidthScheduler::with_defaults(Some(1_000_000), Some(500_000));
let limits = scheduler.get_limits();
assert_eq!(limits.download, Some(1_000_000));
assert_eq!(limits.upload, Some(500_000));
}
#[test]
fn test_scheduler_first_match_wins() {
let rules = vec![
ScheduleRule::all_days(0, 23, Some(1_000_000), None),
ScheduleRule::all_days(0, 23, Some(2_000_000), None),
];
let scheduler = BandwidthScheduler::new(rules, BandwidthLimits::default());
let limits = scheduler.get_limits();
assert_eq!(limits.download, Some(1_000_000));
}
#[test]
fn test_bandwidth_limits_default() {
let limits = BandwidthLimits::default();
assert_eq!(limits.download, None);
assert_eq!(limits.upload, None);
}
}