fix: Retry flags requests on 502 and 504 (#255)

* fix: Retry flags requests on 502 and 504

* test: Cover flags gateway retry exhaustion

* chore: add flags retry changeset

Author: marandaneto

.changeset/gentle-flags-retry.md

@@ -0,0 +1,6 @@
+---
+'PostHog': patch
+'PostHog.AspNetCore': patch
+---
+
+Retry remote feature flag requests after transient 502 and 504 responses.

src/PostHog/Config/PostHogOptions.cs

@@ -174,7 +174,8 @@ public string? ProjectApiKey
     public int MaxRetries { get; set; } = 3;
 
     /// <summary>
-    /// The maximum number of retries for feature flag requests after transient network errors. (Default: 1)
+    /// The maximum number of retries for feature flag requests after transient network errors
+    /// or HTTP 502/504 responses. (Default: 1)
     /// Set to 0 to disable feature flag request retries.
     /// </summary>
     public int FeatureFlagRequestMaxRetries { get; set; } = 1;

src/PostHog/Library/HttpClientExtensions.cs

@@ -46,8 +46,8 @@ internal static class HttpClientExtensions
     }
 
     /// <summary>
-    /// Sends a POST request with retry logic only for network/transport failures and timeouts.
-    /// Non-successful HTTP responses are not retried.
+    /// Sends a POST request with retry logic only for network/transport failures, timeouts,
+    /// and HTTP 502/504 responses.
     /// </summary>
     public static async Task<TBody?> PostJsonWithNetworkRetryAsync<TBody>(
         this HttpClient httpClient,
@@ -141,7 +141,13 @@ async Task<bool> ShouldRetryAfterTransientFailure()
             // be caught by the retry logic above.
             using (response)
             {
-                if (isHalfOpenProbe || response.IsSuccessStatusCode)
+                var isRetryableFlagsStatusCode = ShouldRetryFlagsStatusCode(response.StatusCode);
+                if (isRetryableFlagsStatusCode && await ShouldRetryAfterTransientFailure())
+                {
+                    continue;
+                }
+
+                if ((isHalfOpenProbe && !isRetryableFlagsStatusCode) || response.IsSuccessStatusCode)
                 {
                     circuitBreaker.Close();
                 }
@@ -156,6 +162,9 @@ async Task<bool> ShouldRetryAfterTransientFailure()
         }
     }
 
+    static bool ShouldRetryFlagsStatusCode(HttpStatusCode statusCode)
+        => statusCode is HttpStatusCode.BadGateway or HttpStatusCode.GatewayTimeout;
+
     static bool IsRetryableFlagsHttpRequestException(HttpRequestException exception)
     {
         for (Exception? current = exception; current != null; current = current.InnerException)

tests/UnitTests/Library/HttpClientExtensionsTests.cs

@@ -1153,14 +1153,101 @@ public async Task RetriesOnTaskCanceledExceptionFromTimeoutThenSucceeds()
         Assert.Equal(2, handler.RequestCount);
     }
 
