Handle retries in a more generic manner

mivano · Jun 15, 2024 · f7cc9a0 · f7cc9a0
1 parent a8a5f81
commit f7cc9a0
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/src/Infrastructure/PollyExtensions.cs b/src/Infrastructure/PollyExtensions.cs
@@ -3,24 +3,23 @@
 
 namespace AzureCostCli.Infrastructure;
 
+// ReSharper disable once ClassNeverInstantiated.Global
 public class PollyExtensions
 {
-    private static string RetryAfterHeader = "x-ms-ratelimit-microsoft.costmanagement-clienttype-retry-after";
-    private static string RetryAfterHeader2 = "x-ms-ratelimit-microsoft.costmanagement-entity-retry-after";
-
     public static IAsyncPolicy<HttpResponseMessage> GetRetryAfterPolicy()
     {
         return Policy.HandleResult<HttpResponseMessage>
                 (msg => msg.StatusCode == HttpStatusCode.TooManyRequests)
             .WaitAndRetryAsync(
-                retryCount: 3,
+                retryCount: 5,
                 sleepDurationProvider: (_, response, _) =>
-                    response.Result.Headers.TryGetValues(RetryAfterHeader,
-                        out var seconds)
-                        ? TimeSpan.FromSeconds(int.Parse(seconds.First()))
-                        :  response.Result.Headers.TryGetValues(RetryAfterHeader2,
-                            out var seconds2)
-                            ? TimeSpan.FromSeconds(int.Parse(seconds2.First())): TimeSpan.FromSeconds(5),
+                {
+                    var retryAfterHeader = 
+                        response.Result.Headers.FirstOrDefault(h => h.Key.ToLowerInvariant().Contains("retry-after"));
+                    return retryAfterHeader.Key != null && int.TryParse(retryAfterHeader.Value.First(), out var seconds)
+                        ? TimeSpan.FromSeconds(seconds)
+                        : TimeSpan.FromSeconds(5);
+                },
                 onRetryAsync: (msg, time, retries, context) => Task.CompletedTask
             );
     }