use std::task;
use futures::Future;
use tokio::time::{interval, sleep_until, Duration, Instant, Interval};
#[derive(Default, Debug)]
struct NeverExpire {}
impl Future for NeverExpire {
type Output = Instant;
fn poll(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
task::Poll::Pending
}
}
#[derive(Default, Debug)]
pub enum PeriodicTimer {
Started(Interval),
#[default]
Stopped,
}
impl PeriodicTimer {
pub fn started(period: Duration) -> Self {
Self::Started(interval(period))
}
pub fn stopped() -> Self {
Self::Stopped
}
pub fn start(&mut self, period: Duration) {
*self = Self::started(period);
}
pub fn stop(&mut self) {
*self = Self::stopped()
}
pub async fn tick(&mut self) -> Instant {
match self {
Self::Started(interval) => interval.tick().await,
Self::Stopped => NeverExpire::default().await,
}
}
}
#[derive(Default, Debug)]
pub enum OneshotTimer {
Scheduled(Instant),
#[default]
Expired,
}
impl OneshotTimer {
pub fn scheduled(duration: Duration) -> Self {
Self::Scheduled(Instant::now() + duration)
}
pub fn expired() -> Self {
Self::Expired
}
pub fn schedule(&mut self, duration: Duration) {
*self = Self::scheduled(duration);
}
pub fn cancel(&mut self) {
*self = Self::expired()
}
pub async fn tick(&mut self) {
match self {
Self::Scheduled(instant) => {
sleep_until(*instant).await;
*self = Self::expired();
}
Self::Expired => {
NeverExpire::default().await;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_periodic_timer() {
let mut timer1 = PeriodicTimer::stopped();
let mut timer2 = PeriodicTimer::started(Duration::from_secs(2));
let mut timer1_expired = false;
let mut timer2_expired = false;
tokio::select! {
_ = timer1.tick() => {
timer1_expired = true;
}
_ = timer2.tick() => {
timer2_expired = true;
}
}
assert!(!timer1_expired, "timer1 should not have expired");
assert!(timer2_expired, "timer1 should have expired");
timer1.start(Duration::from_secs(1));
timer2.stop();
timer1_expired = false;
timer2_expired = false;
tokio::select! {
_ = timer1.tick() => {
timer1_expired = true;
}
_ = timer2.tick() => {
timer2_expired = true;
}
}
assert!(timer1_expired, "timer1 should have expired");
assert!(!timer2_expired, "timer2 should not have expired");
}
#[tokio::test]
async fn test_oneshot_timer() {
let mut timer1 = OneshotTimer::expired();
let mut timer2 = OneshotTimer::scheduled(Duration::from_secs(2));
let mut timer1_expired = false;
let mut timer2_expired = false;
tokio::select! {
_ = timer1.tick() => {
timer1_expired = true;
}
_ = timer2.tick() => {
timer2_expired = true;
}
}
assert!(!timer1_expired, "timer1 should not have expired");
assert!(timer2_expired, "timer1 should have expired");
timer1.schedule(Duration::from_secs(1));
timer1_expired = false;
timer2_expired = false;
tokio::select! {
_ = timer1.tick() => {
timer1_expired = true;
}
_ = timer2.tick() => {
timer2_expired = true;
}
}
assert!(timer1_expired, "timer1 should have expired");
assert!(!timer2_expired, "timer2 should not have expired");
timer1.schedule(Duration::from_secs(1));
timer2.schedule(Duration::from_secs(2));
timer1.cancel();
timer1_expired = false;
timer2_expired = false;
tokio::select! {
_ = timer1.tick() => {
timer1_expired = true;
}
_ = timer2.tick() => {
timer2_expired = true;
}
}
assert!(!timer1_expired, "timer1 should not have expired");
assert!(timer2_expired, "timer2 should have expired");
}
#[tokio::test]
async fn test_oneshot_state() {
let mut timer1 = OneshotTimer::scheduled(Duration::from_secs(1));
let result = tokio::time::timeout(Duration::from_millis(1500), timer1.tick()).await;
assert!(result.is_ok(), "Should not timeout");
let mut timer1 = OneshotTimer::scheduled(Duration::from_secs(5));
let mut timer2 = OneshotTimer::scheduled(Duration::from_secs(2));
tokio::select! {
_ = timer1.tick() => {}
_ = timer2.tick() => {}
}
match timer1 {
OneshotTimer::Scheduled(_) => {}
OneshotTimer::Expired => assert!(false, "Should be in scheduled state"),
}
let result = tokio::time::timeout(Duration::from_millis(3500), timer1.tick()).await;
assert!(result.is_ok(), "Should not timeout");
match timer1 {
OneshotTimer::Scheduled(_) => assert!(false, "Timer should be in expired state"),
OneshotTimer::Expired => {}
}
}
#[tokio::test]
async fn test_my_task() {
struct MyTask {
period: PeriodicTimer,
}
impl MyTask {
fn new() -> Self {
Self {
period: PeriodicTimer::started(Duration::from_secs(1)),
}
}
fn do_work(&mut self) {}
}
let mut task = MyTask::new();
let mut sleep = OneshotTimer::scheduled(Duration::from_secs(3));
let result = tokio::time::timeout(Duration::from_secs(10), async move {
for _ in 0..3 {
tokio::select! {
_ = task.period.tick() => {
task.do_work();
task.period.stop();
}
_ = sleep.tick() => {
task.period.start(Duration::from_secs(1));
}
}
}
})
.await;
assert!(result.is_ok(), "Should not timeout");
}
}