Time scheduler for production and testing (completely untested code at this point)

readme.md

Control time during tests

If a system under test (SUT) uses things like Task.Delay, DateTimeOffset.UtcNow, or PeriodicTimer, it becomes hard to create tests that runs fast and predictably.

The idea is to replace the use of e.g. Task.Delay with an abstraction, the ITimeScheduler, that in production is represented by the TimeScheduler, that just uses the real Task.Delay. During testing it is now possible to pass in TestScheduler, that allows the test to control the progress of time, making it possible to skip ahead, e.g. 10 minutes, and also pause time, leading to fast and predictable tests.

As an example, lets test the StuffService below that performs a specific tasks every 10 second:

public class StuffService 
{
  private static readonly TimeSpan doStuffDelay = TimeSpan.FromSeconds(10);
  private readonly ITimeScheduler scheduler;
  private readonly List<string> container;

  public StuffService(ITimeScheduler scheduler, List<string> container)
  {
    this.scheduler = scheduler;
    this.container = container;
  }
  
  public async Task DoStuff(CancellationToken cancelllationToken)
  {
    using var periodicTimer = scheduler.PeriodicTimer(doStuffDelay);
    
    while (await periodicTimer.WaitForNextTickAsync(cancellationToken))
    {
      container.Add("stuff");    
    }
  }
}

The test, using xUnit and FluentAssertions, could look like this:

[Fact]
public void DoStuff_does_stuff_every_10_seconds()
{
  // Arrange
  var scheduler = new TestScheduler();
  var container = new List<string>();  
  var sut = new StuffService(scheduler, container);
  
  // Act
  _ = sut.DoStuff(CancellationToken.None);
  scheduler.ForwardTime(TimeSpan.FromSeconds(10));
  
  // Assert
  container.Should().ContainSingle();
}

ITimeScheduler.cs

namespace Scheduler;

public interface ITimeScheduler
{
    DateTimeOffset UtcNow { get; }

    Task Delay(TimeSpan delay);

    Task Delay(TimeSpan delay, CancellationToken cancellationToken);

    PeriodicTimer PeriodicTimer(TimeSpan period);
}

PeriodicTimer.cs

namespace Scheduler;

/// <summary>Provides a periodic timer that enables waiting asynchronously for timer ticks.</summary>
/// <remarks>
/// This timer is intended to be used only by a single consumer at a time: only one call to <see cref="WaitForNextTickAsync" />
/// may be in flight at any given moment.  <see cref="Dispose()"/> may be used concurrently with an active <see cref="WaitForNextTickAsync"/>
/// to interrupt it and cause it to return false.
/// </remarks>
public abstract class PeriodicTimer : IDisposable
{
    /// <summary>Wait for the next tick of the timer, or for the timer to be stopped.</summary>
    /// <param name="cancellationToken">
    /// A <see cref="CancellationToken"/> to use to cancel the asynchronous wait. If cancellation is requested, it affects only the single wait operation;
    /// the underlying timer continues firing.
    /// </param>
    /// <returns>A task that will be completed due to the timer firing, <see cref="Dispose()"/> being called to stop the timer, or cancellation being requested.</returns>
    /// <remarks>
    /// The <see cref="PeriodicTimer"/> behaves like an auto-reset event, in that multiple ticks are coalesced into a single tick if they occur between
    /// calls to <see cref="WaitForNextTickAsync"/>.  Similarly, a call to <see cref="Dispose()"/> will void any tick not yet consumed. <see cref="WaitForNextTickAsync"/>
    /// may only be used by one consumer at a time, and may be used concurrently with a single call to <see cref="Dispose()"/>.
    /// </remarks>
    public abstract ValueTask<bool> WaitForNextTickAsync(CancellationToken cancellationToken = default);

    /// <summary>Stops the timer and releases associated managed resources.</summary>
    /// <remarks>
    /// <see cref="Dispose()"/> will cause an active wait with <see cref="WaitForNextTickAsync"/> to complete with a value of false.
    /// All subsequent <see cref="WaitForNextTickAsync"/> invocations will produce a value of false.
    /// </remarks>
    public void Dispose()
    {
        Dispose(true);
        GC.SuppressFinalize(this);
    }

    protected virtual void Dispose(bool disposing)
    {
    }
}

TestScheduler.cs

using System.Collections.Concurrent;

namespace Scheduler.Testing;

public sealed partial class TestScheduler : ITimeScheduler, IDisposable
{
    private readonly List<DelayedAction> delayedActions = new();

    public DateTimeOffset UtcNow { get; private set; }

