system_timer/

lib.rs

//! Async timers built on [`SystemTime`].
//!
//! [`tokio::time::Sleep`] operates on [`Instant`](std::time::Instant)s,
//! which are inconvenient when one wants to schedule timers according to wall-clock time.
//!
//! Implemented using [`libc::timerfd_create`] et al.

use std::{
    ffi::c_void,
    future::poll_fn,
    io,
    mem::{self, MaybeUninit},
    ops::Neg,
    os::fd::{AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd},
    pin::Pin,
    ptr,
    task::{Context, Poll, ready},
    time::{Duration, SystemTime},
};

use futures_core::Stream;
use tokio::io::unix::AsyncFd;

/// Expire at a specific [`SystemTime`].
#[derive(Debug)]
pub struct Deadline {
    inner: AsyncFd<RawSystemTimer>,
}

impl Deadline {
    pub fn new(at: SystemTime) -> Self {
        Self {
            inner: {
                let mut raw = RawSystemTimer::new().unwrap();
                raw.set_timing(RawTiming {
                    initial_expiration: Some(at),
                    interval: None,
                });
                AsyncFd::new(raw).unwrap()
            },
        }
    }
    /// # Race
    /// There is a race condition if the deadline is reset after after the reactor has been notified.
    pub fn reset(&mut self, to: SystemTime) {
        // TODO(aatifsyed): is this racy?
        self.inner.get_mut().set_timing(RawTiming {
            initial_expiration: Some(to),
            interval: None,
        });
    }
}

impl Future for Deadline {
    type Output = ();
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let Self { ref inner } = *self;
        ready!(inner.poll_read_ready(cx)).unwrap().clear_ready();
        Poll::Ready(())
    }
}

/// Tick every [`Duration`] since creation.
///
/// Bursts.
#[derive(Debug)]
pub struct Interval {
    missed_ticks: u64,
    inner: AsyncFd<RawSystemTimer>,
}

impl Interval {
    pub fn new(every: Duration) -> Self {
        Self {
            missed_ticks: 0,
            inner: {
                let mut raw = RawSystemTimer::new().unwrap();
                raw.set_timing(RawTiming {
                    initial_expiration: None,
                    interval: Some(every),
                });
                AsyncFd::new(raw).unwrap()
            },
        }
    }
    pub fn poll_tick(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        let Self {
            missed_ticks,
            inner,
        } = &mut *self;
        loop {
            if let Some(new) = missed_ticks.checked_sub(1) {
                *missed_ticks = new;
                return Poll::Ready(());
            }
            match ready!(inner.poll_read_ready_mut(cx)) {
                Ok(mut gd) => {
                    let mut missed = MaybeUninit::<u64>::uninit();
                    let count = mem::size_of_val(&missed);
                    let n = unsafe {
                        libc::read(
                            gd.get_mut().as_raw_fd(),
                            missed.as_mut_ptr().cast::<c_void>(),
                            count,
                        )
                    };
                    assert_eq!(n, count as isize, "{}", io::Error::last_os_error());
                    *missed_ticks = unsafe { missed.assume_init() };
                    gd.clear_ready();
                }
                Err(e) => panic!("io error on TimerFd {inner:?}: {e}"),
            }
        }
    }
    pub async fn tick(&mut self) {
        poll_fn(|cx| Pin::new(&mut *self).poll_tick(cx)).await
    }
}

impl Stream for Interval {
    type Item = ();
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.poll_tick(cx).map(Some)
    }
}

/// Tick every [`Duration`] after a [`SystemTime`].
///
/// Bursts.
#[derive(Debug)]
pub struct IntervalAfter {
    wrapped: Interval,
}

impl IntervalAfter {
    pub fn new(after: SystemTime, every: Duration) -> Self {
        Self {
            wrapped: Interval {
                missed_ticks: 0,
                inner: {
                    let mut raw = RawSystemTimer::new().unwrap();
                    raw.set_timing(RawTiming {
                        initial_expiration: Some(after),
                        interval: Some(every),
                    });
                    AsyncFd::new(raw).unwrap()
                },
            },
        }
    }

    pub fn poll_tick(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        Pin::new(&mut self.wrapped).poll_tick(cx)
    }
    pub async fn tick(&mut self) {
        poll_fn(|cx| Pin::new(&mut *self).poll_tick(cx)).await
    }
}

impl Stream for IntervalAfter {
    type Item = ();
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.poll_tick(cx).map(Some)
    }
}

