Fix fire-and-forget background task in LocalFeatureFlagsLoader

Await the polling task during disposal so in-flight API calls complete
before resources are disposed. Follows the same pattern established in
AsyncBatchHandler (PR #158).

Fixes #159

Author: haacked

src/PostHog/Features/LocalFeatureFlagsLoader.cs

@@ -20,9 +20,11 @@ internal sealed class LocalFeatureFlagsLoader(
     IOptions<PostHogOptions> options,
     ITaskScheduler taskScheduler,
     TimeProvider timeProvider,
-    ILoggerFactory loggerFactory) : IDisposable
+    ILoggerFactory loggerFactory) : IDisposable, IAsyncDisposable
 {
     volatile int _started;
+    volatile int _disposed;
+    Task? _pollingTask;
     LocalEvaluator? _localEvaluator;
     volatile string? _etag; // ETag for conditional requests to reduce bandwidth
     readonly CancellationTokenSource _cancellationTokenSource = new();
@@ -37,7 +39,7 @@ void StartPollingIfNotStarted()
         {
             return;
         }
-        taskScheduler.Run(() => PollForFeatureFlagsAsync(_cancellationTokenSource.Token));
+        _pollingTask = taskScheduler.Run(() => PollForFeatureFlagsAsync(_cancellationTokenSource.Token));
     }
 
     /// <summary>
@@ -143,10 +145,22 @@ async Task PollForFeatureFlagsAsync(CancellationToken cancellationToken)
 
     public bool IsLoaded => _localEvaluator is not null;
 
-    public void Dispose()
+    public void Dispose() => DisposeAsync().AsTask().Wait();
+
+    public async ValueTask DisposeAsync()
     {
-        _cancellationTokenSource.Dispose();
+        if (Interlocked.Exchange(ref _disposed, 1) == 1)
+        {
+            return;
+        }
+
+        // Cancel the token so the polling loop exits, then wait for it to finish.
+        // This ensures any in-flight API call completes before we dispose resources.
+        await _cancellationTokenSource.CancelAsync();
+        await (_pollingTask ?? Task.CompletedTask);
+
         _timer.Dispose();
+        _cancellationTokenSource.Dispose();
     }
 
     public void Clear()

src/PostHog/PostHogClient.cs

@@ -691,11 +691,12 @@ public async Task LoadFeatureFlagsAsync(CancellationToken cancellationToken)
     /// <inheritdoc/>
     public async ValueTask DisposeAsync()
     {
-        // Stop the polling and wait for it.
+        // Stop background tasks first, while the API client is still alive.
+        // The polling task in _featureFlagsLoader may call the API client during shutdown.
         await _asyncBatchHandler.DisposeAsync();
+        await _featureFlagsLoader.DisposeAsync();
         _apiClient.Dispose();
         _featureFlagCalledEventCache.Dispose();
-        _featureFlagsLoader.Dispose();
     }
 
 

tests/UnitTests/Features/LocalFeatureFlagsLoaderTests.cs

@@ -0,0 +1,87 @@
+using System.Net;
+using PostHog;
+using UnitTests.Fakes;
+#if NETCOREAPP3_1
+using TestLibrary.Fakes.Polyfills;
+#endif
+
+namespace LocalFeatureFlagsLoaderTests;
+
+public class TheDisposeAsyncMethod
+{
+    const string LocalEvaluationResponse = """
+        {
+            "flags": [
+                {
+                    "key": "test-flag",
+                    "active": true,
+                    "rollout_percentage": 100,
+                    "filters": {
+                        "groups": [
+                            {
+                                "properties": [],
+                                "rollout_percentage": 100
+                            }
+                        ]
+                    }
+                }
+            ]
+        }
+        """;
+
+    static readonly Uri LocalEvaluationUrl =
+        new("https://us.i.posthog.com/api/feature_flag/local_evaluation?token=fake-project-api-key&send_cohorts");
+
+    [Fact]
+    public async Task CompletesGracefullyDuringInFlightPoll()
+    {
+        var container = new TestContainer("fake-personal-api-key");
+        var pollStarted = new TaskCompletionSource();
+        var pollCanProceed = new TaskCompletionSource();
+
+        // First response succeeds immediately (the initial load).
+        container.FakeHttpMessageHandler.AddLocalEvaluationResponse(LocalEvaluationResponse);
+
+        // Second response (the timer-triggered poll) blocks until we signal it.
+        container.FakeHttpMessageHandler.AddResponse(
+            LocalEvaluationUrl,
+            HttpMethod.Get,
+            async () =>
+            {
+                pollStarted.SetResult();
+                await pollCanProceed.Task;
+                return new HttpResponseMessage(HttpStatusCode.OK)
+                {
+                    Content = new StringContent(
+                        LocalEvaluationResponse,
+                        System.Text.Encoding.UTF8,
+                        "application/json")
+                };
+            });
+
+        var client = container.Activate<PostHogClient>();
+
+        // Initial load starts the polling loop and makes the first API call.
+        await client.LoadFeatureFlagsAsync(CancellationToken.None);
+
+        // Advance past the poll interval so the background poll fires.
+        container.FakeTimeProvider.Advance(TimeSpan.FromSeconds(31));
+
+        // Wait for the poll's API call to begin.
+        await pollStarted.Task;
+
+        // Begin disposal while the poll is mid-flight.
+        var disposeTask = client.DisposeAsync().AsTask();
+
+        // Unblock the in-flight API call so the poll can finish.
+        pollCanProceed.SetResult();
+
+        // Verify disposal completes without deadlock or exception.
+        var timeout = TimeSpan.FromSeconds(5);
+        var completed = await Task.WhenAny(disposeTask, Task.Delay(timeout));
+        if (completed != disposeTask)
+        {
+            throw new TimeoutException("DisposeAsync did not complete within 5 seconds; possible deadlock.");
+        }
+    }
+}