diff --git a/csharp/src/Drivers/Databricks/RetryHttpHandler.cs b/csharp/src/Drivers/Databricks/RetryHttpHandler.cs index 736c51a430..0f7f0b39e4 100644 --- a/csharp/src/Drivers/Databricks/RetryHttpHandler.cs +++ b/csharp/src/Drivers/Databricks/RetryHttpHandler.cs @@ -25,17 +25,19 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks { /// - /// 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. /// internal class RetryHttpHandler : DelegatingHandler { private readonly int _retryTimeoutSeconds; + private readonly int _initialBackoffSeconds = 1; + private readonly int _maxBackoffSeconds = 32; /// /// Initializes a new instance of the class. /// /// The inner handler to delegate to. - /// Whether retry behavior is enabled. /// Maximum total time in seconds to retry before failing. public RetryHttpHandler(HttpMessageHandler innerHandler, int retryTimeoutSeconds) : base(innerHandler) @@ -44,7 +46,7 @@ public RetryHttpHandler(HttpMessageHandler innerHandler, int retryTimeoutSeconds } /// - /// 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. /// protected override async Task SendAsync( HttpRequestMessage request, @@ -58,6 +60,8 @@ protected override async Task SendAsync( HttpResponseMessage response; string? lastErrorMessage = null; DateTime startTime = DateTime.UtcNow; + int attemptCount = 0; + int currentBackoffSeconds = _initialBackoffSeconds; int totalRetrySeconds = 0; do @@ -70,28 +74,48 @@ protected override async Task 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(); @@ -99,16 +123,19 @@ protected override async Task SendAsync( // 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 SendAsync( .SetSqlState("08001"); } + /// + /// Determines if the status code is one that should be retried. + /// + private bool IsRetryableStatusCode(HttpStatusCode statusCode) + { + return statusCode == HttpStatusCode.RequestTimeout || // 408 + statusCode == HttpStatusCode.BadGateway || // 502 + statusCode == HttpStatusCode.ServiceUnavailable || // 503 + statusCode == HttpStatusCode.GatewayTimeout; // 504 + } + + /// + /// Calculates backoff time with jitter to avoid thundering herd problem. + /// + 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); + } + /// /// 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. /// private static async Task CloneHttpContentAsync(HttpContent content) { diff --git a/csharp/test/Drivers/Databricks/Unit/RetryHttpHandlerTest.cs b/csharp/test/Drivers/Databricks/Unit/RetryHttpHandlerTest.cs index 0aa8ba55ad..d21fe04e6d 100644 --- a/csharp/test/Drivers/Databricks/Unit/RetryHttpHandlerTest.cs +++ b/csharp/test/Drivers/Databricks/Unit/RetryHttpHandlerTest.cs @@ -97,10 +97,10 @@ public async Task RetryAfterHandlerThrowsWhenTimeoutExceeded() } /// - /// Tests that the RetryHttpHandler handles non-503 responses correctly. + /// Tests that the RetryHttpHandler handles non-retryable responses correctly. /// [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() } /// - /// 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. /// [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 } /// - /// Tests that the RetryHttpHandler handles invalid Retry-After headers correctly. + /// Tests that the RetryHttpHandler handles invalid Retry-After headers by using exponential backoff. /// [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 + } + + /// + /// Tests that the RetryHttpHandler properly processes retryable status codes. + /// + [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 + } + + /// + /// Tests that the RetryHttpHandler properly handles multiple retries with exponential backoff. + /// + [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 + } + + /// + /// Tests that the RetryHttpHandler throws an exception when the server keeps returning errors + /// and we reach the timeout with exponential backoff. + /// + [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(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}"); } ///