    public TestScheduler()
    {
        UtcNow = DateTimeOffset.UtcNow;
    }

    public TestScheduler(DateTimeOffset startDateTime)
    {
        UtcNow = startDateTime;
    }

    public Task Delay(TimeSpan delay)
    {
        return Delay(delay, CancellationToken.None);
    }

    public Task Delay(TimeSpan delay, CancellationToken cancellationToken)
    {
        var tcs = new TaskCompletionSource();
        RegisterDelayedTask(UtcNow + delay, () => tcs.TrySetResult(), () => tcs.TrySetCanceled(), cancellationToken);
        return tcs.Task;
    }

    public PeriodicTimer PeriodicTimer(TimeSpan period)
    {
        return new TestPeriodicTimer(period, this);
    }

    public void ForwardTime(TimeSpan timeSpan)
    {
        UtcNow = UtcNow + timeSpan;
        CompleteDelayedTasks();
    }

    private void CompleteDelayedTasks()
    {
        var tasksToComplete = delayedActions
            .Where(x => x.CompletionTime <= UtcNow)
            .OrderBy(x => x.CompletionTime)
            .ToArray();

        foreach (var pair in tasksToComplete)
        {
            pair.Complete();
            delayedActions.Remove(pair);
        }
    }

    private void RegisterDelayedTask(
        DateTimeOffset completionTime,
        Action complete,
        Action cancel,
        CancellationToken cleanupToken = default)
    {
        var delayedAction = new DelayedAction(completionTime, complete, cancel);

        cleanupToken.Register(() =>
        {
            cancel();
            delayedActions.Remove(delayedAction);
        });

        delayedActions.Add(delayedAction);
    }

    public void Dispose()
    {
        foreach (var delayedAction in delayedActions)
        {
            delayedAction.Cancel();
        }
    }

    private sealed class DelayedAction
    {
        private readonly Action complete;
        private readonly Action cancel;
        private int completed;

        public DateTimeOffset CompletionTime { get; }

        public DelayedAction(DateTimeOffset completionTime, Action complete, Action cancel)
        {
            CompletionTime = completionTime;
            this.complete = complete;
            this.cancel = cancel;
        }

        public void Complete()
        {
            // Ensure that complete/cancel is only being called once.
            if (Interlocked.CompareExchange(ref completed, 1, 0) == 0)
            {
                complete();
            }
        }

        public void Cancel()
        {
            // Ensure that cancel/cancel is only being called once.
            if (Interlocked.CompareExchange(ref completed, 2, 0) == 0)
            {
                cancel();
            }
        }
    }
}

TestScheduler.TestPeriodicTimer.cs

namespace Scheduler.Testing;

public partial class TestScheduler : ITimeScheduler, IDisposable
{
    private sealed class TestPeriodicTimer : PeriodicTimer
    {
        private readonly TimeSpan period;
        private readonly TestScheduler owner;
        private bool stopped;
        private TaskCompletionSource<bool>? completionSource;
        private DateTimeOffset nextSignal;

        public TestPeriodicTimer(TimeSpan period, TestScheduler owner)
        {
            this.period = period;
            this.owner = owner;
            SetNextSignalTime();

        }

        public override ValueTask<bool> WaitForNextTickAsync(CancellationToken cancellationToken = default)
        {
            if (completionSource is not null)
            {
                throw new InvalidOperationException("WaitForNextTickAsync should only be used by one consumer at a time. Failing to do so is an error.");
            }

            if (cancellationToken.IsCancellationRequested)
            {
                return ValueTask.FromCanceled<bool>(cancellationToken);
            }

            if (!stopped && owner.UtcNow >= nextSignal)
            {
                SetNextSignalTime();
                return new ValueTask<bool>(!stopped);
            }

            completionSource = new TaskCompletionSource<bool>();

            owner.RegisterDelayedTask(
                nextSignal,
                () => Signal(),
                () => completionSource?.TrySetCanceled(),
                cancellationToken);

            return new ValueTask<bool>(completionSource.Task);
        }

        private void Signal()
        {
            if (completionSource is not null)
            {
                completionSource.TrySetResult(!stopped);
                completionSource = null;
            }

            SetNextSignalTime();
        }

        private void SetNextSignalTime()
        {
            nextSignal = owner.UtcNow + period;
        }

        protected override void Dispose(bool disposing)
        {
            stopped = true;
            Signal();
            base.Dispose(disposing);
        }
    }
}

TestSchedulerTests.cs

namespace Scheduler.Testing;