#[derive(Debug)]
pub struct RawSystemTimer {
    fd: OwnedFd,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
pub struct RawTiming {
    pub initial_expiration: Option<SystemTime>,
    pub interval: Option<Duration>,
}

impl RawTiming {
    pub const DISABLED: Self = Self {
        initial_expiration: None,
        interval: None,
    };
    pub fn is_disabled(&self) -> bool {
        matches!(*self, Self::DISABLED)
    }
    pub fn is_enabled(&self) -> bool {
        !self.is_disabled()
    }
}

impl FromRawFd for RawSystemTimer {
    unsafe fn from_raw_fd(fd: RawFd) -> Self {
        Self {
            fd: unsafe { OwnedFd::from_raw_fd(fd) },
        }
    }
}

impl IntoRawFd for RawSystemTimer {
    fn into_raw_fd(self) -> RawFd {
        let Self { fd } = self;
        fd.into_raw_fd()
    }
}

const TIMESPEC_ZERO: libc::timespec = libc::timespec {
    tv_sec: 0,
    tv_nsec: 0,
};

impl RawSystemTimer {
    pub fn new() -> io::Result<Self> {
        let res = unsafe { libc::timerfd_create(libc::CLOCK_REALTIME, libc::TFD_NONBLOCK) };
        match res == -1 {
            false => Ok(Self {
                fd: unsafe { OwnedFd::from_raw_fd(res) },
            }),
            true => Err(io::Error::last_os_error()),
        }
    }
    pub fn timing(&self) -> RawTiming {
        let Self { fd } = self;
        let mut out = MaybeUninit::uninit();
        let rc = unsafe { libc::timerfd_gettime(fd.as_raw_fd(), out.as_mut_ptr()) };
        assert_eq!(
            rc,
            0,
            "timerfd_gettime({fd:?}) threw {}",
            io::Error::last_os_error()
        );
        let libc::itimerspec {
            it_value,
            it_interval,
        } = unsafe { out.assume_init() };

        let initial_expiration = match it_value {
            libc::timespec {
                tv_sec: 0,
                tv_nsec: 0,
            } => None,
            libc::timespec { tv_sec, tv_nsec } => {
                let mut instant = SystemTime::UNIX_EPOCH;
                match tv_sec.is_positive() {
                    true => instant += Duration::from_secs(tv_sec as u64),
                    false => instant -= Duration::from_secs(tv_sec.unsigned_abs()),
                }
                match tv_nsec.is_positive() {
                    true => instant += Duration::from_nanos(tv_sec as u64),
                    false => instant -= Duration::from_nanos(tv_sec.unsigned_abs()),
                }
                Some(instant)
            }
        };

        let interval = match it_interval {
            libc::timespec {
                tv_sec: 0,
                tv_nsec: 0,
            } => None,
            libc::timespec { tv_sec, tv_nsec } => {
                let mut duration = Duration::ZERO;
                match tv_sec.is_positive() {
                    true => duration += Duration::from_secs(tv_sec as u64),
                    false => duration -= Duration::from_secs(tv_sec.unsigned_abs()),
                }
                match tv_nsec.is_positive() {
                    true => duration += Duration::from_nanos(tv_sec as u64),
                    false => duration -= Duration::from_nanos(tv_sec.unsigned_abs()),
                }
                Some(duration)
            }
        };

        RawTiming {
            initial_expiration,
            interval,
        }
    }
    pub fn set_timing(&mut self, to: RawTiming) {
        let Self { fd } = self;
        let RawTiming {
            initial_expiration,
            interval,
        } = to;
        let it_value = match initial_expiration {
            Some(SystemTime::UNIX_EPOCH) => panic!(
                "Due to platform limitations, SystemTime::UNIX_EPOCH is not supported as an `initial_expiration`"
            ),
            Some(time) => {
                let (neg, dur) = match time.duration_since(SystemTime::UNIX_EPOCH) {
                    Ok(pos) => (false, pos),
                    Err(e) => (true, e.duration()),
                };
                libc::timespec {
                    tv_sec: neg_if(
                        i64::try_from(dur.as_secs()).unwrap_or_else(|e| {
                            panic!("can't fit duration {dur:?} in a timespec: {e}")
                        }),
                        neg,
                    ),
                    tv_nsec: neg_if(i64::from(dur.subsec_nanos()), neg),
                }
            }
            None => TIMESPEC_ZERO,
        };
        let it_interval = match interval {
            Some(dur) => libc::timespec {
                tv_sec: dur
                    .as_secs()
                    .try_into()
                    .unwrap_or_else(|e| panic!("can't fit duration {dur:?} in a timespec: {e}")),
                tv_nsec: i64::from(dur.subsec_nanos()),
            },
            None => TIMESPEC_ZERO,
        };
        let rc = unsafe {
            libc::timerfd_settime(
                fd.as_raw_fd(),
                libc::TFD_TIMER_ABSTIME,
                &libc::itimerspec {
                    it_interval,
                    it_value,
                },
                ptr::null_mut(),
            )
        };
        assert_eq!(
            rc,
            0,
            "timerfd_settime({fd:?}, {to:?}) threw {}",
            io::Error::last_os_error()
        );
    }
}

fn neg_if<T: Neg<Output = T>>(t: T, neg: bool) -> T {
    match neg {
        true => -t,
        false => t,
    }
}

impl AsRawFd for RawSystemTimer {
    fn as_raw_fd(&self) -> RawFd {
        self.fd.as_raw_fd()
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;

    use super::*;

    #[test]
    fn raw() {
        let mut timer = RawSystemTimer::new().unwrap();
        let timing = timer.timing();
        assert_eq!(timing, RawTiming::DISABLED);
        timer.set_timing(RawTiming {
            initial_expiration: Some(SystemTime::now()),
            interval: Some(Duration::from_mins(60)),
        });
        assert_matches!(
            timer.timing(),
            RawTiming {
                initial_expiration: Some(_),
                interval: Some(_)
            }
        );
    }

    #[tokio::test]
    async fn deadline() {
        let t0 = SystemTime::now();
        let duration = Duration::from_secs(5);
        Deadline::new(t0 + duration).await;
        assert!(t0.elapsed().unwrap() >= duration);
    }

    #[tokio::test]
    async fn interval_since() {
        let mut interval = IntervalAfter::new(
            SystemTime::UNIX_EPOCH + Duration::from_nanos(1),
            Duration::from_hours(1),
        );
        interval.tick().await;
        assert!(interval.wrapped.missed_ticks > 100);
    }
}