Search code examples
c#.netrate-limitingpollyretry-logic

Adjusting code to use the built-in RateLimit policy


I want to achieve the same behavior as the policy below with the built-in RateLimit policy, i.e. the logger message and to read the Retry-After header and wait the exact seconds that were needed to wait for but using the built-in RateLimit policy.

Attempt

// TODO: No logger message and not sure if it waits the time taken from the Retry-After header.
public static AsyncRateLimitPolicy Limit<T>(ILogger<T> logger)
{
    return Policy.RateLimitAsync(RateLimitRetryCount, TimeSpan.FromSeconds(5));
}

Works

public static AsyncRetryPolicy<RestResponse> AsyncRateLimit<T>(ILogger<T> logger)
{
    return Policy.HandleResult<RestResponse>(response => response.StatusCode == HttpStatusCode.TooManyRequests)
        .WaitAndRetryAsync(RateLimitRetryCount,
            (attemptCount, restResponse, _) =>
            {
                var retryAfterHeader = restResponse?.Result?.Headers?.SingleOrDefault(h => h.Name == "Retry-After");
                double secondsInterval = 0;

                if (retryAfterHeader != null)
                {
                    var value = retryAfterHeader.Value?.ToString();
                    if (!double.TryParse(value, out secondsInterval))
                    {
                        secondsInterval = Math.Pow(2, attemptCount);
                    }
                }

                return TimeSpan.FromSeconds(secondsInterval);
            },
            (response, timeSpan, retryCount, _) =>
            {
                logger.LogTrace(
                    "The API request has been rate limited. HttpStatusCode={StatusCode}. Waiting {Seconds} seconds before retry. Number attempt {RetryCount}. Uri={Url}; RequestResponse={Content}",
                    response.Result.StatusCode, timeSpan.TotalSeconds, retryCount, response.Result.ResponseUri, response.Result.Content);

                return Task.CompletedTask;
            });
}

Solution

  • There were multiple questions so let me answer all of them.

    1) How to inject logger to a Policy?

    You need to use Polly's context for that.

    The context is created outside of the policy. It is used as a container to store any arbitrary information

    var context = new Context().WithLogger(logger);
    

    Then it is passed through the Execute/ExecuteAsync call

    await policy.ExecuteAsync(ctx => FooAsync(), context);
    

    Finally you can use the context in any user delegate (like onRetry/onRetryAsync) to retrieve the passed object

    (exception, timeSpan, retryCount, context) =>
    {
      var logger = context.GetLogger();
      logger?.LogWarning(...);
      ...
    }
    

    The WithLogger and GetLogger extension methods

    public static class ContextExtensions
    {
        private static readonly string LoggerKey = "LoggerKey";
    
        public static Context WithLogger(this Context context, ILogger logger)
        {
            context[LoggerKey] = logger;
            return context;
        }
    
        public static ILogger GetLogger(this Context context)
        {
            if (context.TryGetValue(LoggerKey, out object logger))
            {
                return logger as ILogger;
            }
            return null;
        }
    }
    

    2) Does the above rate limiter work in the same way as the retry?

    No. The rate limiter is a proactive policy which can be useful to prevent resource abuse. That means it will throw an RateLimitRejectedException if the predefined limit is exceeded.

    Whenever we are talking about resilience strategy we are referring to a predefined protocol between the two parties to overcome on transient failures. So the rate limiter is the server-side of this story whereas the retry (reactive policy) is the client-side.

    If you want to set the RetryAfter header in your rate limiter then you can do that like this

    IAsyncPolicy<HttpResponseMessage> limit = Policy
        .RateLimitAsync(RateLimitRetryCount, TimeSpan.FromSeconds(5), RateLimitRetryCount,
            (retryAfter, context) => {
                var response = new HttpResponseMessage(System.Net.HttpStatusCode.TooManyRequests);
                response.Headers.Add("Retry-After", retryAfter.TotalSeconds.ToString());
                return response;
            });
    

    Then on the client-side inside your retry's sleepDurationProvider delegate you can retrieve that value like this if the response is a DelegateResult<HttpResponseMessage>

    response.Result.Headers.RetryAfter.Delta ?? TimeSpan.FromSeconds(0)