diff --git a/.changeset/happy-pens-attend.md b/.changeset/happy-pens-attend.md new file mode 100644 index 00000000..f961e90e --- /dev/null +++ b/.changeset/happy-pens-attend.md @@ -0,0 +1,6 @@ +--- +"PostHog": patch +"PostHog.AspNetCore": patch +--- + +Fix `AsyncBatchHandler` background flushing so transient batch send failures no longer permanently stop future flushes. `FlushAsync()` now also waits for an in-progress flush instead of returning early without doing work. diff --git a/src/PostHog/Library/AsyncBatchHandler.cs b/src/PostHog/Library/AsyncBatchHandler.cs index fdf90f1f..51c990c9 100644 --- a/src/PostHog/Library/AsyncBatchHandler.cs +++ b/src/PostHog/Library/AsyncBatchHandler.cs @@ -43,10 +43,10 @@ internal sealed class AsyncBatchHandler : IDisposable, IAs readonly PeriodicTimer _timer; readonly CancellationTokenSource _cancellationTokenSource = new(); readonly SemaphoreSlim _flushSignal = new(0); // Used to signal when a flush is needed + readonly SemaphoreSlim _flushLock = new(1, 1); // Ensures only one flush drains the queue at a time. readonly Task _timerTask; readonly Task _flushSignalTask; volatile int _disposed; - volatile int _flushing; public AsyncBatchHandler( Func, Task> batchHandlerFunc, @@ -132,27 +132,28 @@ void SignalFlush() async Task HandleFlushSignal(CancellationToken cancellationToken) { - try + while (!cancellationToken.IsCancellationRequested) { - while (!cancellationToken.IsCancellationRequested) + try { await _flushSignal.WaitAsync(cancellationToken); - await FlushBatchesAsync(); + await TryFlushBatchesAsync(); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + _logger.LogTraceOperationCancelled(nameof(HandleFlushSignal)); + break; } - } - catch (OperationCanceledException) - { - _logger.LogTraceOperationCancelled(nameof(HandleFlushSignal)); - } #pragma warning disable CA1031 // Do not catch general exception types - catch (Exception ex) + catch (Exception ex) #pragma warning restore CA1031 - { - // When running locally we want this to throw so we can see the exception. - Debug.Assert(ex is not ArgumentException and not NullReferenceException, - $"Unexpected {ex.GetType().FullName} occurred during async batch handling."); + { + // When running locally we want this to throw so we can see the exception. + Debug.Assert(ex is not ArgumentException and not NullReferenceException, + $"Unexpected {ex.GetType().FullName} occurred during async batch handling."); - _logger.LogErrorUnexpectedException(ex); + _logger.LogErrorUnexpectedException(ex); + } } } @@ -195,25 +196,43 @@ public async Task FlushAsync() async Task FlushBatchesAsync() { - // If we're flushing, don't start another flush. - if (Interlocked.CompareExchange(ref _flushing, 1, 0) == 1) + await _flushLock.WaitAsync(); + try + { + await DrainBatchesAsync(); + } + finally + { + _flushLock.Release(); + } + } + + async Task TryFlushBatchesAsync() + { + // Background flushes coalesce when another flush is already in progress. + if (!await _flushLock.WaitAsync(0)) { return; } try { - var batchContext = _batchContextFunc(); - while (_channel.Reader.TryReadBatch(_options.Value.MaxBatchSize, out var batch)) - { - var tasks = batch.Select(item => item.ResolveItem(batchContext)); - var resolved = await Task.WhenAll(tasks); - await SendBatch(resolved); - } + await DrainBatchesAsync(); } finally { - Interlocked.Exchange(ref _flushing, 0); + _flushLock.Release(); + } + } + + async Task DrainBatchesAsync() + { + var batchContext = _batchContextFunc(); + while (_channel.Reader.TryReadBatch(_options.Value.MaxBatchSize, out var batch)) + { + var tasks = batch.Select(item => item.ResolveItem(batchContext)); + var resolved = await Task.WhenAll(tasks); + await SendBatch(resolved); } return; @@ -247,8 +266,7 @@ public async ValueTask DisposeAsync() _logger.LogInfoDisposeAsyncCalled(); // 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. + // This ensures any in-flight flush completes before we attempt the final flush below. try { await _cancellationTokenSource.CancelAsync(); @@ -277,6 +295,10 @@ public async ValueTask DisposeAsync() _logger.LogErrorUnexpectedException(e); } + finally + { + _flushLock.Dispose(); + } } } diff --git a/tests/UnitTests/Library/AsyncBatchHandlerTests.cs b/tests/UnitTests/Library/AsyncBatchHandlerTests.cs index ebdfe805..49488f9e 100644 --- a/tests/UnitTests/Library/AsyncBatchHandlerTests.cs +++ b/tests/UnitTests/Library/AsyncBatchHandlerTests.cs @@ -147,6 +147,89 @@ public async Task FlushBatchAsyncContinuesAfterException() Assert.Equal([3], items); } + [Fact] + public async Task BackgroundFlushContinuesAfterException() + { + var options = new FakeOptions(new() + { + FlushAt = 2, + MaxBatchSize = 2, + FlushInterval = TimeSpan.FromHours(3) + }); + var items = new List(); + var firstFlushStarted = new TaskCompletionSource(); + var secondFlushCompleted = new TaskCompletionSource(); + int callCount = 0; + + Func, Task> handlerFunc = batch => + { + callCount++; + if (callCount == 1) + { + firstFlushStarted.SetResult(); + throw new InvalidOperationException("Test exception"); + } + + items.AddRange(batch); + secondFlushCompleted.SetResult(); + return Task.CompletedTask; + }; + + await using var batchHandler = new AsyncBatchHandler( + handlerFunc, + () => new object(), + new FakeTimeProvider(), + options); + + batchHandler.Enqueue(Task.FromResult(1)); + batchHandler.Enqueue(Task.FromResult(2)); + await CompleteWithin(firstFlushStarted.Task, TimeSpan.FromSeconds(1)); + + batchHandler.Enqueue(Task.FromResult(3)); + batchHandler.Enqueue(Task.FromResult(4)); + await CompleteWithin(secondFlushCompleted.Task, TimeSpan.FromSeconds(1)); + + Assert.Equal([3, 4], items); + } + + [Fact] + public async Task FlushAsyncWaitsForInProgressFlush() + { + var options = new FakeOptions(new() + { + FlushAt = 1, + FlushInterval = TimeSpan.FromHours(3) + }); + var items = new List(); + var flushStarted = new TaskCompletionSource(); + var flushCanProceed = new TaskCompletionSource(); + + Func, Task> handlerFunc = async batch => + { + flushStarted.SetResult(); + await flushCanProceed.Task; + items.AddRange(batch); + }; + + await using var batchHandler = new AsyncBatchHandler( + handlerFunc, + () => new object(), + new FakeTimeProvider(), + options); + + batchHandler.Enqueue(Task.FromResult(42)); + await CompleteWithin(flushStarted.Task, TimeSpan.FromSeconds(1)); + + var flushTask = batchHandler.FlushAsync(); + var completedEarly = await Task.WhenAny(flushTask, Task.Delay(TimeSpan.FromMilliseconds(100))); + Assert.NotSame(flushTask, completedEarly); + + flushCanProceed.SetResult(); + await CompleteWithin(flushTask, TimeSpan.FromSeconds(1)); + + Assert.Equal([42], items); + } + [Fact] public async Task DropsOlderEventsWhenMaxQueueMet() { @@ -260,6 +343,17 @@ public async Task IgnoresEnqueuedItemAfterDispose() Assert.Equal(0, batchHandler.Count); } + + static async Task CompleteWithin(Task task, TimeSpan timeout) + { + var completedTask = await Task.WhenAny(task, Task.Delay(timeout)); + if (completedTask != task) + { + throw new TimeoutException("The operation timed out."); + } + + await task; + } } public class TheDisposeAsyncMethod