use std::collections::BTreeMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct TimeoutCookie(pub u64);
#[derive(Default)]
pub struct LuaTimeoutManager {
registrations: BTreeMap<TimeoutCookie, i64>,
next_cookie: u64,
}
impl LuaTimeoutManager {
pub fn set_cookie(
&mut self,
cookie: TimeoutCookie,
now_monotonic_millis: i64,
timeout_millis: i64,
) {
self
.registrations
.insert(cookie, now_monotonic_millis + timeout_millis.max(0));
}
pub fn advance_timeout(&mut self, now_monotonic_millis: i64, millis: i64) {
let deadline = now_monotonic_millis + millis.max(0);
for deadline_slot in self.registrations.values_mut() {
*deadline_slot = deadline;
}
}
pub fn register_for_timeout(
&mut self,
now_monotonic_millis: i64,
timeout_millis: i64,
) -> TimeoutCookie {
self.next_cookie += 1;
let cookie = TimeoutCookie(self.next_cookie);
self.set_cookie(cookie, now_monotonic_millis, timeout_millis);
cookie
}
pub fn remove_registration(&mut self, cookie: TimeoutCookie) {
self.registrations.remove(&cookie);
}
pub fn tick_timeouts(&mut self, now_monotonic_millis: i64) -> Vec<TimeoutCookie> {
let expired: Vec<TimeoutCookie> = self
.registrations
.iter()
.filter(|(_, deadline)| **deadline <= now_monotonic_millis)
.map(|(cookie, _)| *cookie)
.collect();
for cookie in &expired {
self.registrations.remove(cookie);
}
expired
}
pub fn active_count(&self) -> usize {
self.registrations.len()
}
}
#[cfg(test)]
mod tests {
use super::{LuaTimeoutManager, TimeoutCookie};
#[test]
fn register_tick_expire() {
let mut mgr = LuaTimeoutManager::default();
let c1 = mgr.register_for_timeout(1_000, 500);
let c2 = mgr.register_for_timeout(1_000, 2_000);
assert_eq!(mgr.active_count(), 2);
assert!(mgr.tick_timeouts(1_400).is_empty());
assert_eq!(mgr.tick_timeouts(1_500), vec![c1]);
assert_eq!(mgr.active_count(), 1);
assert_eq!(mgr.tick_timeouts(3_000), vec![c2]);
}
#[test]
fn set_cookie_and_advance() {
let mut mgr = LuaTimeoutManager::default();
let cookie = TimeoutCookie(7);
mgr.set_cookie(cookie, 1_000, 1_000);
mgr.advance_timeout(1_000, 100);
assert_eq!(mgr.tick_timeouts(1_100), vec![cookie]);
mgr.remove_registration(cookie);
assert_eq!(mgr.active_count(), 0);
}
}