+    [Theory]
+    [InlineData(HttpStatusCode.BadGateway)] // 502
+    [InlineData(HttpStatusCode.GatewayTimeout)] // 504
+    public async Task RetriesOnGatewayHttpStatusCodeThenSucceeds(HttpStatusCode statusCode)
+    {
+        var handler = new FakeRetryHttpMessageHandler();
+        handler.AddResponse(statusCode, new { type = "error", detail = "server error" });
+        handler.AddResponse(HttpStatusCode.OK, new { featureFlags = new { retry_flag = true } });
+        using var httpClient = CreateHttpClient(handler);
+        var options = CreateOptions();
+        var timeProvider = new FakeTimeProvider();
+
+        var task = httpClient.PostJsonWithNetworkRetryAsync<FlagsApiResult>(
+            FlagsUrl,
+            new { api_key = "test", distinct_id = "user-1" },
+            timeProvider,
+            options,
+            new FeatureFlagRequestCircuitBreaker(),
+            CancellationToken.None);
+
+        await handler.WaitForRequestCountAsync(1);
+        Assert.Equal(1, handler.RequestCount);
+        timeProvider.Advance(TimeSpan.FromMilliseconds(1));
+        var result = await task;
+
+        Assert.NotNull(result);
+        Assert.NotNull(result!.FeatureFlags);
+        Assert.True(result.FeatureFlags!.TryGetValue("retry_flag", out var retryFlag));
+        Assert.True(retryFlag == true);
+        Assert.Equal(2, handler.RequestCount);
+    }
+
+    [Theory]
+    [InlineData(HttpStatusCode.BadGateway)] // 502
+    [InlineData(HttpStatusCode.GatewayTimeout)] // 504
+    public async Task DoesNotRetryGatewayHttpStatusCodeWhenFeatureFlagRequestMaxRetriesIsZero(HttpStatusCode statusCode)
+    {
+        var handler = new FakeRetryHttpMessageHandler();
+        handler.AddResponse(statusCode, new { type = "error", detail = "server error" });
+        handler.AddResponse(HttpStatusCode.OK, new { featureFlags = new { retry_flag = true } }); // Should never be reached
+        using var httpClient = CreateHttpClient(handler);
+        var options = CreateOptions(maxRetries: 0);
+        var timeProvider = new FakeTimeProvider();
+
+        await Assert.ThrowsAsync<ApiException>(() =>
+            httpClient.PostJsonWithNetworkRetryAsync<FlagsApiResult>(
+                FlagsUrl,
+                new { api_key = "test", distinct_id = "user-1" },
+                timeProvider,
+                options,
+                new FeatureFlagRequestCircuitBreaker(),
+                CancellationToken.None));
+
+        Assert.Equal(1, handler.RequestCount);
+    }
+
+    [Theory]
+    [InlineData(HttpStatusCode.BadGateway)] // 502
+    [InlineData(HttpStatusCode.GatewayTimeout)] // 504
+    public async Task ThrowsAfterFeatureFlagRequestMaxRetriesWhenGatewayHttpStatusCodesKeepFailing(
+        HttpStatusCode statusCode)
+    {
+        var handler = new FakeRetryHttpMessageHandler();
+        handler.AddResponse(statusCode, new { type = "error", detail = "server error" });
+        handler.AddResponse(statusCode, new { type = "error", detail = "server error" });
+        handler.AddResponse(statusCode, new { type = "error", detail = "server error" });
+        using var httpClient = CreateHttpClient(handler);
+        var options = CreateOptions(maxRetries: 2);
+        var timeProvider = new FakeTimeProvider();
+
+        var task = httpClient.PostJsonWithNetworkRetryAsync<FlagsApiResult>(
+            FlagsUrl,
+            new { api_key = "test", distinct_id = "user-1" },
+            timeProvider,
+            options,
+            new FeatureFlagRequestCircuitBreaker(),
+            CancellationToken.None);
+
+        for (var i = 1; i <= 3 && !task.IsCompleted; i++)
+        {
+            await handler.WaitForRequestCountAsync(i);
+            timeProvider.Advance(TimeSpan.FromSeconds(1));
+        }
+
+        var exception = await Assert.ThrowsAsync<ApiException>(() => task);
+        Assert.Equal(statusCode, exception.Status);
+        Assert.Equal(3, handler.RequestCount);
+    }
+
     [Theory]
     [InlineData(HttpStatusCode.RequestTimeout)] // 408
     [InlineData(HttpStatusCode.TooManyRequests)] // 429
     [InlineData(HttpStatusCode.InternalServerError)] // 500
-    [InlineData(HttpStatusCode.BadGateway)] // 502
     [InlineData(HttpStatusCode.ServiceUnavailable)] // 503
-    [InlineData(HttpStatusCode.GatewayTimeout)] // 504
-    public async Task DoesNotRetryOnHttpErrorStatusCodes(HttpStatusCode statusCode)
+    public async Task DoesNotRetryOnOtherHttpErrorStatusCodes(HttpStatusCode statusCode)
     {
         var handler = new FakeRetryHttpMessageHandler();
         handler.AddResponse(statusCode, new { type = "error", detail = "server error" });