SCM: Add support for Retry-After header to default ClientRetryPolicy (#…

…45078) * Add support for Retry-After header to default ClientRetryPolicy * Add tests * updates * nits * support http-date format as well * update CHANGELOG * add header collection to mock
Azure · Jul 19, 2024 · e1d67ba · e1d67ba
1 parent c9fd7c6
commit e1d67ba
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 7 deletions.
diff --git a/sdk/core/System.ClientModel/CHANGELOG.md b/sdk/core/System.ClientModel/CHANGELOG.md
@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Added support for delaying retrying a request until after the interval specified on a response `Retry-After` header.
+
 ### Other Changes
 
 ## 1.1.0-beta.5 (2024-07-11)

diff --git a/sdk/core/System.ClientModel/src/Pipeline/ClientRetryPolicy.cs b/sdk/core/System.ClientModel/src/Pipeline/ClientRetryPolicy.cs
@@ -24,6 +24,7 @@ public class ClientRetryPolicy : PipelinePolicy
 
     private const int DefaultMaxRetries = 3;
     private static readonly TimeSpan DefaultInitialDelay = TimeSpan.FromSeconds(0.8);
+    private const string RetryAfterHeaderName = "Retry-After";
 
     private readonly int _maxRetries;
     private readonly TimeSpan _initialDelay;
@@ -270,8 +271,18 @@ protected virtual ValueTask<bool> ShouldRetryAsync(PipelineMessage message, Exce
     /// </returns>
     protected virtual TimeSpan GetNextDelay(PipelineMessage message, int tryCount)
     {
-        // Default implementation is exponential backoff
-        return TimeSpan.FromMilliseconds((1 << (tryCount - 1)) * _initialDelay.TotalMilliseconds);
+        // Default implementation is exponential backoff, unless the response
+        // has a retry-after header.
+        double nextDelayMilliseconds = (1 << (tryCount - 1)) * _initialDelay.TotalMilliseconds;
+
+        if (message.Response is not null &&
+            TryGetRetryAfter(message.Response, out TimeSpan retryAfter) &&
+            retryAfter.TotalMilliseconds > nextDelayMilliseconds)
+        {
+            return retryAfter;
+        }
+
+        return TimeSpan.FromMilliseconds(nextDelayMilliseconds);
     }
 
     /// <summary>
@@ -308,4 +319,26 @@ protected virtual void Wait(TimeSpan time, CancellationToken cancellationToken)
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
         }
     }
+
+    private static bool TryGetRetryAfter(PipelineResponse response, out TimeSpan value)
+    {
+        // See: https://www.rfc-editor.org/rfc/rfc7231#section-7.1.3
+        if (response.Headers.TryGetValue(RetryAfterHeaderName, out string? retryAfter))
+        {
+            if (int.TryParse(retryAfter, out var delaySeconds))
+            {
+                value = TimeSpan.FromSeconds(delaySeconds);
+                return true;
+            }
+
+            if (DateTimeOffset.TryParse(retryAfter, out DateTimeOffset retryAfterDate))
+            {
+                value = retryAfterDate - DateTimeOffset.Now;
+                return true;
+            }
+        }
+
+        value = default;
+        return false;
+    }
 }
diff --git a/sdk/core/System.ClientModel/tests/Pipeline/ClientRetryPolicyTests.cs b/sdk/core/System.ClientModel/tests/Pipeline/ClientRetryPolicyTests.cs
@@ -1,14 +1,14 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-using ClientModel.Tests;
-using ClientModel.Tests.Mocks;
-using NUnit.Framework;
 using System.ClientModel.Primitives;
 using System.Collections.Generic;
 using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
+using ClientModel.Tests;
+using ClientModel.Tests.Mocks;
+using NUnit.Framework;
 
 namespace System.ClientModel.Tests.Pipeline;
 
@@ -129,6 +129,42 @@ public async Task OnlyRetriesRetriableCodes()
         Assert.AreEqual(501, message.Response!.Status);
     }
 
