ResourceThrottleRetryPolicy: Fixes cumulativeRetryDelay tracking when x-ms-retry-after-ms header is absent

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;
             }
 

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;