Fix disposal race condition in AsyncBatchHandler

src/PostHog/Library/AsyncBatchHandler.cs

@@ -43,6 +43,8 @@ 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 Task _timerTask;
+    readonly Task _flushSignalTask;
     volatile int _disposed;
     volatile int _flushing;
 
@@ -63,8 +65,8 @@ public AsyncBatchHandler(
             FullMode = BoundedChannelFullMode.DropOldest
         });
         _timer = new PeriodicTimer(options.Value.FlushInterval, timeProvider);
-        taskScheduler.Run(() => HandleTimer(_cancellationTokenSource.Token));
-        taskScheduler.Run(() => HandleFlushSignal(_cancellationTokenSource.Token));
+        _timerTask = taskScheduler.Run(() => HandleTimer(_cancellationTokenSource.Token));
+        _flushSignalTask = taskScheduler.Run(() => HandleFlushSignal(_cancellationTokenSource.Token));
     }
 
     public AsyncBatchHandler(
@@ -244,17 +246,20 @@ public async ValueTask DisposeAsync()
 
         _logger.LogInfoDisposeAsyncCalled();
 
-        // Ensures that both the HandleFlushSignal and HandleTimer throw
-        // OperationCancelledException which is handled gracefully.
+        // 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();
-        _cancellationTokenSource.Dispose();
+        await Task.WhenAll(_timerTask, _flushSignalTask);
+
         _timer.Dispose();
         _flushSignal.Dispose();
+        _cancellationTokenSource.Dispose();
         _channel.Writer.Complete();
         try
         {
             _logger.LogTraceFlushCalledInDispose(Count);
-            // Flush the last remaining items.
+            // Flush any remaining items that weren't picked up by the background tasks.
             await FlushBatchesAsync();
         }
 #pragma warning disable CA1031 // Do not catch general exception types

tests/UnitTests/Library/AsyncBatchHandlerTests.cs

@@ -330,6 +330,48 @@ public async Task DoesNotDisposeTwice()
         Assert.Equal([1, 2], items);
     }
 
+    [Fact]
+    public async Task FlushesItemWhenDisposedDuringInFlightFlush()
+    {
+        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);
+        };
+
+        var batchHandler = new AsyncBatchHandler<int, object>(
+            handlerFunc, () => new object(), new FakeTimeProvider(), options);
+
+        // Enqueue triggers a flush (FlushAt = 1), which will block on flushCanProceed.
+        batchHandler.Enqueue(Task.FromResult(42));
+        await flushStarted.Task;
+
+        // Start disposal while the flush is still in progress.
+        var disposeTask = batchHandler.DisposeAsync().AsTask();
+
+        // Unblock the flush handler so it can complete.
+        flushCanProceed.SetResult();
+
+        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.");
+        }
+
+        Assert.Equal([42], items);
+    }
+
     [Fact]
     public async Task HandlesExceptionsInFlushBatchAsync()
     {