Lines
87.39 %
Functions
61.6 %
Branches
100 %
//! Functionality for simulating the passage of time in unit tests.
//!
//! We do this by providing [`MockSleepProvider`], a "SleepProvider"
//! instance that can simulate timeouts and retries without requiring
//! the actual system clock to advance.
#![allow(clippy::missing_docs_in_private_items)]
use std::{
cmp::{Eq, Ordering, PartialEq, PartialOrd},
collections::BinaryHeap,
pin::Pin,
sync::{Arc, Mutex, Weak},
task::{Context, Poll, Waker},
time::{Duration, Instant, SystemTime},
};
use futures::Future;
use tracing::trace;
use std::collections::HashSet;
use tor_rtcompat::SleepProvider;
/// A dummy [`SleepProvider`] instance for testing.
///
/// The MockSleepProvider ignores the current time, and instead keeps
/// its own view of the current `Instant` and `SystemTime`. You
/// can advance them in-step by calling `advance()`, and you can simulate
/// jumps in the system clock by calling `jump()`.
/// This is *not* for production use.
#[derive(Clone)]
pub struct MockSleepProvider {
/// The shared backend for this MockSleepProvider and its futures.
state: Arc<Mutex<SleepSchedule>>,
}
/// Shared backend for sleep provider and Sleeping futures.
struct SleepSchedule {
/// What time do we pretend it is (monotonic)? This value only
/// moves forward.
instant: Instant,
/// What time do we pretend it is (wall clock)? This value can move
/// in any way, but usually moves in step with `instant`.
wallclock: SystemTime,
/// Priority queue of events, in the order that we should wake them.
sleepers: BinaryHeap<SleepEntry>,
/// If the mock time system is being driven by a `WaitFor`, holds a `Waker` to wake up that
/// `WaitFor` in order for it to make more progress.
waitfor_waker: Option<Waker>,
/// Number of sleepers instantiated.
sleepers_made: usize,
/// Number of sleepers polled.
sleepers_polled: usize,
/// Whether an advance is needed.
should_advance: bool,
/// A set of reasons why advances shouldn't be allowed right now.
blocked_advance: HashSet<String>,
/// A time up to which advances are allowed, irrespective of them being blocked.
allowed_advance: Duration,
/// An entry telling us when to wake which future up.
struct SleepEntry {
/// The time at which this entry should wake
when: Instant,
/// The Waker to call when the instant has passed.
waker: Waker,
/// A future returned by [`MockSleepProvider::sleep()`].
pub struct Sleeping {
/// The instant when we should become ready.
/// True if we have pushed this into the queue.
inserted: bool,
/// The schedule to queue ourselves in if we're polled before we're ready.
provider: Weak<Mutex<SleepSchedule>>,
impl MockSleepProvider {
/// Create a new MockSleepProvider, starting at a given wall-clock time.
pub fn new(wallclock: SystemTime) -> Self {
let instant = Instant::now();
let sleepers = BinaryHeap::new();
let state = SleepSchedule {
instant,
wallclock,
sleepers,
waitfor_waker: None,
sleepers_made: 0,
sleepers_polled: 0,
should_advance: false,
blocked_advance: HashSet::new(),
allowed_advance: Duration::from_nanos(0),
MockSleepProvider {
state: Arc::new(Mutex::new(state)),
/// Advance the simulated timeline forward by `dur`.
/// Calling this function will wake any pending futures as
/// appropriate, and yield to the scheduler so they get a chance
/// to run.
/// # Limitations
/// This function advances time in one big step. We might instead
/// want to advance in small steps and make sure that each step's
/// futures can get run before the ones scheduled to run after it.
pub async fn advance(&self, dur: Duration) {
self.advance_noyield(dur);
tor_rtcompat::task::yield_now().await;
/// appropriate, but not yield to the scheduler. Mostly you
/// should call [`advance`](Self::advance) instead.
pub(crate) fn advance_noyield(&self, dur: Duration) {
// It's not so great to unwrap here in general, but since this is
// only testing code we don't really care.
let mut state = self.state.lock().expect("Poisoned lock for state");
state.wallclock += dur;
state.instant += dur;
state.fire();
/// Simulate a discontinuity in the system clock, by jumping to
/// `new_wallclock`.
/// # Panics
/// Panics if we have already panicked while holding the lock on
/// the internal timer state, and the lock is poisoned.
pub fn jump_to(&self, new_wallclock: SystemTime) {
state.wallclock = new_wallclock;
/// Return the amount of virtual time until the next timeout
/// should elapse.
/// If there are no more timeouts, return None. If the next
/// timeout should elapse right now, return Some(0).
pub(crate) fn time_until_next_timeout(&self) -> Option<Duration> {
let state = self.state.lock().expect("Poisoned lock for state");
let now = state.instant;
state
.sleepers
.peek()
.map(|sleepent| sleepent.when.saturating_duration_since(now))
/// Return true if a `WaitFor` driving this sleep provider should advance time in order for
/// futures blocked on sleeping to make progress.
/// NOTE: This function has side-effects; if it returns true, the caller is expected to do an
/// advance before calling it again.
pub(crate) fn should_advance(&mut self) -> bool {
if !state.blocked_advance.is_empty() && state.allowed_advance == Duration::from_nanos(0) {
// We've had advances blocked, and don't have any quota for doing allowances while
// blocked left.
trace!(
"should_advance = false: blocked by {:?}",
state.blocked_advance
);
return false;
if !state.should_advance {
// The advance flag wasn't set.
trace!("should_advance = false; bit not previously set");
// Clear the advance flag; we'll either return true and cause an advance to happen,
// or the reasons to return false below also imply that the advance flag will be set again
// later on.
state.should_advance = false;
if state.sleepers_polled < state.sleepers_made {
// Something did set the advance flag before, but it's not valid any more now because
// more unpolled sleepers were created.
trace!("should_advance = false; advancing no longer valid");
if !state.blocked_advance.is_empty() && state.allowed_advance > Duration::from_nanos(0) {
// If we're here, we would've returned earlier due to having advances blocked, but
// we have quota to advance up to a certain time while advances are blocked.
// Let's see when the next timeout is, and whether it falls within that quota.
let next_timeout = {
let next_timeout = match next_timeout {
Some(x) => x,
None => {
// There's no timeout set, so we really shouldn't be here anyway.
trace!("should_advance = false; allow_one set but no timeout yet");
if next_timeout <= state.allowed_advance {
// We can advance up to the next timeout, since it's in our quota.
// Subtract the amount we're going to advance by from said quota.
state.allowed_advance -= next_timeout;
"WARNING: allowing advance due to allow_one; new allowed is {:?}",
state.allowed_advance
} else {
// The next timeout is too far in the future.
"should_advance = false; allow_one set but only up to {:?}, next is {:?}",
state.allowed_advance,
next_timeout
true
/// Register a `Waker` to be woken up when an advance in time is required to make progress.
/// This is used by `WaitFor`.
pub(crate) fn register_waitfor_waker(&mut self, waker: Waker) {
state.waitfor_waker = Some(waker);
/// Remove a previously registered `Waker` registered with `register_waitfor_waker()`.
pub(crate) fn clear_waitfor_waker(&mut self) {
state.waitfor_waker = None;
/// Returns true if a `Waker` has been registered with `register_waitfor_waker()`.
/// This is used to ensure that you don't have two concurrent `WaitFor`s running.
pub(crate) fn has_waitfor_waker(&self) -> bool {
state.waitfor_waker.is_some()
impl SleepSchedule {
/// Wake any pending events that are ready according to the
/// current simulated time.
fn fire(&mut self) {
use std::collections::binary_heap::PeekMut;
let now = self.instant;
while let Some(top) = self.sleepers.peek_mut() {
if now < top.when {
return;
PeekMut::pop(top).waker.wake();
/// Add a new SleepEntry to this schedule.
fn push(&mut self, ent: SleepEntry) {
self.sleepers.push(ent);
/// If all sleepers made have been polled, set the advance flag and wake up any `WaitFor` that
/// might be waiting.
fn maybe_advance(&mut self) {
if self.sleepers_polled >= self.sleepers_made {
if let Some(ref waker) = self.waitfor_waker {
trace!("setting advance flag");
self.should_advance = true;
waker.wake_by_ref();
trace!("would advance, but no waker");
/// Register a sleeper as having been polled, and advance if necessary.
fn increment_poll_count(&mut self) {
self.sleepers_polled += 1;
"sleeper polled, {}/{}",
self.sleepers_polled,
self.sleepers_made
self.maybe_advance();
impl SleepProvider for MockSleepProvider {
type SleepFuture = Sleeping;
fn sleep(&self, duration: Duration) -> Self::SleepFuture {
let mut provider = self.state.lock().expect("Poisoned lock for state");
let when = provider.instant + duration;
// We're making a new sleeper, so register this in the state.
provider.sleepers_made += 1;
"sleeper made for {:?}, {}/{}",
duration,
provider.sleepers_polled,
provider.sleepers_made
Sleeping {
when,
inserted: false,
provider: Arc::downgrade(&self.state),
fn block_advance<T: Into<String>>(&self, reason: T) {
let reason = reason.into();
trace!("advancing blocked: {}", reason);
provider.blocked_advance.insert(reason);
fn release_advance<T: Into<String>>(&self, reason: T) {
trace!("advancing released: {}", reason);
provider.blocked_advance.remove(&reason);
if provider.blocked_advance.is_empty() {
provider.maybe_advance();
fn allow_one_advance(&self, dur: Duration) {
provider.allowed_advance = Duration::max(provider.allowed_advance, dur);
"** allow_one_advance fired; may advance up to {:?} **",
provider.allowed_advance
fn now(&self) -> Instant {
self.state.lock().expect("Poisoned lock for state").instant
fn wallclock(&self) -> SystemTime {
self.state
.lock()
.expect("Poisoned lock for state")
.wallclock
impl PartialEq for SleepEntry {
fn eq(&self, other: &Self) -> bool {
self.when == other.when
impl Eq for SleepEntry {}
impl PartialOrd for SleepEntry {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
impl Ord for SleepEntry {
fn cmp(&self, other: &Self) -> Ordering {
self.when.cmp(&other.when).reverse()
impl Drop for Sleeping {
fn drop(&mut self) {
if let Some(provider) = Weak::upgrade(&self.provider) {
let mut provider = provider.lock().expect("Poisoned lock for provider");
if !self.inserted {
// A sleeper being dropped will never be polled, so there's no point waiting;
// act as if it's been polled in order to avoid waiting forever.
trace!("sleeper dropped, incrementing count");
provider.increment_poll_count();
self.inserted = true;
impl Future for Sleeping {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let now = provider.instant;
if now >= self.when {
// The sleep time's elapsed.
// If we never registered this sleeper as being polled, do so now.
if !provider.should_advance {
// The first advance during a `WaitFor` gets triggered by all sleepers that
// have been created being polled.
// However, this only happens once.
// What we do to get around this is have sleepers that return Ready kick off
// another advance, in order to wake the next waiting sleeper.
return Poll::Ready(());
// dbg!("sleep check with", self.when-now);
let entry = SleepEntry {
when: self.when,
waker: cx.waker().clone(),
provider.push(entry);
// Register this sleeper as having been polled.
// dbg!(provider.sleepers.len());
Poll::Pending
#[cfg(test)]
mod test {
#![allow(clippy::unwrap_used)]
use super::*;
use tor_rtcompat::test_with_all_runtimes;
#[test]
fn basics_of_time_travel() {
let w1 = SystemTime::now();
let sp = MockSleepProvider::new(w1);
let i1 = sp.now();
assert_eq!(sp.wallclock(), w1);
let interval = Duration::new(4 * 3600 + 13 * 60, 0);
sp.advance_noyield(interval);
assert_eq!(sp.now(), i1 + interval);
assert_eq!(sp.wallclock(), w1 + interval);
sp.jump_to(w1 + interval * 3);
assert_eq!(sp.wallclock(), w1 + interval * 3);
fn time_moves_on() {
test_with_all_runtimes!(|_| async {
use futures::channel::oneshot;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
let sp = MockSleepProvider::new(SystemTime::now());
let one_hour = Duration::new(3600, 0);
let (s1, r1) = oneshot::channel();
let (s2, r2) = oneshot::channel();
let (s3, r3) = oneshot::channel();
let b1 = AtomicBool::new(false);
let b2 = AtomicBool::new(false);
let b3 = AtomicBool::new(false);
let real_start = Instant::now();
futures::join!(
async {
sp.sleep(one_hour).await;
b1.store(true, Ordering::SeqCst);
s1.send(()).unwrap();
},
sp.sleep(one_hour * 3).await;
b2.store(true, Ordering::SeqCst);
s2.send(()).unwrap();
sp.sleep(one_hour * 5).await;
b3.store(true, Ordering::SeqCst);
s3.send(()).unwrap();
sp.advance(one_hour * 2).await;
r1.await.unwrap();
assert!(b1.load(Ordering::SeqCst));
assert!(!b2.load(Ordering::SeqCst));
assert!(!b3.load(Ordering::SeqCst));
r2.await.unwrap();
assert!(b2.load(Ordering::SeqCst));
r3.await.unwrap();
assert!(b3.load(Ordering::SeqCst));
let real_end = Instant::now();
assert!(real_end - real_start < one_hour);
std::io::Result::Ok(())
});