public class TestSchedulerTests
{
    [Fact]
    public void ForwardTime_updates_UtcNow()
    {
        var startTime = DateTimeOffset.UtcNow;
        using var sut = new TestScheduler(startTime);

        sut.ForwardTime(TimeSpan.FromTicks(1));

        sut.UtcNow.Should().Be(startTime + TimeSpan.FromTicks(1));
    }

    [Fact]
    public void Delayed_task_is_completes()
    {
        var startTime = DateTimeOffset.UtcNow;
        var future = TimeSpan.FromTicks(1);
        using var sut = new TestScheduler(startTime);
        var task = sut.Delay(TimeSpan.FromTicks(1));

        sut.ForwardTime(future);

        task.Status.Should().Be(TaskStatus.RanToCompletion);
    }

    [Fact]
    public void Delayed_task_is_cancelled()
    {
        using var cts = new CancellationTokenSource();
        using var sut = new TestScheduler();
        var task = sut.Delay(TimeSpan.FromTicks(1), cts.Token);

        cts.Cancel();

        task.Status.Should().Be(TaskStatus.Canceled);
    }

    [Fact]
    public void PeriodicTimer_WaitForNextTickAsync_cancelled_immidiately()
    {
        using var cts = new CancellationTokenSource();
        using var sut = new TestScheduler();
        using var periodicTimer = sut.PeriodicTimer(TimeSpan.FromTicks(1));

        cts.Cancel();
        var task = periodicTimer.WaitForNextTickAsync(cts.Token);

        task.IsCanceled.Should().BeTrue();
    }

    [Fact]
    public async Task PeriodicTimer_WaitForNextTickAsync_complete_immidiately()
    {
        using var sut = new TestScheduler();
        using var periodicTimer = sut.PeriodicTimer(TimeSpan.FromTicks(1));

        sut.ForwardTime(TimeSpan.FromTicks(1));
        var task = periodicTimer.WaitForNextTickAsync();

        (await task).Should().BeTrue();
    }

    [Fact]
    public async Task PeriodicTimer_WaitForNextTickAsync_completes()
    {
        var startTime = DateTimeOffset.UtcNow;
        var future = TimeSpan.FromTicks(1);
        using var sut = new TestScheduler(startTime);
        using var periodicTimer = sut.PeriodicTimer(TimeSpan.FromTicks(1));
        var task = periodicTimer.WaitForNextTickAsync();

        sut.ForwardTime(future);

        (await task).Should().BeTrue();
    }

    [Fact]
    public async Task PeriodicTimer_WaitForNextTickAsync_completes_after_dispose()
    {
        var startTime = DateTimeOffset.UtcNow;
        using var sut = new TestScheduler(startTime);
        var periodicTimer = sut.PeriodicTimer(TimeSpan.FromTicks(1));
        var task = periodicTimer.WaitForNextTickAsync();

        periodicTimer.Dispose();

        (await task).Should().BeFalse();
    }

    [Fact]
    public async Task PeriodicTimer_WaitForNextTickAsync_cancelled_with_exception()
    {
        using var cts = new CancellationTokenSource();
        using var sut = new TestScheduler();
        using var periodicTimer = sut.PeriodicTimer(TimeSpan.FromTicks(1));
        var task = periodicTimer.WaitForNextTickAsync(cts.Token);
        cts.CancelAfter(TimeSpan.Zero);

        var throws = async () => await task;

        await throws.Should().ThrowAsync<OperationCanceledException>();
    }
}

TimeScheduler.cs

namespace Scheduler;

public partial class TimeScheduler : ITimeScheduler
{
    public DateTimeOffset UtcNow => DateTimeOffset.UtcNow;

    public Task Delay(TimeSpan delay) => Task.Delay(delay);

    public Task Delay(TimeSpan delay, CancellationToken cancellationToken)
        => Task.Delay(delay, cancellationToken);

    public PeriodicTimer PeriodicTimer(TimeSpan period)
        => new PeriodicTimerWrapper(period);
}

TimeScheduler.PeriodicTimerWrapper.cs

namespace Scheduler;

public partial class TimeScheduler : ITimeScheduler
{
    private sealed class PeriodicTimerWrapper : PeriodicTimer
    {
        private readonly System.Threading.PeriodicTimer timer;

        public PeriodicTimerWrapper(TimeSpan period)
        {
            timer = new System.Threading.PeriodicTimer(period);
        }

        public override ValueTask<bool> WaitForNextTickAsync(CancellationToken cancellationToken = default)
            => timer.WaitForNextTickAsync(cancellationToken);

        protected override void Dispose(bool disposing)
        {
            timer.Dispose();
            base.Dispose(disposing);
        }
    }
}