C# Programming

Extending Polly Retry Policies to cover IAsyncEnumerables

I recently encountered an issue where a Polly retry policy wasn’t covering a chunk of code the way that I expected. The reason being that I was wrapping a method that yielded elements to produce an IAsyncEnumerable. This caused the retry policy to complete before I reached the first element in the collection. I only received an exception when I attempted to enumerate the IAsyncEnumerable. This can be remedied by creating an extension method for an AsyncRetryPolicy that ensures the collection’s first yield is triggered before exiting the policy.

    public static async IAsyncEnumerable<TItem> ExecuteEnumerableAsync<TItem>(
        this AsyncRetryPolicy policy,
        Func<Context, CancellationToken, IAsyncEnumerable<TItem>> action,
        Context context,
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var (enumerator, movedNext) = await policy.ExecuteAsync(
            async (ctx, ct) =>
            {
                var asyncEnumerable = action(ctx, ct);
                var asyncEnumerator = asyncEnumerable.GetAsyncEnumerator();
                (IAsyncEnumerator<TItem> Enumerator, bool MovedNext) result =
                    (enumerator, await asyncEnumerator.MoveNextAsync(ct));

                return result;
            },
            context,
            cancellationToken);

        if (movedNext)
        {
            do
            {
                yield return enumerator.Current;
            }
            while (await enumerator.MoveNextAsync(cancellationToken));
        }
    }

Calling this extension method looks just like calling a normal retry policy.

var results = policy.ExecuteEnumerableAsync(
    (ctx, ct) => GetCollectionAsync(ct),
    new Context("some-context-name"),
    cancellationToken);