+    [Test]
+    [TestCaseSource(nameof(RetryAfterTestValues))]
+    public void RespectsRetryAfterHeader(string headerName, string headerValue, double expected)
+    {
+        MockRetryPolicy retryPolicy = new();
+        MockPipelineMessage message = new();
+        MockPipelineResponse response = new();
+        response.SetHeader(headerName, headerValue);
+        message.SetResponse(response);
+
+        // Default delay with exponential backoff for second try is 1600 ms.
+        double delayMillis = retryPolicy.GetNextDelayMilliseconds(message, 2);
+        Assert.AreEqual(expected, delayMillis);
+    }
+
+    [Test]
+    public void RespectsRetryAfterDateHeader()
+    {
+        MockRetryPolicy retryPolicy = new();
+        MockPipelineMessage message = new();
+        MockPipelineResponse response = new();
+
+        // Retry after 100 seconds from now
+        response.SetHeader(
+            "Retry-After",
+            (DateTimeOffset.Now + TimeSpan.FromSeconds(100)).ToString("R"));
+        message.SetResponse(response);
+
+        // Default delay with exponential backoff for second try is 1600 ms.
+        double delayMillis = retryPolicy.GetNextDelayMilliseconds(message, 2);
+
+        // Retry-After header is larger - wait the Retry-After time, which
+        // should be approx 100s, so test for > 20s.
+        Assert.GreaterOrEqual(delayMillis, 20 * 1000);
+    }
+
     [Test]
     public async Task ShouldRetryIsCalledOnlyForErrors()
     {
@@ -356,4 +392,26 @@ public void WaitThrowsOnCancellation()
         Assert.ThrowsAsync<TaskCanceledException>(async () =>
             await retryPolicy.WaitSyncOrAsync(delay, cts.Token, IsAsync));
     }
+
+    #region Helpers
+    public static IEnumerable<object[]> RetryAfterTestValues()
+    {
+        // Retry-After header is larger - wait Retry-After time
+        yield return new object[] { "Retry-After", "5", 5000 };
+
+        // Retry-After header is smaller - wait exponential backoff time
+        yield return new object[] { "Retry-After", "1", 1600 };
+
+        // Not standard HTTP header - wait exponential backoff time
+        yield return new object[] { "retry-after-ms", "5", 1600 };
+
+        // No Retry-After header - wait exponential backoff time
+        yield return new object[] { "Content-Type", "application/json", 1600 };
+
+        // Retry-After header is smaller - wait exponential backoff
+        yield return new object[] { "Retry-After",
+            (DateTimeOffset.Now + TimeSpan.FromSeconds(1)).ToString("R"),
+            1600 };
+    }
+    #endregion
 }
diff --git a/sdk/core/System.ClientModel/tests/TestFramework/Mocks/MockPipelineResponse.cs b/sdk/core/System.ClientModel/tests/TestFramework/Mocks/MockPipelineResponse.cs
@@ -17,7 +17,7 @@ public class MockPipelineResponse : PipelineResponse
     private Stream? _contentStream;
     private BinaryData? _bufferedContent;
 
-    private readonly PipelineResponseHeaders _headers;
+    private readonly MockResponseHeaders _headers;
 
     private bool _disposed;
 
@@ -36,6 +36,9 @@ public MockPipelineResponse(int status = 0, string reasonPhrase = "")
 
     public void SetReasonPhrase(string value) => _reasonPhrase = value;
 
+    public void SetHeader(string name, string value)
+        => _headers.SetHeader(name, value);
+
     public void SetContent(byte[] content)
     {
         ContentStream = new MemoryStream(content, 0, content.Length, false, true);

diff --git a/sdk/core/System.ClientModel/tests/TestFramework/Mocks/MockPipelineTransport.cs b/sdk/core/System.ClientModel/tests/TestFramework/Mocks/MockPipelineTransport.cs
@@ -172,7 +172,7 @@ public override Stream? ContentStream
         public override BinaryData Content => throw new NotImplementedException();
 
         protected override PipelineResponseHeaders HeadersCore
-            => throw new NotImplementedException();
+            => new MockResponseHeaders();
 
         public override void Dispose() { }
 

diff --git a/sdk/core/System.ClientModel/tests/TestFramework/Mocks/MockResponseHeaders.cs b/sdk/core/System.ClientModel/tests/TestFramework/Mocks/MockResponseHeaders.cs
@@ -16,6 +16,9 @@ public MockResponseHeaders()
         _headers = new Dictionary<string, string>();
     }
 
+    public void SetHeader(string name, string value)
+        => _headers[name] = value;
+
     public override IEnumerator<KeyValuePair<string, string>> GetEnumerator()
     {
         throw new NotImplementedException();

diff --git a/sdk/core/System.ClientModel/tests/TestFramework/Mocks/MockRetryPolicy.cs b/sdk/core/System.ClientModel/tests/TestFramework/Mocks/MockRetryPolicy.cs
@@ -86,6 +86,9 @@ protected override ValueTask OnSendingRequestAsync(PipelineMessage message)
         return base.OnSendingRequestAsync(message);
     }
 
+    public double GetNextDelayMilliseconds(PipelineMessage message, int tryCount)
+        => GetNextDelay(message, tryCount).TotalMilliseconds;
+
     protected override TimeSpan GetNextDelay(PipelineMessage message, int tryCount)
     {
         if (_delayFactory is not null)