[Service Bus] Retry when throttled with short try (Azure#50394)

* [Service Bus] Retry when throttled with short try

The focus of these changes is to improve the retry
logic in the Service Bus SDK when throttling occurs.
Behavior has been adjusted to allow short `TryTimeout`
configurations to consume retries when waiting
for service throttling rather than triggering an
immediate `ServiceBusy` exception.

Author: jsquire

sdk/servicebus/Azure.Messaging.ServiceBus/CHANGELOG.md

@@ -12,6 +12,8 @@
 
 ### Other Changes
 
+- Updated retry policy behavior when the service is throttling and the `TryTimeout` is shorter than the standard throttling time of 30 seconds.  Previously, the operation was immediately canceled with a server busy exception.  With these changes, the operation will begin consuming retry attempts while throttling until either the server busy state is cleared or all configured retry attempts are exhausted.  ([#50121](https://github.com/Azure/azure-sdk-for-net/issues/50121))
+
 ## 7.19.0 (2025-04-08)
 
 ### Features Added

sdk/servicebus/Azure.Messaging.ServiceBus/src/Core/BasicRetryPolicy.cs

@@ -29,7 +29,7 @@ internal class BasicRetryPolicy : ServiceBusRetryPolicy
         private static readonly ThreadLocal<Random> RandomNumberGenerator = new ThreadLocal<Random>(() => new Random(Interlocked.Increment(ref s_randomSeed)), false);
 
         /// <summary>The maximum number of seconds allowed for a <see cref="TimeSpan" />.</summary>
-        private static double MaximumTimeSpanSeconds = TimeSpan.MaxValue.TotalSeconds;
+        private static readonly double MaximumTimeSpanSeconds = TimeSpan.MaxValue.TotalSeconds;
 
         /// <summary>
         ///   The set of options responsible for configuring the retry

sdk/servicebus/Azure.Messaging.ServiceBus/src/Primitives/ServiceBusRetryPolicy.cs

@@ -35,6 +35,9 @@ public abstract class ServiceBusRetryPolicy
         private const int ServerNotBusyState = 0; // default value of serverBusy
         private const int ServerBusyState = 1;
 
+        /// <summary>A generic retriable busy exception to use for delay calculations.</summary>
+        private static readonly ServiceBusException DefaultServiceBusyException = new ServiceBusException(Resources.DefaultServerBusyException, ServiceBusFailureReason.ServiceBusy);
+
         /// <summary>
         /// Determines whether or not the server returned a busy error.
         /// </summary>
@@ -130,17 +133,42 @@ internal async ValueTask<TResult> RunOperation<T1, TResult>(
             bool logTimeoutRetriesAsVerbose = false)
       {
             var failedAttemptCount = 0;
+            var tryTimeout = CalculateTryTimeout(0);
 
-            TimeSpan tryTimeout = CalculateTryTimeout(0);
             if (IsServerBusy && tryTimeout < ServerBusyBaseSleepTime)
             {
-                // We are in a server busy state before we start processing.
-                // Since ServerBusyBaseSleepTime > remaining time for the operation, we don't wait for the entire Sleep time.
-                await Task.Delay(tryTimeout, cancellationToken).ConfigureAwait(false);
-                throw new ServiceBusException(
-                    ServerBusyExceptionMessage,
-                    ServiceBusFailureReason.ServiceBusy);
+                while (IsServerBusy && !cancellationToken.IsCancellationRequested)
+                {
+                    // If we are in a server busy state, we will wait for the try timeout.
+
+                    await Task.Delay(tryTimeout, cancellationToken).ConfigureAwait(false);
+
+                    // If the server is still busy, consider this a retry attempt and wait for the
+                    // calculated retry delay before trying again.
+
+                    if (IsServerBusy)
+                    {
+                        ++failedAttemptCount;
+                        var delay = CalculateRetryDelay(DefaultServiceBusyException, failedAttemptCount);
+
+                        if (delay.HasValue)
+                        {
+                            Logger.RunOperationExceptionEncountered(DefaultServiceBusyException.ToString());
+
+                            await Task.Delay(delay.Value, cancellationToken).ConfigureAwait(false);
+                            tryTimeout = CalculateTryTimeout(failedAttemptCount);
+                        }
+                        else
+                        {
+                            // If there are no retries left, then fail because the server is busy.
+                            throw new ServiceBusException(
+                                ServerBusyExceptionMessage,
+                                ServiceBusFailureReason.ServiceBusy);
+                        }
+                    }
+                }
             }
+
             while (!cancellationToken.IsCancellationRequested)
             {
                 if (IsServerBusy)

sdk/servicebus/Azure.Messaging.ServiceBus/tests/Primitives/ServiceBusRetryPolicyTests.cs

@@ -102,6 +102,67 @@ await policy.RunOperation((state, timeout, token) =>
             Assert.That(policy.IsServerBusy, Is.False);
         }
 
+        [Test]
+        public void RunOperationServerBusyTryTimeoutShorterThanBaseWaitTimeThrowsAfterRetries()
+        {
+            var policy = new CustomServerBusyMockRetryPolicy
+            {
+                ServerBusyBaseSleepTime = TimeSpan.FromSeconds(30)
+            };
+
+            // Set the server busy state before running the operation
+            policy.SetServerBusyForTest();
+
+            var operationCallCount = 0;
+
+            // The operation is never actually called in the busy block, but we provide a dummy
+            Func<object, TimeSpan, CancellationToken, ValueTask<bool>> operation = (state, timeout, token) =>
+            {
+                operationCallCount++;
+                return new ValueTask<bool>(true);
+            };
+
+            // Should throw after retries are exhausted
+            var ex = Assert.ThrowsAsync<ServiceBusException>(async () =>
+            {
+                await policy.RunOperation(operation, null, null, CancellationToken.None);
+            });
+
+            Assert.That(ex.Reason, Is.EqualTo(ServiceBusFailureReason.ServiceBusy));
+            Assert.That(policy.CalculateRetryDelayCallCount, Is.EqualTo(CustomServerBusyMockRetryPolicy.MaxRetries + 1));
+        }
+
+        [Test]
+        public async Task RunOperationServerBusyResolvesBeforeRetriesInvokesOperation()
+        {
+            var policy = new ServerBusyResolvesMockRetryPolicy
+            {
+                ServerBusyBaseSleepTime = TimeSpan.FromMilliseconds(10)
+            };
+
+            // Set the server busy state before running the operation
+            policy.SetServerBusyForTest();
+
+            var operationCallCount = 0;
+            var operationInvoked = false;
+
+            Func<object, TimeSpan, CancellationToken, ValueTask<bool>> operation = (state, timeout, token) =>
+            {
+                ++operationCallCount;
+                operationInvoked = true;
+                return new ValueTask<bool>(true);
+            };
+
+            // Should eventually invoke the operation and return true
+            var result = await policy.RunOperation(operation, null, null, CancellationToken.None);
+
+            Assert.That(result, Is.True);
+            Assert.That(operationInvoked, Is.True);
+            Assert.That(operationCallCount, Is.EqualTo(1));
+            Assert.That(policy.CalculateRetryDelayCallCount, Is.EqualTo(ServerBusyResolvesMockRetryPolicy.BusyResolveAfterRetries + 1));
+        }
+
+        // Private test helper classes as nested types
         private class MockServiceBusRetryPolicy : ServiceBusRetryPolicy
         {
             public override TimeSpan CalculateTryTimeout(int attemptCount)
@@ -114,5 +175,81 @@ public override TimeSpan CalculateTryTimeout(int attemptCount)
                 return null;
             }
         }
+
+        private class CustomServerBusyMockRetryPolicy : ServiceBusRetryPolicy
+        {
+            public const int MaxRetries = 3;
+
+            private int _retryCount = 0;
+
+            public int CalculateRetryDelayCallCount { get; private set; }
+
+            public void SetServerBusyForTest()
+            {
+                // Set the private _serverBusyState to ServerBusyState (1)
+                var field = typeof(ServiceBusRetryPolicy).GetField("_serverBusyState", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
+                field.SetValue(this, 1);
+            }
+
+            public override TimeSpan CalculateTryTimeout(int attemptCount)
+            {
+                // Always return 1 second, which is less than the base sleep time
+                return TimeSpan.FromSeconds(1);
+            }
+
+            public override TimeSpan? CalculateRetryDelay(Exception lastException, int attemptCount)
+            {
+                ++CalculateRetryDelayCallCount;
+
+                if (++_retryCount <= MaxRetries)
+                {
+                    return TimeSpan.FromMilliseconds(10);
+                }
+
+                return null;
+            }
+        }
+
+        private class ServerBusyResolvesMockRetryPolicy : ServiceBusRetryPolicy
+        {
+            public const int BusyResolveAfterRetries = 2;
+
+            private int _retryCount = 0;
+            private bool _serverBusyCleared = false;
+
+            public int CalculateRetryDelayCallCount { get; private set; }
+
+            public void SetServerBusyForTest(bool isBusy = true)
+            {
+                var field = typeof(ServiceBusRetryPolicy).GetField("_serverBusyState", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
+                field.SetValue(this, isBusy ? 1 : 0);
+            }
+
+            public override TimeSpan CalculateTryTimeout(int attemptCount)
+            {
+                return TimeSpan.FromMilliseconds(1);
+            }
+
+            public override TimeSpan? CalculateRetryDelay(Exception lastException, int attemptCount)
+            {
+                ++CalculateRetryDelayCallCount;
+
+                if (++_retryCount <= BusyResolveAfterRetries)
+                {
+                    return TimeSpan.FromMilliseconds(1);
+                }
+
+                // Simulate server busy resolving
+                if (!_serverBusyCleared)
+                {
+                    SetServerBusyForTest(false); // Clear the server busy state
+                    _serverBusyCleared = true;
+
+                    return TimeSpan.FromMilliseconds(1);
+                }
+
+                return null;
+            }
+        }
     }
 }