Search code examples
c#unit-testing.net-8.0pollyretry-logic

Why is Task.Delay(1) necessary to advance clock when unit testing with .NET TimeProvider?


I have configured a retry based pipeline with Polly.

// Using Polly for retry logic
private readonly ResiliencePipeline _retryPipeline = new ResiliencePipelineBuilder { TimeProvider = timeProvider }
    .AddRetry(new RetryStrategyOptions
    {
        ShouldHandle = new PredicateBuilder().Handle<ConditionalCheckFailedException>(),
        Delay = TimeSpan.FromMilliseconds(_backoffFactor),
        MaxRetryAttempts = _maxAttempts - 1,

        // Linear backoff increases the delay each time by the backoff factor
        BackoffType = DelayBackoffType.Linear,
        OnRetry = onRetryArguments =>
        {
            logger.LogWarning(
                "Failed to acquire lock. Retrying. {@LogContext}",
                new { onRetryArguments });
            return ValueTask.CompletedTask;
        }
    })
    .Build();

Which I execute using

// Attempt to store the lock with backoff retry 
LockResult result = await _retryPipeline.ExecuteAsync(
   async _ => await AttemptLockStorageAsync(lockId, expiryMilliseconds, attempts++),
   cancellationTokenSource.Token);

When unit testing, I find that I have to add a Task.Delay(1) in order for Polly to perform the retries

// Act
Func<Task<DistributedLock>> func = async () =>
{
    Task<DistributedLock> result = _distributedLockService.AcquireLockAsync(lockId);
    for (int i = 1; i <= 4; i++)
    {
        _timeProvider.Advance(TimeSpan.FromMilliseconds(1000 * i + 1));
        await Task.Delay(1);
    }

    return await result;
};

// Assert
// We expect that we should be able to attempt 5 full times, rather than getting a TaskCancelledException.
(await func.Should().ThrowAsync<TimeoutException>()).WithMessage(
    $"Could not acquire lock {lockId}. Attempted 5 times.");

Why is the Task.Delay necessary?

Edit

TimeProvider provided to the SUT via the primary constructor.

public class DistributedLockService(
    IDistributedLockRepository distributedLockRepository,
    ILogger<DistributedLockService> logger,
    TimeProvider timeProvider)
    : IDisposable, IDistributedLockService

FakeTimer is provided in the unit test constructor

    private readonly FakeTimeProvider _timeProvider = new();

    public DistributedLockServiceTests()
    {
        _timeProvider.SetUtcNow(DateTimeOffset.Parse("2024-01-23", CultureInfo.InvariantCulture));
        _distributedLockService = new DistributedLockService(
            _distributedLockRepository.Object,
            _logger.Object,
            _timeProvider);
    }

Filed a bug report based on minimal reproduction https://github.com/App-vNext/Polly/issues/1932


Solution

  • As documented here

    https://github.com/dotnet/extensions/pull/5169/files

    The Advance method is used to simulate the passage of time. This can be useful in tests where you need to control the timing of asynchronous operations. When awaiting a task in a test that uses FakeTimeProvider, it's important to use ConfigureAwait(true).

    Here's an example:

    await provider.Delay(TimeSpan.FromSeconds(delay)).ConfigureAwait(true);
    

    This ensures that the continuation of the awaited task (i.e., the code that comes after the await statement) runs in the original context.

    For a more realistic example, consider the following test using Polly:

    using Polly;
    using Polly.Retry;
    
    public class SomeService(TimeProvider timeProvider)
    {
        // Don't do this in real life, not thread safe
        public int Tries { get; private set; }
    
        private readonly ResiliencePipeline _retryPipeline = new ResiliencePipelineBuilder { TimeProvider = timeProvider }
            .AddRetry(
                new RetryStrategyOptions
                {
                    ShouldHandle = new PredicateBuilder().Handle<InvalidOperationException>(),
                    Delay = TimeSpan.FromSeconds(1),
                    MaxRetryAttempts = 2,
                    BackoffType = DelayBackoffType.Linear,
                })
            .Build();
    
        public async Task<int> PollyRetry(double taskDelay, double cancellationSeconds)
        {
            CancellationTokenSource cts = new(TimeSpan.FromSeconds(cancellationSeconds), timeProvider);
            Tries = 0;
    
            // get a context from the pool and return it when done
            var context = ResilienceContextPool.Shared.Get(
                // ensure execution continues on captured context 
                continueOnCapturedContext: true, 
                cancellationToken: cts.Token);
    
            var result = await _retryPipeline.ExecuteAsync(
                async _ =>
                {
                    Tries++;
    
                    // Simulate a task that takes some time to complete
                    await Task.Delay(TimeSpan.FromSeconds(taskDelay), timeProvider).ConfigureAwait(true);
    
                    if (Tries <= 2)
                    {
                        throw new InvalidOperationException();
                    }
    
                    return Tries;
                },
                context);
    
            ResilienceContextPool.Shared.Return(context);
    
            return result;
        }
    }
    
    using Microsoft.Extensions.Time.Testing;
    
    public class SomeServiceTests
    {
        [Fact]
        public void PollyRetry_ShouldHave2Tries()
        {
            var timeProvider = new FakeTimeProvider();
            var someService = new SomeService(timeProvider);
    
            // Act
            var result = someService.PollyRetry(taskDelay: 1, cancellationSeconds: 6);
    
            // Advancing the time more than one second should resolves the first execution delay.
            timeProvider.Advance(TimeSpan.FromMilliseconds(1001));
    
            // Advancing the time more than the retry delay time of 1s,
            // and less then the task execution delay should start the second try
            timeProvider.Advance(TimeSpan.FromMilliseconds(1050));
    
            // Assert
            result.IsCompleted.Should().BeFalse();
            someService.Tries.Should().Be(2);
        }
    }