Improve PublishBuildToMaestro retries (#16020)

Author: zivkan

src/Microsoft.DotNet.Build.Tasks.Feed.Tests/PublishBuildToMaestroTests.cs

@@ -0,0 +1,86 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Net;
+using FluentAssertions;
+using Xunit;
+
+namespace Microsoft.DotNet.Build.Tasks.Feed.Tests
+{
+    public class PublishBuildToMaestroTests
+    {
+        [Fact]
+        public void GetTimeToWait_ReadsGitHubRateLimitingHeaders()
+        {
+            // Unauthenticated requests to github that are rate limited don't use the Retry-After header,
+            // so we need to be able to handle their X-RateLimit-Reset header.
+
+            DateTime resetTime = DateTime.UtcNow.AddSeconds(5.0);
+            long unixTime = (long)(resetTime - DateTime.UnixEpoch).TotalSeconds;
+            var response = new System.Net.Http.HttpResponseMessage(HttpStatusCode.Forbidden)
+            {
+                Headers =
+                {
+                    { "X-RateLimit-Remaining", "0" },
+                    { "X-RateLimit-Reset", unixTime.ToString() },
+                }
+            };
+
+            TimeSpan? actual = PublishBuildToMaestro.GetTimeToWait(response);
+
+            actual.Should().NotBeNull();
+            actual.Should().BeCloseTo(resetTime - DateTime.UtcNow, TimeSpan.FromSeconds(1));
+        }
+
+        [Fact]
+        public void GetTimeToWait_NotWaitWhenNotLimited()
+        {
+            DateTime resetTime = DateTime.UtcNow.AddSeconds(5.0);
+            long unixTime = (long)(resetTime - DateTime.UnixEpoch).TotalSeconds;
+            var response = new System.Net.Http.HttpResponseMessage(HttpStatusCode.Forbidden)
+            {
+                Headers =
+                {
+                    { "X-RateLimit-Remaining", "1" },
+                    { "X-RateLimit-Reset", unixTime.ToString() },
+                }
+            };
+
+            TimeSpan? actual = PublishBuildToMaestro.GetTimeToWait(response);
+
+            actual.Should().BeNull();
+        }
+
+        [Fact]
+        public void GetTimeToWait_UsesRetryAfterDuration()
+        {
+            var response = new System.Net.Http.HttpResponseMessage(HttpStatusCode.Forbidden)
+            {
+                Headers =
+                {
+                    { "Retry-After", "5" },
+                }
+            };
+            TimeSpan? actual = PublishBuildToMaestro.GetTimeToWait(response);
+            actual.Should().NotBeNull();
+            actual.Should().BeCloseTo(TimeSpan.FromSeconds(5), TimeSpan.FromSeconds(1));
+        }
+
+        [Fact]
+        public void GetTimeToWait_UsesRetryAfterDate()
+        {
+            DateTime resetTime = DateTime.UtcNow.AddSeconds(5.0);
+            var response = new System.Net.Http.HttpResponseMessage(HttpStatusCode.Forbidden)
+            {
+                Headers =
+                {
+                    { "Retry-After", resetTime.ToString("R") },
+                }
+            };
+            TimeSpan? actual = PublishBuildToMaestro.GetTimeToWait(response);
+            actual.Should().NotBeNull();
+            actual.Should().BeCloseTo(resetTime - DateTime.UtcNow, TimeSpan.FromSeconds(1));
+        }
+    }
+}

src/Microsoft.DotNet.Build.Tasks.Feed/src/PublishBuildToMaestro.cs

@@ -529,6 +529,12 @@ private void LookupForMatchingGitHubRepository(BuildIdentity buildIdentity)
                 string repoIdentity = string.Empty;
                 string gitHubHost = "github.com";
 
+                if (Environment.GetEnvironmentVariable("GITHUB_TOKEN") is string envGitHubToken)
+                {
+                    client.DefaultRequestHeaders.Authorization =
+                        new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", envGitHubToken);
+                }
+
                 if (!Uri.TryCreate(buildIdentity.AzureDevOpsRepository, UriKind.Absolute, out Uri repoAddr))
                 {
                     throw new Exception($"Can't parse the repository URL: {buildIdentity.AzureDevOpsRepository}");
@@ -550,26 +556,20 @@ private void LookupForMatchingGitHubRepository(BuildIdentity buildIdentity)
                 HttpResponseMessage response = client.GetAsync(url).Result;
 
                 const int MaxRetries = 5;
-                for (int retry = 1; retry <= MaxRetries && response.StatusCode == System.Net.HttpStatusCode.TooManyRequests; retry++)
+                for (int retry = 1; 
+                    retry <= MaxRetries && (response.StatusCode == System.Net.HttpStatusCode.TooManyRequests || response.StatusCode == System.Net.HttpStatusCode.Forbidden);
+                    retry++)
                 {
-                    TimeSpan timeSpan;
-                    if (response.Headers.RetryAfter?.Delta != null)
-                    {
-                        timeSpan = response.Headers.RetryAfter.Delta.Value;
-                    }
-                    else if (response.Headers.RetryAfter?.Date != null)
-                    {
-                        timeSpan = response.Headers.RetryAfter.Date.Value - DateTimeOffset.UtcNow;
-                    }
-                    else
+                    TimeSpan? timeSpan = GetTimeToWait(response);
+
+                    if (!timeSpan.HasValue)
                     {
-                        const int defaultRetryAfterSeconds = 10;
-                        timeSpan = TimeSpan.FromSeconds(defaultRetryAfterSeconds);
+                        break;
                     }
 
                     Log.LogMessage(MessageImportance.High,
-                        $"API rate limit exceeded, retrying in {timeSpan.TotalSeconds} seconds. Retry attempt: {retry}");
-                    Thread.Sleep(timeSpan);
+                        $"API rate limit exceeded, retrying in {timeSpan.Value.TotalSeconds} seconds. Retry attempt: {retry}");
+                    Thread.Sleep(timeSpan.Value);
                     response = client.GetAsync(url).Result;
                 }
 
@@ -594,6 +594,37 @@ private void LookupForMatchingGitHubRepository(BuildIdentity buildIdentity)
             }
         }
 
+        public static TimeSpan? GetTimeToWait(HttpResponseMessage response)
+        {
+            if (response.Headers.RetryAfter != null)
+            {
+                if (response.Headers.RetryAfter.Delta.HasValue)
+                {
+                    return response.Headers.RetryAfter.Delta.Value;
+                }
+                else if (response.Headers.RetryAfter.Date.HasValue)
+                {
+                    return response.Headers.RetryAfter.Date.Value - DateTimeOffset.Now;
+                }
+            }
+
+            if (response.Headers.TryGetValues("X-RateLimit-Remaining", out var values))
+            {
+                if (values.FirstOrDefault() != "0")
+                {
+                    // apparently not rate limited, so no need to wait
+                    return null;
+                }
+
+                values = response.Headers.GetValues("X-RateLimit-Reset");
+                if (long.TryParse(values.FirstOrDefault(), out long resetTime))
+                {
+                    return DateTimeOffset.FromUnixTimeSeconds(resetTime) - DateTimeOffset.Now;
+                }
+            }
+            return null;
+        }
+
         public static string GetRepoName(string repoUrl)
         {
             // In case the URL comes in ending with a '/', prevent an indexing exception