Fix fire-and-forget background task in LocalFeatureFlagsLoader

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;
+    volatile 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,27 @@ async Task PollForFeatureFlagsAsync(CancellationToken cancellationToken)
 
     public bool IsLoaded => _localEvaluator is not null;
 
-    public void Dispose()
+    public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();
+
+    public async ValueTask DisposeAsync()
     {
-        _cancellationTokenSource.Dispose();
-        _timer.Dispose();
+        if (Interlocked.Exchange(ref _disposed, 1) == 1)
+        {
+            return;
+        }
+
+        // Cancel the token so the polling loop exits, then wait for it to finish
+        // (either by completing normally or via cancellation) before disposing resources.
+        try
+        {
+            await _cancellationTokenSource.CancelAsync();
+            await (_pollingTask ?? Task.CompletedTask);
+        }
+        finally
+        {
+            _timer.Dispose();
+            _cancellationTokenSource.Dispose();
+        }
     }
 
     public void Clear()

src/PostHog/Library/AsyncBatchHandler.cs

@@ -233,7 +233,7 @@ async Task SendBatch(IReadOnlyCollection<TItem> batch)
 
     public void Dispose()
     {
-        DisposeAsync().AsTask().Wait();
+        DisposeAsync().AsTask().GetAwaiter().GetResult();
     }
 
     public async ValueTask DisposeAsync()
@@ -249,13 +249,18 @@ public async ValueTask DisposeAsync()
         // Cancel the token so both background loops exit, then wait for them to finish.
         // This ensures any in-flight flush completes and _flushing returns to 0 before
         // we attempt the final flush below.
-        await _cancellationTokenSource.CancelAsync();
-        await Task.WhenAll(_timerTask, _flushSignalTask);
-
-        _timer.Dispose();
-        _flushSignal.Dispose();
-        _cancellationTokenSource.Dispose();
-        _channel.Writer.Complete();
+        try
+        {
+            await _cancellationTokenSource.CancelAsync();
+            await Task.WhenAll(_timerTask, _flushSignalTask);
+        }
+        finally
+        {
+            _timer.Dispose();
+            _flushSignal.Dispose();
+            _cancellationTokenSource.Dispose();
+            _channel.Writer.Complete();
+        }
         try
         {
             _logger.LogTraceFlushCalledInDispose(Count);

src/PostHog/PostHogClient.cs

@@ -681,7 +681,7 @@ public async Task LoadFeatureFlagsAsync(CancellationToken cancellationToken)
     }
 
     /// <inheritdoc/>
-    public void Dispose() => DisposeAsync().AsTask().Wait();
+    public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult();
 
     /// <summary>
     /// Clears the local flags cache.
@@ -691,11 +691,18 @@ public async Task LoadFeatureFlagsAsync(CancellationToken cancellationToken)
     /// <inheritdoc/>
     public async ValueTask DisposeAsync()
     {
-        // Stop the polling and wait for it.
-        await _asyncBatchHandler.DisposeAsync();
-        _apiClient.Dispose();
-        _featureFlagCalledEventCache.Dispose();
-        _featureFlagsLoader.Dispose();
+        // Stop background tasks first, while the API client is still alive.
+        // The polling task in _featureFlagsLoader may call the API client during shutdown.
+        try
+        {
+            await _asyncBatchHandler.DisposeAsync();
+            await _featureFlagsLoader.DisposeAsync();
+        }
+        finally
+        {
+            _apiClient.Dispose();
+            _featureFlagCalledEventCache.Dispose();
+        }
     }
 
 

tests/UnitTests/Features/LocalFeatureFlagsLoaderTests.cs

@@ -0,0 +1,114 @@
+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.");
+        }
+
+        // Surface any exception thrown during disposal.
+        await disposeTask;
+    }
+
+    [Fact]
+    public async Task DoesNotDisposeTwice()
+    {
+        var container = new TestContainer("fake-personal-api-key");
+        container.FakeHttpMessageHandler.AddLocalEvaluationResponse(LocalEvaluationResponse);
+
+        var client = container.Activate<PostHogClient>();
+        await client.LoadFeatureFlagsAsync(CancellationToken.None);
+
+        await Task.WhenAll(
+            client.DisposeAsync().AsTask(),
+            client.DisposeAsync().AsTask());
+    }
+
+    [Fact]
+    public async Task CompletesGracefullyWhenPollingNeverStarted()
+    {
+        var container = new TestContainer();
+        var client = container.Activate<PostHogClient>();
+
+        // Dispose without ever calling LoadFeatureFlagsAsync.
+        await client.DisposeAsync();
+    }
+}