chore(deps): update dependency tunit to 1.14.0

.github/workflows/ci.yml

@@ -72,6 +72,16 @@ jobs:
           if-no-files-found: ignore
           retention-days: 30
 
+      - name: Upload HangDump artifacts
+        uses: actions/upload-artifact@v6
+        if: always()
+        with:
+          name: hangdump-artifacts-${{ matrix.category }}
+          path: |
+            **/*.dmp
+          if-no-files-found: ignore
+          retention-days: 7
+
   integration-tests:
     name: Integration Tests (${{ matrix.category }})
     needs: build-and-unit-test

src/Dekaf/Consumer/KafkaConsumer.cs

@@ -653,34 +653,26 @@ public async IAsyncEnumerable<ConsumeResult<TKey, TValue>> ConsumeAsync(
                     else
                     {
                         // Wait for prefetch with timeout, then try direct fetch
-                        // Rent CTS from pool to avoid allocation
-                        var timeoutCts = _ctsPool.Rent();
+                        using var timeoutCts = _ctsPool.Rent();
+                        timeoutCts.CancelAfter(TimeSpan.FromMilliseconds(_options.FetchMaxWaitMs));
+                        using var reg = cancellationToken.CanBeCanceled
+                            ? cancellationToken.Register(static s => ((CancellationTokenSource)s!).Cancel(), timeoutCts)
+                            : default;
                         try
                         {
-                            timeoutCts.CancelAfter(TimeSpan.FromMilliseconds(_options.FetchMaxWaitMs));
-                            using var reg = cancellationToken.CanBeCanceled
-                                ? cancellationToken.Register(static s => ((CancellationTokenSource)s!).Cancel(), timeoutCts)
-                                : default;
-                            try
-                            {
-                                var fetched = await _prefetchChannel.Reader.ReadAsync(timeoutCts.Token).ConfigureAwait(false);
-                                _pendingFetches.Enqueue(fetched);
-                                TrackPrefetchedBytes(fetched, release: true);
-                            }
-                            catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
-                            {
-                                // Prefetch not ready - check for EOF events before continuing
-                                // (EOF events are queued by prefetch loop when partition is caught up)
-                            }
-                            catch (ChannelClosedException ex) when (ex.InnerException is KafkaException kafkaEx)
-                            {
-                                // Rethrow the original KafkaException from the prefetch task
-                                throw kafkaEx;
-                            }
+                            var fetched = await _prefetchChannel.Reader.ReadAsync(timeoutCts.Token).ConfigureAwait(false);
+                            _pendingFetches.Enqueue(fetched);
+                            TrackPrefetchedBytes(fetched, release: true);
                         }
-                        finally
+                        catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
                         {
-                            _ctsPool.Return(timeoutCts);
+                            // Prefetch not ready - check for EOF events before continuing
+                            // (EOF events are queued by prefetch loop when partition is caught up)
+                        }
+                        catch (ChannelClosedException ex) when (ex.InnerException is KafkaException kafkaEx)
+                        {
+                            // Rethrow the original KafkaException from the prefetch task
+                            throw kafkaEx;
                         }
                     }
                 }
@@ -888,8 +880,8 @@ private async Task PrefetchLoopAsync(CancellationToken cancellationToken)
     private async ValueTask PrefetchRecordsAsync(CancellationToken cancellationToken)
     {
         // Rent wakeup CTS from pool — used as the combined cancellation source
-        // instead of allocating a LinkedCTS
-        var wakeupCts = _ctsPool.Rent();
+        // instead of allocating a LinkedCTS. Dispose auto-returns to pool.
+        using var wakeupCts = _ctsPool.Rent();
         _wakeupCts = wakeupCts;
 
         try
@@ -932,7 +924,6 @@ private async ValueTask PrefetchRecordsAsync(CancellationToken cancellationToken
         finally
         {
             _wakeupCts = null;
-            _ctsPool.Return(wakeupCts);
         }
     }
 
@@ -1184,50 +1175,42 @@ private static long EstimatePendingFetchBytes(PendingFetchData pending)
         TimeSpan timeout,
         CancellationToken cancellationToken = default)
     {
-        // Rent CTS from pool to avoid allocation
-        var timeoutCts = _ctsPool.Rent();
-        try
-        {
-            timeoutCts.CancelAfter(timeout);
-
-            // Fast path: if no external cancellation, use timeout CTS directly (avoids allocation)
-            if (!cancellationToken.CanBeCanceled)
-            {
-                try
-                {
-                    await foreach (var result in ConsumeAsync(timeoutCts.Token).ConfigureAwait(false))
-                    {
-                        return result;
-                    }
-                }
-                catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
-                {
-                    // Timeout expired with no messages - return null instead of throwing
-                }
-                return null;
-            }
-
-            // Slow path: need to link external cancellation with timeout
-            using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);
+        using var timeoutCts = _ctsPool.Rent();
+        timeoutCts.CancelAfter(timeout);
 
+        // Fast path: if no external cancellation, use timeout CTS directly (avoids allocation)
+        if (!cancellationToken.CanBeCanceled)
+        {
             try
             {
-                await foreach (var result in ConsumeAsync(linkedCts.Token).ConfigureAwait(false))
+                await foreach (var result in ConsumeAsync(timeoutCts.Token).ConfigureAwait(false))
                 {
                     return result;
                 }
             }
-            catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+            catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
             {
-                // Our timeout expired (not user cancellation) with no messages - return null
+                // Timeout expired with no messages - return null instead of throwing
             }
-
             return null;
         }
-        finally
+
+        // Slow path: need to link external cancellation with timeout
+        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);
+
+        try
         {
-            _ctsPool.Return(timeoutCts);
+            await foreach (var result in ConsumeAsync(linkedCts.Token).ConfigureAwait(false))
+            {
+                return result;
+            }
         }
+        catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+        {
+            // Our timeout expired (not user cancellation) with no messages - return null
+        }
+
+        return null;
     }
 
     public async ValueTask CommitAsync(CancellationToken cancellationToken = default)
@@ -1995,8 +1978,8 @@ private async ValueTask<long> ResolveOffsetAsync(TopicPartition partition, long
 
     private async ValueTask FetchRecordsAsync(CancellationToken cancellationToken)
     {
-        // Rent wakeup CTS from pool to avoid allocation
-        var wakeupCts = _ctsPool.Rent();
+        // Rent wakeup CTS from pool — Dispose auto-returns to pool.
+        using var wakeupCts = _ctsPool.Rent();
         _wakeupCts = wakeupCts;
 
         try
@@ -2049,7 +2032,6 @@ private async ValueTask FetchRecordsAsync(CancellationToken cancellationToken)
         finally
         {
             _wakeupCts = null;
-            _ctsPool.Return(wakeupCts);
         }
     }
 

src/Dekaf/Internal/CancellationTokenSourcePool.cs

@@ -4,63 +4,78 @@ namespace Dekaf.Internal;
 
 /// <summary>
 /// Thread-safe pool for CancellationTokenSource instances to avoid allocations in hot paths.
+/// Rent() returns a <see cref="PooledCancellationTokenSource"/> whose Dispose() auto-returns
+/// to the pool, making it safe to use with <c>using var</c>.
 /// </summary>
+/// <remarks>
+/// Design based on Microsoft's ASP.NET Core CancellationTokenSourcePool.
+/// TryReset() is called on return (not rent) to keep the pool clean — cancelled
+/// instances are disposed immediately and never occupy a pool slot.
+/// </remarks>
 internal sealed class CancellationTokenSourcePool
 {
-    private readonly ConcurrentBag<CancellationTokenSource> _pool = new();
     // Pool size must accommodate peak concurrent requests across all connections.
     // At 250K msg/sec with batching and multiple partitions, concurrent in-flight
     // requests can spike to 500+ during bursts or when broker responses are slow.
     // 512 provides headroom for high-throughput scenarios while bounding memory usage.
     private const int MaxPoolSize = 512;
+    private readonly ConcurrentQueue<PooledCancellationTokenSource> _queue = new();
     private int _count;
 
-    public CancellationTokenSource Rent()
+    public PooledCancellationTokenSource Rent()
     {
-        if (_pool.TryTake(out var cts))
+        if (_queue.TryDequeue(out var cts))
         {
             Interlocked.Decrement(ref _count);
-            if (cts.TryReset())
-            {
-                return cts;
-            }
-            // Reset failed, dispose and create new
-            cts.Dispose();
+            Volatile.Write(ref cts._returned, 0); // Reset guard for new rental
+            return cts;
         }
 
-        return new CancellationTokenSource();
+        return new PooledCancellationTokenSource(this);
     }
 
-    public void Return(CancellationTokenSource cts)
+    private bool Return(PooledCancellationTokenSource cts)
     {
-        if (cts.IsCancellationRequested)
-        {
-            cts.Dispose();
-            return;
-        }
-
-        // Atomic check-and-increment to prevent race condition
-        var currentCount = Volatile.Read(ref _count);
-        while (currentCount < MaxPoolSize)
+        if (Interlocked.Increment(ref _count) > MaxPoolSize || !cts.TryReset())
         {
-            if (Interlocked.CompareExchange(ref _count, currentCount + 1, currentCount) == currentCount)
-            {
-                _pool.Add(cts);
-                return;
-            }
-            currentCount = Volatile.Read(ref _count);
+            Interlocked.Decrement(ref _count);
+            return false;
         }
 
-        // Pool is full, dispose
-        cts.Dispose();
+        _queue.Enqueue(cts);
+        return true;
     }
 
     public void Clear()
     {
-        while (_pool.TryTake(out var cts))
+        while (_queue.TryDequeue(out var cts))
         {
-            cts.Dispose();
+            cts.DisposeWithoutPooling();
         }
         _count = 0;
     }
+
+    internal sealed class PooledCancellationTokenSource : CancellationTokenSource
+    {
+        private readonly CancellationTokenSourcePool _pool;
+        internal int _returned; // Guard against double-dispose returning to pool twice
+
+        internal PooledCancellationTokenSource(CancellationTokenSourcePool pool) => _pool = pool;
+
+        protected override void Dispose(bool disposing)
+        {
+            if (disposing && Interlocked.Exchange(ref _returned, 1) == 0)
+            {
+                if (!_pool.Return(this))
+                {
+                    base.Dispose(disposing);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Disposes without returning to pool. Used by <see cref="CancellationTokenSourcePool.Clear"/>.
+        /// </summary>
+        internal void DisposeWithoutPooling() => base.Dispose(true);
+    }
 }