diff --git a/Microsoft.Azure.Cosmos/src/ResourceThrottleRetryPolicy.cs b/Microsoft.Azure.Cosmos/src/ResourceThrottleRetryPolicy.cs index 8dae27d467..4f8cfa84de 100644 --- a/Microsoft.Azure.Cosmos/src/ResourceThrottleRetryPolicy.cs +++ b/Microsoft.Azure.Cosmos/src/ResourceThrottleRetryPolicy.cs @@ -159,16 +159,16 @@ private bool CheckIfRetryNeeded( retryDelay = TimeSpan.FromTicks(retryDelay.Ticks * this.backoffDelayFactor); } + if (retryDelay == TimeSpan.Zero) + { + // we should never reach here as BE should turn non-zero of retryDelay + DefaultTrace.TraceInformation("Received retryDelay of 0 with Http 429: {0}", retryAfter); + retryDelay = TimeSpan.FromSeconds(DefaultRetryInSeconds); + } + if (retryDelay < this.maxWaitTimeInMilliseconds && this.maxWaitTimeInMilliseconds >= (this.cumulativeRetryDelay = retryDelay.Add(this.cumulativeRetryDelay))) { - if (retryDelay == TimeSpan.Zero) - { - // we should never reach here as BE should turn non-zero of retryDelay - DefaultTrace.TraceInformation("Received retryDelay of 0 with Http 429: {0}", retryAfter); - retryDelay = TimeSpan.FromSeconds(DefaultRetryInSeconds); - } - return true; } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ResourceThrottleRetryPolicyTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ResourceThrottleRetryPolicyTests.cs index dca4796b18..c6d6fc5ec8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ResourceThrottleRetryPolicyTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ResourceThrottleRetryPolicyTests.cs @@ -7,8 +7,11 @@ namespace Microsoft.Azure.Cosmos.Tests using System; using System.Collections.Generic; using System.Diagnostics; + using System.Net; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.Core.Trace; + using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; [TestClass] @@ -63,6 +66,75 @@ public async Task DoesSerializeExceptionOnTracingEnabled() Assert.AreEqual(1, exception.ToStringCount, "Exception was not serialized"); } + [TestMethod] + public async Task MissingRetryAfterHeader_RespectsMaxWaitTime() + { + // maxWaitTime = 12 seconds, maxAttempts = 9 + // With 5-second default fallback, only 2 retries should fit within 12s (5+5=10 <= 12, 5+5+5=15 > 12) + int maxAttempts = 9; + int maxWaitTimeInSeconds = 12; + ResourceThrottleRetryPolicy policy = new ResourceThrottleRetryPolicy( + maxAttempts, + maxWaitTimeInSeconds); + + int retryCount = 0; + for (int i = 0; i < maxAttempts + 1; i++) + { + ResponseMessage throttledResponse = ResourceThrottleRetryPolicyTests.CreateThrottledResponseWithoutRetryAfter(); + ShouldRetryResult result = await policy.ShouldRetryAsync(throttledResponse, CancellationToken.None); + + if (result.ShouldRetry) + { + retryCount++; + Assert.AreEqual(TimeSpan.FromSeconds(5), result.BackoffTime, + $"Retry {retryCount}: Expected 5-second fallback delay"); + } + else + { + break; + } + } + + // With 12s budget and 5s per retry: 5s + 5s = 10s (ok), 5s + 5s + 5s = 15s (exceeds 12s) + Assert.AreEqual(2, retryCount, + "Expected exactly 2 retries within 12-second budget with 5-second fallback delays"); + } + + [TestMethod] + public async Task PresentRetryAfterHeader_UseServerProvidedDelay() + { + int maxAttempts = 9; + int maxWaitTimeInSeconds = 30; + ResourceThrottleRetryPolicy policy = new ResourceThrottleRetryPolicy( + maxAttempts, + maxWaitTimeInSeconds); + + TimeSpan serverRetryAfter = TimeSpan.FromMilliseconds(1000); + ResponseMessage throttledResponse = ResourceThrottleRetryPolicyTests.CreateThrottledResponseWithRetryAfter(serverRetryAfter); + + ShouldRetryResult result = await policy.ShouldRetryAsync(throttledResponse, CancellationToken.None); + + Assert.IsTrue(result.ShouldRetry); + Assert.AreEqual(serverRetryAfter, result.BackoffTime, + "Should use server-provided RetryAfter delay"); + } + + private static ResponseMessage CreateThrottledResponseWithoutRetryAfter() + { + ResponseMessage response = new ResponseMessage((HttpStatusCode)429); + // Do NOT set RetryAfterInMilliseconds header — simulating missing x-ms-retry-after-ms + return response; + } + + private static ResponseMessage CreateThrottledResponseWithRetryAfter(TimeSpan retryAfter) + { + ResponseMessage response = new ResponseMessage((HttpStatusCode)429); + response.Headers.Set( + HttpConstants.HttpHeaders.RetryAfterInMilliseconds, + retryAfter.TotalMilliseconds.ToString(System.Globalization.CultureInfo.InvariantCulture)); + return response; + } + private class CustomException : Exception { public int ToStringCount { get; private set; } = 0;