fix(csharp/test/Drivers/Databricks): Enrich RetryHttpHandler with other status codes

csharp/src/Drivers/Databricks/RetryHttpHandler.cs

@@ -25,17 +25,19 @@
 namespace Apache.Arrow.Adbc.Drivers.Databricks
 {
     /// <summary>
-    /// HTTP handler that implements retry behavior for 503 responses with Retry-After headers.
+    /// HTTP handler that implements retry behavior for 408, 502, 503, and 504 responses.
+    /// Uses Retry-After header if present, otherwise uses exponential backoff.
     /// </summary>
     internal class RetryHttpHandler : DelegatingHandler
     {
         private readonly int _retryTimeoutSeconds;
+        private readonly int _initialBackoffSeconds = 1;
+        private readonly int _maxBackoffSeconds = 32;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="RetryHttpHandler"/> class.
         /// </summary>
         /// <param name="innerHandler">The inner handler to delegate to.</param>
-        /// <param name="retryEnabled">Whether retry behavior is enabled.</param>
         /// <param name="retryTimeoutSeconds">Maximum total time in seconds to retry before failing.</param>
         public RetryHttpHandler(HttpMessageHandler innerHandler, int retryTimeoutSeconds)
             : base(innerHandler)
@@ -44,7 +46,7 @@ public RetryHttpHandler(HttpMessageHandler innerHandler, int retryTimeoutSeconds
         }
 
         /// <summary>
-        /// Sends an HTTP request to the inner handler with retry logic for 503 responses.
+        /// Sends an HTTP request to the inner handler with retry logic for retryable status codes.
         /// </summary>
         protected override async Task<HttpResponseMessage> SendAsync(
             HttpRequestMessage request,
@@ -58,6 +60,8 @@ protected override async Task<HttpResponseMessage> SendAsync(
             HttpResponseMessage response;
             string? lastErrorMessage = null;
             DateTime startTime = DateTime.UtcNow;
+            int attemptCount = 0;
+            int currentBackoffSeconds = _initialBackoffSeconds;
             int totalRetrySeconds = 0;
 
             do
@@ -70,45 +74,68 @@ protected override async Task<HttpResponseMessage> SendAsync(
 
                 response = await base.SendAsync(request, cancellationToken);
 
-                // If it's not a 503 response, return immediately
-                if (response.StatusCode != HttpStatusCode.ServiceUnavailable)
+                // If it's not a retryable status code, return immediately
+                if (!IsRetryableStatusCode(response.StatusCode))
                 {
                     return response;
                 }
 
-                // Check for Retry-After header
-                if (!response.Headers.TryGetValues("Retry-After", out var retryAfterValues))
+                attemptCount++;
+
+                // Check if we've exceeded the timeout
+                TimeSpan elapsedTime = DateTime.UtcNow - startTime;
+                if (_retryTimeoutSeconds > 0 && elapsedTime.TotalSeconds > _retryTimeoutSeconds)
                 {
-                    // No Retry-After header, so return the response as is
-                    return response;
+                    // We've exceeded the timeout, so break out of the loop
+                    break;
                 }
 
-                // Parse the Retry-After value
-                string retryAfterValue = string.Join(",", retryAfterValues);
-                if (!int.TryParse(retryAfterValue, out int retryAfterSeconds) || retryAfterSeconds <= 0)
+                int waitSeconds;
+
+                // Check for Retry-After header
+                if (response.Headers.TryGetValues("Retry-After", out var retryAfterValues))
                 {
-                    // Invalid Retry-After value, return the response as is
-                    return response;
+                    // Parse the Retry-After value
+                    string retryAfterValue = string.Join(",", retryAfterValues);
+                    if (int.TryParse(retryAfterValue, out int retryAfterSeconds) && retryAfterSeconds > 0)
+                    {
+                        // Use the Retry-After value
+                        waitSeconds = retryAfterSeconds;
+                        lastErrorMessage = $"Service temporarily unavailable (HTTP {(int)response.StatusCode}). Using server-specified retry after {waitSeconds} seconds. Attempt {attemptCount}.";
+                    }
+                    else
+                    {
+                        // Invalid Retry-After value, use exponential backoff
+                        waitSeconds = CalculateBackoffWithJitter(currentBackoffSeconds);
+                        lastErrorMessage = $"Service temporarily unavailable (HTTP {(int)response.StatusCode}). Invalid Retry-After header, using exponential backoff of {waitSeconds} seconds. Attempt {attemptCount}.";
+                    }
+                }
+                else
+                {
+                    // No Retry-After header, use exponential backoff
+                    waitSeconds = CalculateBackoffWithJitter(currentBackoffSeconds);
+                    lastErrorMessage = $"Service temporarily unavailable (HTTP {(int)response.StatusCode}). Using exponential backoff of {waitSeconds} seconds. Attempt {attemptCount}.";
                 }
-
-                lastErrorMessage = $"Service temporarily unavailable (HTTP 503). Retry after {retryAfterSeconds} seconds.";
 
                 // Dispose the response before retrying
                 response.Dispose();
 
                 // Reset the request content for the next attempt
                 request.Content = null;
 
-                // Check if we've exceeded the timeout
-                totalRetrySeconds += retryAfterSeconds;
+                // Update total retry time
+                totalRetrySeconds += waitSeconds;
                 if (_retryTimeoutSeconds > 0 && totalRetrySeconds > _retryTimeoutSeconds)
                 {
                     // We've exceeded the timeout, so break out of the loop
                     break;
                 }
 
-                // Wait for the specified retry time
-                await Task.Delay(TimeSpan.FromSeconds(retryAfterSeconds), cancellationToken);
+                // Wait for the calculated time
+                await Task.Delay(TimeSpan.FromSeconds(waitSeconds), cancellationToken);
+
+                // Increase backoff for next attempt (exponential backoff)
+                currentBackoffSeconds = Math.Min(currentBackoffSeconds * 2, _maxBackoffSeconds);
             } while (!cancellationToken.IsCancellationRequested);
 
             // If we get here, we've either exceeded the timeout or been cancelled
@@ -121,10 +148,32 @@ protected override async Task<HttpResponseMessage> SendAsync(
                 .SetSqlState("08001");
         }
 
+        /// <summary>
+        /// Determines if the status code is one that should be retried.
+        /// </summary>
+        private bool IsRetryableStatusCode(HttpStatusCode statusCode)
+        {
+            return statusCode == HttpStatusCode.RequestTimeout ||      // 408
+                   statusCode == HttpStatusCode.BadGateway ||          // 502
+                   statusCode == HttpStatusCode.ServiceUnavailable ||  // 503
+                   statusCode == HttpStatusCode.GatewayTimeout;        // 504
+        }
+
+        /// <summary>
+        /// Calculates backoff time with jitter to avoid thundering herd problem.
+        /// </summary>
+        private int CalculateBackoffWithJitter(int baseBackoffSeconds)
+        {
+            // Add jitter by randomizing between 80-120% of the base backoff time
+            Random random = new Random();
+            double jitterFactor = 0.8 + (random.NextDouble() * 0.4); // Between 0.8 and 1.2
+            return (int)Math.Max(1, baseBackoffSeconds * jitterFactor);
+        }
+
         /// <summary>
         /// Clones an HttpContent object so it can be reused for retries.
-        /// per .net guidance, we should not reuse the http content across multiple
-        /// request, as it maybe disposed.
+        /// Per .NET guidance, we should not reuse the HTTP content across multiple
+        /// requests, as it may be disposed.
         /// </summary>
         private static async Task<HttpContent> CloneHttpContentAsync(HttpContent content)
         {

csharp/test/Drivers/Databricks/Unit/RetryHttpHandlerTest.cs

@@ -97,10 +97,10 @@ public async Task RetryAfterHandlerThrowsWhenTimeoutExceeded()
         }
 
         /// <summary>
-        /// Tests that the RetryHttpHandler handles non-503 responses correctly.
+        /// Tests that the RetryHttpHandler handles non-retryable responses correctly.
         /// </summary>
         [Fact]
-        public async Task RetryAfterHandlerHandlesNon503Response()
+        public async Task RetryAfterHandlerHandlesNonRetryableResponse()
         {
             // Create a mock handler that returns a 404 response
             var mockHandler = new MockHttpMessageHandler(
@@ -125,10 +125,10 @@ public async Task RetryAfterHandlerHandlesNon503Response()
         }
 
         /// <summary>
-        /// Tests that the RetryHttpHandler handles 503 responses without Retry-After headers correctly.
+        /// Tests that the RetryHttpHandler handles 503 responses without Retry-After headers using exponential backoff.
         /// </summary>
         [Fact]
-        public async Task RetryAfterHandlerHandles503WithoutRetryAfterHeader()
+        public async Task RetryHandlerUsesExponentialBackoffFor503WithoutRetryAfterHeader()
         {
             // Create a mock handler that returns a 503 response without a Retry-After header
             var mockHandler = new MockHttpMessageHandler(
@@ -143,20 +143,26 @@ public async Task RetryAfterHandlerHandles503WithoutRetryAfterHeader()
             // Create an HttpClient with our handler
             var httpClient = new HttpClient(retryHandler);
 
+            // Set the mock handler to return a success response after the second retry
+            mockHandler.SetResponseAfterRetryCount(2, new HttpResponseMessage(HttpStatusCode.OK)
+            {
+                Content = new StringContent("Success")
+            });
+
             // Send a request
             var response = await httpClient.GetAsync("http://test.com");
 
-            // Verify the response is 503
-            Assert.Equal(HttpStatusCode.ServiceUnavailable, response.StatusCode);
-            Assert.Equal("Service Unavailable", await response.Content.ReadAsStringAsync());
-            Assert.Equal(1, mockHandler.RequestCount); // Only the initial request, no retries
+            // Verify the response is OK
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+            Assert.Equal("Success", await response.Content.ReadAsStringAsync());
+            Assert.Equal(3, mockHandler.RequestCount); // Initial request + 2 retries
         }
 
         /// <summary>
-        /// Tests that the RetryHttpHandler handles invalid Retry-After headers correctly.
+        /// Tests that the RetryHttpHandler handles invalid Retry-After headers by using exponential backoff.
         /// </summary>
         [Fact]
-        public async Task RetryAfterHandlerHandlesInvalidRetryAfterHeader()
+        public async Task RetryHandlerUsesExponentialBackoffForInvalidRetryAfterHeader()
         {
             // Create a mock handler that returns a 503 response with an invalid Retry-After header
             var mockHandler = new MockHttpMessageHandler(
@@ -173,6 +179,12 @@ public async Task RetryAfterHandlerHandlesInvalidRetryAfterHeader()
             response.Headers.TryAddWithoutValidation("Retry-After", "invalid");
             mockHandler.SetResponseAfterRetryCount(0, response);
 
+            // Set the mock handler to return a success response after the first retry
+            mockHandler.SetResponseAfterRetryCount(1, new HttpResponseMessage(HttpStatusCode.OK)
+            {
+                Content = new StringContent("Success")
+            });
+
             // Create the RetryHttpHandler with retry enabled
             var retryHandler = new RetryHttpHandler(mockHandler, 5);
 
@@ -182,10 +194,118 @@ public async Task RetryAfterHandlerHandlesInvalidRetryAfterHeader()
             // Send a request
             response = await httpClient.GetAsync("http://test.com");
 
-            // Verify the response is 503
-            Assert.Equal(HttpStatusCode.ServiceUnavailable, response.StatusCode);
-            Assert.Equal("Service Unavailable", await response.Content.ReadAsStringAsync());
-            Assert.Equal(1, mockHandler.RequestCount); // Only the initial request, no retries
+            // Verify the response is OK
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+            Assert.Equal("Success", await response.Content.ReadAsStringAsync());
+            Assert.Equal(2, mockHandler.RequestCount); // Initial request + 1 retry
+        }
+
+        /// <summary>
+        /// Tests that the RetryHttpHandler properly processes retryable status codes.
+        /// </summary>
+        [Theory]
+        [InlineData(HttpStatusCode.RequestTimeout, "Request Timeout")]      // 408
+        [InlineData(HttpStatusCode.BadGateway, "Bad Gateway")]              // 502
+        [InlineData(HttpStatusCode.ServiceUnavailable, "Service Unavailable")] // 503
+        [InlineData(HttpStatusCode.GatewayTimeout, "Gateway Timeout")]      // 504
+        public async Task RetryHandlerProcessesRetryableStatusCodes(HttpStatusCode statusCode, string errorMessage)
+        {
+            // Create a mock handler that returns the specified status code
+            var mockHandler = new MockHttpMessageHandler(
+                new HttpResponseMessage(statusCode)
+                {
+                    Content = new StringContent(errorMessage)
+                });
+
+            // Create the RetryHttpHandler with retry enabled
+            var retryHandler = new RetryHttpHandler(mockHandler, 5);
+
+            // Create an HttpClient with our handler
+            var httpClient = new HttpClient(retryHandler);
+
+            // Set the mock handler to return a success response after the first retry
+            mockHandler.SetResponseAfterRetryCount(1, new HttpResponseMessage(HttpStatusCode.OK)
+            {
+                Content = new StringContent("Success")
+            });
+
+            // Send a request
+            var response = await httpClient.GetAsync("http://test.com");
+
+            // Verify the response is OK
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+            Assert.Equal("Success", await response.Content.ReadAsStringAsync());
+            Assert.Equal(2, mockHandler.RequestCount); // Initial request + 1 retry
+        }
+
+        /// <summary>
+        /// Tests that the RetryHttpHandler properly handles multiple retries with exponential backoff.
+        /// </summary>
+        [Fact]
+        public async Task RetryHandlerHandlesMultipleRetriesWithExponentialBackoff()
+        {
+            // Create a mock handler that returns a 503 response without a Retry-After header
+            var mockHandler = new MockHttpMessageHandler(
+                new HttpResponseMessage(HttpStatusCode.ServiceUnavailable)
+                {
+                    Content = new StringContent("Service Unavailable")
+                });
+
+            // Create the RetryHttpHandler with retry enabled and a generous timeout
+            var retryHandler = new RetryHttpHandler(mockHandler, 10);
+
+            // Create an HttpClient with our handler
+            var httpClient = new HttpClient(retryHandler);
+
+            // Set the mock handler to return a success response after the third retry
+            mockHandler.SetResponseAfterRetryCount(3, new HttpResponseMessage(HttpStatusCode.OK)
+            {
+                Content = new StringContent("Success")
+            });
+
+            // Send a request
+            var response = await httpClient.GetAsync("http://test.com");
+
+            // Verify the response is OK
+            Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+            Assert.Equal("Success", await response.Content.ReadAsStringAsync());
+            Assert.Equal(4, mockHandler.RequestCount); // Initial request + 3 retries
+        }
+
+        /// <summary>
+        /// Tests that the RetryHttpHandler throws an exception when the server keeps returning errors
+        /// and we reach the timeout with exponential backoff.
+        /// </summary>
+        [Theory]
+        [InlineData(HttpStatusCode.RequestTimeout)]      // 408
+        [InlineData(HttpStatusCode.BadGateway)]          // 502
+        [InlineData(HttpStatusCode.ServiceUnavailable)]  // 503
+        [InlineData(HttpStatusCode.GatewayTimeout)]      // 504
+        public async Task RetryHandlerThrowsWhenServerNeverRecovers(HttpStatusCode statusCode)
+        {
+            // Create a mock handler that always returns the error status code
+            var mockHandler = new MockHttpMessageHandler(
+                new HttpResponseMessage(statusCode)
+                {
+                    Content = new StringContent($"Error: {statusCode}")
+                });
+
+            // Create the RetryHttpHandler with a short timeout to make the test run faster
+            var retryHandler = new RetryHttpHandler(mockHandler, 3);
+
+            // Create an HttpClient with our handler
+            var httpClient = new HttpClient(retryHandler);
+
+            // Send a request and expect a DatabricksException
+            var exception = await Assert.ThrowsAsync<DatabricksException>(async () =>
+                await httpClient.GetAsync("http://test.com"));
+
+            // Verify the exception has the correct SQL state
+            Assert.Contains("08001", exception.SqlState);
+            Assert.Equal(AdbcStatusCode.IOError, exception.Status);
+
+            // Verify we tried multiple times before giving up
+            Assert.True(mockHandler.RequestCount > 1, $"Expected multiple requests, but got {mockHandler.RequestCount}");
         }
 
         /// <summary>