fix: recover after flush handler exceptions

.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.

src/PostHog/Library/AsyncBatchHandler.cs

@@ -43,10 +43,10 @@ internal sealed class AsyncBatchHandler<TItem, TBatchContext> : 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<IEnumerable<TItem>, 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();
+        }
     }
 }
 

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<PostHogOptions>(new()
+        {
+            FlushAt = 2,
+            MaxBatchSize = 2,
+            FlushInterval = TimeSpan.FromHours(3)
+        });
+        var items = new List<int>();
+        var firstFlushStarted = new TaskCompletionSource();
+        var secondFlushCompleted = new TaskCompletionSource();
+        int callCount = 0;
+
+        Func<IEnumerable<int>, 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<int, object>(
+            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<PostHogOptions>(new()
+        {
+            FlushAt = 1,
+            FlushInterval = TimeSpan.FromHours(3)
+        });
+        var items = new List<int>();
+        var flushStarted = new TaskCompletionSource();
+        var flushCanProceed = new TaskCompletionSource();
+
+        Func<IEnumerable<int>, Task> handlerFunc = async batch =>
+        {
+            flushStarted.SetResult();
+            await flushCanProceed.Task;
+            items.AddRange(batch);
+        };
+
+        await using var batchHandler = new AsyncBatchHandler<int, object>(
+            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