diff --git a/src/PostHog/Features/LocalFeatureFlagsLoader.cs b/src/PostHog/Features/LocalFeatureFlagsLoader.cs index 75689354..06c92f12 100644 --- a/src/PostHog/Features/LocalFeatureFlagsLoader.cs +++ b/src/PostHog/Features/LocalFeatureFlagsLoader.cs @@ -20,9 +20,11 @@ internal sealed class LocalFeatureFlagsLoader( IOptions 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)); } /// @@ -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() diff --git a/src/PostHog/Library/AsyncBatchHandler.cs b/src/PostHog/Library/AsyncBatchHandler.cs index 419b356a..fdf90f1f 100644 --- a/src/PostHog/Library/AsyncBatchHandler.cs +++ b/src/PostHog/Library/AsyncBatchHandler.cs @@ -233,7 +233,7 @@ async Task SendBatch(IReadOnlyCollection 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); diff --git a/src/PostHog/PostHogClient.cs b/src/PostHog/PostHogClient.cs index 789e496a..53686fb6 100644 --- a/src/PostHog/PostHogClient.cs +++ b/src/PostHog/PostHogClient.cs @@ -681,7 +681,7 @@ public async Task LoadFeatureFlagsAsync(CancellationToken cancellationToken) } /// - public void Dispose() => DisposeAsync().AsTask().Wait(); + public void Dispose() => DisposeAsync().AsTask().GetAwaiter().GetResult(); /// /// Clears the local flags cache. @@ -691,11 +691,18 @@ public async Task LoadFeatureFlagsAsync(CancellationToken cancellationToken) /// 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(); + } } diff --git a/tests/UnitTests/Features/LocalFeatureFlagsLoaderTests.cs b/tests/UnitTests/Features/LocalFeatureFlagsLoaderTests.cs new file mode 100644 index 00000000..b04c53ad --- /dev/null +++ b/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(); + + // 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(); + 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(); + + // Dispose without ever calling LoadFeatureFlagsAsync. + await client.DisposeAsync(); + } +}