Improved Memoryblock tracking

Also fixed 3 tests that weren't handling block properly

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/MemoryPool.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
 {
@@ -60,20 +61,44 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
         /// </summary>
         private bool _disposedValue = false; // To detect redundant calls
 
+#if DEBUG
+        public MemoryPoolBlock Lease([CallerMemberName] string memberName = "",
+                             [CallerFilePath] string sourceFilePath = "",
+                             [CallerLineNumber] int sourceLineNumber = 0)
+        {
+            Debug.Assert(!_disposedValue, "Block being leased from disposed pool!");
+#else
+
         /// <summary>
         /// Called to take a block from the pool.
         /// </summary>
         /// <returns>The block that is reserved for the called. It must be passed to Return when it is no longer being used.</returns>
         public MemoryPoolBlock Lease()
         {
+#endif
             MemoryPoolBlock block;
             if (_blocks.TryDequeue(out block))
             {
                 // block successfully taken from the stack - return it
+#if DEBUG
+                block.Leaser = memberName + ", " + sourceFilePath + ", " + sourceLineNumber;
+                block.IsLeased = true;
+#if !NETSTANDARD1_3
+                block.StackTrace = new StackTrace(true).ToString();
+#endif
+#endif
                 return block;
             }
             // no blocks available - grow the pool
-            return AllocateSlab();
+            block = AllocateSlab();
+#if DEBUG
+            block.Leaser = memberName + ", " + sourceFilePath + ", " + sourceLineNumber;
+            block.IsLeased = true;
+#if !NETSTANDARD1_3
+            block.StackTrace = new StackTrace(true).ToString();
+#endif
+#endif
+            return block;
         }
 
         /// <summary>
@@ -100,6 +125,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
                     basePtr,
                     this,
                     slab);
+#if DEBUG
+                block.IsLeased = true;
+#endif
                 Return(block);
             }
 
@@ -123,19 +151,33 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
         /// <param name="block">The block to return. It must have been acquired by calling Lease on the same memory pool instance.</param>
         public void Return(MemoryPoolBlock block)
         {
+#if DEBUG
             Debug.Assert(block.Pool == this, "Returned block was not leased from this pool");
+            Debug.Assert(block.IsLeased, "Block being returned to pool twice: " + block.Leaser + "\n" + block.StackTrace);
+            block.IsLeased = false;
+#endif
 
             if (block.Slab != null && block.Slab.IsActive)
             {
                 block.Reset();
                 _blocks.Enqueue(block);
             }
+            else
+            {
+                GC.SuppressFinalize(block);
+            }
         }
 
         protected virtual void Dispose(bool disposing)
         {
             if (!_disposedValue)
             {
+                _disposedValue = true;
+#if DEBUG
+                GC.Collect();
+                GC.WaitForPendingFinalizers();
+                GC.Collect();
+#endif
                 if (disposing)
                 {
                     MemoryPoolSlab slab;
@@ -146,7 +188,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
                     }
                 }
 
-                foreach (var block in _blocks)
+                // Discard blocks in pool
+                MemoryPoolBlock block;
+                while (_blocks.TryDequeue(out block))
                 {
                     GC.SuppressFinalize(block);
                 }
@@ -155,7 +199,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
 
                 // N/A: set large fields to null.
 
-                _disposedValue = true;
             }
         }
 

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/MemoryPoolBlock.cs

@@ -71,10 +71,17 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
         /// </summary>
         public MemoryPoolBlock Next;
 
+#if DEBUG
+        public bool IsLeased { get; set; }
+        public string Leaser { get; set; }
+        public string StackTrace { get; set; }
+#endif
+
         ~MemoryPoolBlock()
         {
-            Debug.Assert(Slab == null || !Slab.IsActive, "Block being garbage collected instead of returned to pool");
+#if DEBUG
-
+            Debug.Assert(Slab == null || !Slab.IsActive, "Block being garbage collected instead of returned to pool: " + Leaser + "\n" + StackTrace);
+#endif
             if (Slab != null && Slab.IsActive)
             {
                 Pool.Return(new MemoryPoolBlock(DataArrayPtr)

src/Microsoft.AspNetCore.Server.Kestrel/Internal/Infrastructure/MemoryPoolSlab.cs

@@ -61,6 +61,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
         {
             if (!_disposedValue)
             {
+                _disposedValue = true;
+
                 if (disposing)
                 {
                     // N/A: dispose managed state (managed objects).
@@ -72,8 +74,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal.Infrastructure
 
                 // set large fields to null.
                 Array = null;
-
-                _disposedValue = true;
             }
         }
 

src/Microsoft.AspNetCore.Server.Kestrel/Internal/KestrelEngine.cs

@@ -46,6 +46,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Internal
                 thread.Stop(TimeSpan.FromSeconds(2.5));
             }
             Threads.Clear();
+#if DEBUG
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+            GC.Collect();
+#endif
         }
 
         public IDisposable CreateServer(ServerAddress address)

test/Microsoft.AspNetCore.Server.KestrelTests/AsciiDecoding.cs

@@ -59,6 +59,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                         var end = GetIterator(begin, byteRange.Length);
 
                         Assert.Throws<BadHttpRequestException>(() => begin.GetAsciiString(end));
+
+                        pool.Return(mem);
                     }
                 }
             }
@@ -150,7 +152,10 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 {
                     var returnBlock = block;
                     block = block.Next;
-                    pool.Return(returnBlock);
+                    if (returnBlock != mem0 && returnBlock != mem1)
+                    {
+                        pool.Return(returnBlock);
+                    }
                 }
 
                 pool.Return(mem0);

test/Microsoft.AspNetCore.Server.KestrelTests/MemoryPoolIteratorTests.cs

@@ -317,6 +317,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             Assert.ThrowsAny<InvalidOperationException>(() => scan.Skip(8));
 
             _pool.Return(block);
+            _pool.Return(nextBlock);
         }
 
         [Theory]