Fix dispose hang in network handler and buffer pool cleanup

Two bugs caused LimitedFixedBufferPool.Dispose() to hang indefinitely
during server teardown (blocking ClusterResetHardDuringDisklessReplicationAttach):

1. TcpNetworkHandlerBase.Dispose() never called DisposeImpl(), so when a
   handler thread was blocked synchronously (e.g. in TryBeginDisklessSync),
   the CTS was never cancelled and activeHandlerCount was never decremented.
   DisposeActiveHandlers() would spin forever waiting for it to reach 0.

2. GarnetTcpNetworkSender.DisposeNetworkSender() disposed the saeaStack
   but not the current responseObject, leaking a PoolEntry that was never
   returned to the pool. LimitedFixedBufferPool.Dispose() then spun
   forever waiting for totalReferences to reach 0.

Also adds PoolEntry source tracking infrastructure (PoolEntryBufferType
and PoolOwnerType enums) with DEBUG-only diagnostics that log unreturned
buffer details after a 5-second timeout during pool disposal.

Co-authored-by: Copilot <[email protected]>

Author: vazois

libs/client/ClientSession/GarnetClientSession.cs

@@ -127,7 +127,7 @@ public GarnetClientSession(
 
             this.usingManagedNetworkPool = networkPool != null;
             this.networkBufferSettings = networkBufferSettings;
-            this.networkPool = networkPool ?? networkBufferSettings.CreateBufferPool();
+            this.networkPool = networkPool ?? networkBufferSettings.CreateBufferPool(ownerType: PoolOwnerType.ClientSession);
             this.bufferSizeDigits = NumUtils.CountDigits(this.networkBufferSettings.sendBufferSize);
 
             this.logger = logger;

libs/cluster/Server/Replication/ReplicationManager.cs

@@ -102,7 +102,7 @@ public ReplicationManager(ClusterProvider clusterProvider, ILogger logger = null
             this.pageSizeBits = storeWrapper.appendOnlyFile == null ? 0 : storeWrapper.appendOnlyFile.UnsafeGetLogPageSizeBits();
 
             networkBufferSettings.Log(logger, nameof(ReplicationManager));
-            this.networkPool = networkBufferSettings.CreateBufferPool(logger: logger);
+            this.networkPool = networkBufferSettings.CreateBufferPool(logger: logger, ownerType: PoolOwnerType.Replication);
             ValidateNetworkBufferSettings();
 
             aofProcessor = new AofProcessor(storeWrapper, recordToAof: false, clusterProvider: clusterProvider, logger: logger);

libs/common/Memory/LimitedFixedBufferPool.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 using System;
+using System.Collections.Concurrent;
 using System.Diagnostics;
 using System.Numerics;
 using System.Runtime.CompilerServices;
@@ -26,6 +27,11 @@ public sealed class LimitedFixedBufferPool : IDisposable
         readonly int maxAllocationSize;
         readonly ILogger logger;
 
+        /// <summary>
+        /// Pool owner type, packed into byte 1 of each <see cref="PoolEntry.source"/>.
+        /// </summary>
+        readonly int ownerByte;
+
         /// <summary>
         /// Min allocation size
         /// </summary>
@@ -41,16 +47,29 @@ public sealed class LimitedFixedBufferPool : IDisposable
         /// </summary>
         int totalOutOfBoundAllocations;
 
+#if DEBUG
+        /// <summary>
+        /// Tracks all outstanding (checked-out) pool entries for leak diagnosis.
+        /// </summary>
+        readonly ConcurrentDictionary<PoolEntry, byte> outstandingEntries = new();
+
+        /// <summary>
+        /// Timeout in milliseconds for Dispose to wait before logging outstanding entries.
+        /// </summary>
+        const int DisposeWaitDiagnosticMs = 5_000;
+#endif
+
         /// <summary>
         /// Constructor
         /// </summary>
-        public LimitedFixedBufferPool(int minAllocationSize, int maxEntriesPerLevel = 16, int numLevels = 4, ILogger logger = null)
+        public LimitedFixedBufferPool(int minAllocationSize, int maxEntriesPerLevel = 16, int numLevels = 4, PoolOwnerType ownerType = PoolOwnerType.Unknown, ILogger logger = null)
         {
             this.minAllocationSize = minAllocationSize;
             this.maxAllocationSize = minAllocationSize << (numLevels - 1);
             this.maxEntriesPerLevel = maxEntriesPerLevel;
             this.numLevels = numLevels;
             this.logger = logger;
+            this.ownerByte = (int)ownerType << 8;
             pool = new PoolLevel[numLevels];
         }
 
@@ -85,6 +104,9 @@ public bool Validate(NetworkBufferSettings settings)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void Return(PoolEntry buffer)
         {
+#if DEBUG
+            outstandingEntries.TryRemove(buffer, out _);
+#endif
             var level = Position(buffer.entry.Length);
             if (level >= 0)
             {
@@ -107,9 +129,10 @@ public void Return(PoolEntry buffer)
         /// Get buffer
         /// </summary>
         /// <param name="size"></param>
+        /// <param name="bufferType">Identifies the caller for leak diagnosis.</param>
         /// <returns></returns>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public unsafe PoolEntry Get(int size)
+        public unsafe PoolEntry Get(int size, PoolEntryBufferType bufferType = PoolEntryBufferType.Unknown)
         {
             if (Interlocked.Increment(ref totalReferences) < 0)
             {
@@ -118,6 +141,8 @@ public unsafe PoolEntry Get(int size)
                 return null;
             }
 
+            var source = ownerByte | (int)bufferType;
+
             var level = Position(size);
             if (level == -1) Interlocked.Increment(ref totalOutOfBoundAllocations);
 
@@ -132,10 +157,19 @@ public unsafe PoolEntry Get(int size)
                 {
                     Interlocked.Decrement(ref pool[level].size);
                     page.Reuse();
+                    page.source = source;
+#if DEBUG
+                    outstandingEntries[page] = 0;
+#endif
                     return page;
                 }
             }
-            return new PoolEntry(size, this);
+            var entry = new PoolEntry(size, this);
+            entry.source = source;
+#if DEBUG
+            outstandingEntries[entry] = 0;
+#endif
+            return entry;
         }
 
         /// <summary>
@@ -157,23 +191,37 @@ public void Purge()
         }
 
         /// <summary>
-        /// Dipose pool entries from all levels
+        /// Dispose pool entries from all levels
         /// NOTE:
         ///     This is used to destroy the instance and reclaim all allocated buffer pool entries.
         ///     As a consequence it spin waits until totalReferences goes back down to 0 and blocks any future allocations.
+        ///     In DEBUG builds, logs outstanding unreturned entries after a timeout for leak diagnosis.
         /// </summary>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void Dispose()
         {
-#if HANGDETECT
-            int count = 0;
+#if DEBUG
+            var sw = Stopwatch.StartNew();
+            var diagnosed = false;
 #endif
             while (totalReferences > int.MinValue &&
                 Interlocked.CompareExchange(ref totalReferences, int.MinValue, 0) != 0)
             {
-#if HANGDETECT
-                    if (++count % 10000 == 0)
-                        logger?.LogTrace("Dispose iteration {count}, {activeHandlerCount}", count, activeHandlerCount);
+#if DEBUG
+                if (!diagnosed && sw.ElapsedMilliseconds > DisposeWaitDiagnosticMs)
+                {
+                    diagnosed = true;
+                    var remaining = totalReferences;
+                    var ownerType = (PoolOwnerType)(ownerByte >> 8);
+                    logger?.LogError("LimitedFixedBufferPool.Dispose blocked with {remaining} unreturned references (poolOwner={ownerType}). Outstanding entries:", remaining, ownerType);
+                    foreach (var kvp in outstandingEntries)
+                    {
+                        var entryBufferType = (PoolEntryBufferType)(kvp.Key.source & 0xFF);
+                        var entryOwnerType = (PoolOwnerType)((kvp.Key.source >> 8) & 0xFF);
+                        logger?.LogError("  Unreturned buffer: ownerType={ownerType}, bufferType={bufferType}, size={size}",
+                            entryOwnerType, entryBufferType, kvp.Key.entry.Length);
+                    }
+                }
 #endif
                 Thread.Yield();
             }

libs/common/Memory/PoolEntry.cs

@@ -25,6 +25,12 @@ public unsafe class PoolEntry : IDisposable
         readonly LimitedFixedBufferPool pool;
         bool disposed;
 
+        /// <summary>
+        /// Packed source identifier: low byte = <see cref="PoolEntryBufferType"/>, byte 1 = <see cref="PoolOwnerType"/>.
+        /// Set when the entry is acquired via <see cref="LimitedFixedBufferPool.Get"/>.
+        /// </summary>
+        internal int source;
+
         /// <summary>
         /// Constructor
         /// </summary>

libs/common/Memory/PoolEntryTypes.cs

@@ -0,0 +1,54 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+namespace Garnet.common
+{
+    /// <summary>
+    /// Identifies the buffer role when a <see cref="PoolEntry"/> is acquired from <see cref="LimitedFixedBufferPool"/>.
+    /// </summary>
+    public enum PoolEntryBufferType : int
+    {
+        /// <summary>Default/unknown buffer type.</summary>
+        Unknown = 0,
+
+        /// <summary>Initial network receive buffer (TcpNetworkHandlerBase).</summary>
+        NetworkReceiveBuffer = 1,
+
+        /// <summary>Transport receive buffer for TLS (NetworkHandler).</summary>
+        TransportReceiveBuffer = 2,
+
+        /// <summary>Transport send buffer for TLS (NetworkHandler).</summary>
+        TransportSendBuffer = 3,
+
+        /// <summary>Doubled network receive buffer (NetworkHandler).</summary>
+        DoubleNetworkReceiveBuffer = 4,
+
+        /// <summary>Shrunk network receive buffer (NetworkHandler).</summary>
+        ShrinkNetworkReceiveBuffer = 5,
+
+        /// <summary>Doubled transport receive buffer for TLS (NetworkHandler).</summary>
+        DoubleTransportReceiveBuffer = 6,
+
+        /// <summary>Send buffer for async socket operations (GarnetSaeaBuffer).</summary>
+        SaeaSendBuffer = 7,
+    }
+
+    /// <summary>
+    /// Identifies the owner of a <see cref="LimitedFixedBufferPool"/> instance.
+    /// Set at pool construction time to indicate which subsystem created the pool.
+    /// </summary>
+    public enum PoolOwnerType : int
+    {
+        /// <summary>Default/unknown owner.</summary>
+        Unknown = 0,
+
+        /// <summary>Server-side network pool (GarnetServerTcp).</summary>
+        ServerNetwork = 1,
+
+        /// <summary>Replication network pool (ReplicationManager).</summary>
+        Replication = 2,
+
+        /// <summary>Client-side network pool (GarnetClientSession, self-managed).</summary>
+        ClientSession = 3,
+    }
+}

libs/common/NetworkBufferSettings.cs

@@ -73,17 +73,18 @@ public static NetworkBufferSettings GetInclusive(NetworkBufferSettings[] setting
         /// Allocate network buffer pool
         /// </summary>
         /// <param name="maxEntriesPerLevel"></param>
+        /// <param name="ownerType"></param>
         /// <param name="logger"></param>
         /// <returns></returns>
-        public LimitedFixedBufferPool CreateBufferPool(int maxEntriesPerLevel = 16, ILogger logger = null)
+        public LimitedFixedBufferPool CreateBufferPool(int maxEntriesPerLevel = 16, PoolOwnerType ownerType = PoolOwnerType.Unknown, ILogger logger = null)
         {
             var minSize = Math.Min(Math.Min(sendBufferSize, initialReceiveBufferSize), maxReceiveBufferSize);
             var maxSize = Math.Max(Math.Max(sendBufferSize, initialReceiveBufferSize), maxReceiveBufferSize);
 
             var levels = LimitedFixedBufferPool.GetLevel(minSize, maxSize) + 1;
             Debug.Assert(levels >= 0);
             levels = Math.Max(4, levels);
-            return new LimitedFixedBufferPool(minSize, maxEntriesPerLevel: maxEntriesPerLevel, numLevels: levels, logger: logger);
+            return new LimitedFixedBufferPool(minSize, maxEntriesPerLevel: maxEntriesPerLevel, numLevels: levels, logger: logger, ownerType: ownerType);
         }
 
         public void Log(ILogger logger, string category)

libs/common/Networking/GarnetSaeaBuffer.cs

@@ -30,7 +30,7 @@ public GarnetSaeaBuffer(EventHandler<SocketAsyncEventArgs> eventHandler, Network
         {
             socketEventAsyncArgs = new SocketAsyncEventArgs();
 
-            buffer = networkPool.Get(networkBufferSettings.sendBufferSize);
+            buffer = networkPool.Get(networkBufferSettings.sendBufferSize, PoolEntryBufferType.SaeaSendBuffer);
             socketEventAsyncArgs.SetBuffer(buffer.entry, 0, buffer.entry.Length);
             socketEventAsyncArgs.Completed += eventHandler;
         }

libs/common/Networking/GarnetTcpNetworkSender.cs

@@ -252,6 +252,10 @@ public override void DisposeNetworkSender(bool waitForSendCompletion)
             // Wait for ongoing sends to complete
             while (throttleCount >= 0 && Interlocked.CompareExchange(ref throttleCount, int.MinValue, 0) != 0) Thread.Yield();
 
+            // Dispose the current response object if one is held
+            responseObject?.Dispose();
+            responseObject = null;
+
             // Empty and dispose the stack
             saeaStack.Dispose();
 

libs/common/Networking/NetworkHandler.cs

@@ -129,11 +129,11 @@ public unsafe NetworkHandler(TServerHook serverHook, TNetworkSender networkSende
                 expectingData = new SemaphoreSlim(0);
                 cancellationTokenSource = new();
 
-                transportReceiveBufferEntry = this.networkPool.Get(this.networkBufferSettings.initialReceiveBufferSize);
+                transportReceiveBufferEntry = this.networkPool.Get(this.networkBufferSettings.initialReceiveBufferSize, PoolEntryBufferType.TransportReceiveBuffer);
                 transportReceiveBuffer = transportReceiveBufferEntry.entry;
                 transportReceiveBufferPtr = transportReceiveBufferEntry.entryPtr;
 
-                transportSendBufferEntry = this.networkPool.Get(this.networkBufferSettings.sendBufferSize);
+                transportSendBufferEntry = this.networkPool.Get(this.networkBufferSettings.sendBufferSize, PoolEntryBufferType.TransportSendBuffer);
                 transportSendBuffer = transportSendBufferEntry.entry;
                 transportSendBufferPtr = transportSendBufferEntry.entryPtr;
             }
@@ -501,7 +501,7 @@ unsafe bool TryProcessRequest()
 
         unsafe void DoubleNetworkReceiveBuffer()
         {
-            var tmp = networkPool.Get(networkReceiveBuffer.Length * 2);
+            var tmp = networkPool.Get(networkReceiveBuffer.Length * 2, PoolEntryBufferType.DoubleNetworkReceiveBuffer);
             Array.Copy(networkReceiveBuffer, tmp.entry, networkReceiveBuffer.Length);
             networkReceiveBufferEntry.Dispose();
             networkReceiveBufferEntry = tmp;
@@ -515,7 +515,7 @@ unsafe void ShrinkNetworkReceiveBuffer()
         {
             Debug.Assert(networkReadHead == 0, "Shouldn't call if remaining data not already moved to head of receive buffer");
 
-            var tmp = networkPool.Get(networkBufferSettings.maxReceiveBufferSize);
+            var tmp = networkPool.Get(networkBufferSettings.maxReceiveBufferSize, PoolEntryBufferType.ShrinkNetworkReceiveBuffer);
             if (networkBytesRead > 0)
             {
                 Array.Copy(networkReceiveBuffer, tmp.entry, networkBytesRead);
@@ -543,7 +543,7 @@ unsafe void DoubleTransportReceiveBuffer()
         {
             if (sslStream != null)
             {
-                var tmp = networkPool.Get(transportReceiveBuffer.Length * 2);
+                var tmp = networkPool.Get(transportReceiveBuffer.Length * 2, PoolEntryBufferType.DoubleTransportReceiveBuffer);
                 Array.Copy(transportReceiveBuffer, tmp.entry, transportReceiveBuffer.Length);
                 transportReceiveBufferEntry.Dispose();
                 transportReceiveBufferEntry = tmp;

libs/common/Networking/TcpNetworkHandlerBase.cs

@@ -170,6 +170,11 @@ public override void Dispose()
                 // Dispose of the socket to free up unmanaged resources
                 socket.Dispose();
             }
+
+            // Ensure the handler is fully cleaned up: cancel CTS, remove from activeHandlers,
+            // and return pool buffers. DisposeImpl is guarded by disposeCount so it is safe to
+            // call even when the SAEA callback path has already invoked it.
+            DisposeImpl();
         }
 
         /// <summary>
@@ -252,7 +257,7 @@ void HandleReceiveFailure(Exception ex, SocketAsyncEventArgs e)
 
         unsafe void AllocateNetworkReceiveBuffer()
         {
-            networkReceiveBufferEntry = networkPool.Get(networkBufferSettings.initialReceiveBufferSize);
+            networkReceiveBufferEntry = networkPool.Get(networkBufferSettings.initialReceiveBufferSize, PoolEntryBufferType.NetworkReceiveBuffer);
             networkReceiveBuffer = networkReceiveBufferEntry.entry;
             networkReceiveBufferPtr = networkReceiveBufferEntry.entryPtr;
         }