Verify that memory pool blocks aren't leaked in tests

src/Microsoft.AspNetCore.Server.Kestrel/Infrastructure/MemoryPool2.cs

@@ -161,6 +161,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Infrastructure
                     }
                 }
 
+                foreach (var block in _blocks)
+                {
+                    GC.SuppressFinalize(block);
+                }
+
                 // N/A: free unmanaged resources (unmanaged objects) and override a finalizer below.
 
                 // N/A: set large fields to null.

src/Microsoft.AspNetCore.Server.Kestrel/Infrastructure/MemoryPoolBlock2.cs

@@ -77,7 +77,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Infrastructure
         ~MemoryPoolBlock2()
         {
             Debug.Assert(!_pinHandle.IsAllocated, "Ad-hoc memory block wasn't unpinned");
-            // Debug.Assert(Slab == null || !Slab.IsActive, "Block being garbage collected instead of returned to pool");
+            Debug.Assert(Slab == null, "Block being garbage collected instead of returned to pool");
 
             if (_pinHandle.IsAllocated)
             {

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

@@ -5,7 +5,6 @@ using System;
 using System.IO;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
-using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Filter;
 using Microsoft.AspNetCore.Testing.xunit;
 using Xunit;

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

@@ -113,7 +113,6 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             var address = ServerAddress.FromUrl($"http://localhost:{port}/");
             var started = engine.CreateServer(address);
 
-            Console.WriteLine("Started");
             var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
             socket.Connect(new IPEndPoint(IPAddress.Loopback, port));
             socket.Send(Encoding.ASCII.GetBytes("POST / HTTP/1.0\r\n\r\nHello World"));
@@ -1058,7 +1057,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                         throw;
                     }
 
-                    readTcs.SetCanceled();
+                    readTcs.SetException(new Exception("This shouldn't be reached."));
                 }
             }, testContext))
             {
@@ -1122,6 +1121,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                     for (int i = 0; i < 100; i++)
                     {
                         await response.WriteAsync(largeString, lifetime.RequestAborted);
+                        registrationWh.Wait(1000);
                     }
                 }
                 catch (Exception ex)

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

@@ -26,20 +26,23 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
         {
             var trace = new KestrelTrace(new TestKestrelTrace());
             var ltp = new LoggingThreadPool(trace);
-            var socketInput = new SocketInput(new MemoryPool2(), ltp);
+            using (var pool = new MemoryPool2())
-            var headerCollection = new FrameRequestHeaders();
+            using (var socketInput = new SocketInput(pool, ltp))
+            {
+                var headerCollection = new FrameRequestHeaders();
 
-            var headerArray = Encoding.ASCII.GetBytes(rawHeaders);
+                var headerArray = Encoding.ASCII.GetBytes(rawHeaders);
-            socketInput.IncomingData(headerArray, 0, headerArray.Length);
+                socketInput.IncomingData(headerArray, 0, headerArray.Length);
 
-            var success = Frame.TakeMessageHeaders(socketInput, headerCollection);
+                var success = Frame.TakeMessageHeaders(socketInput, headerCollection);
 
-            Assert.True(success);
+                Assert.True(success);
-            Assert.Equal(numHeaders, headerCollection.Count());
+                Assert.Equal(numHeaders, headerCollection.Count());
 
-            // Assert TakeMessageHeaders consumed all the input
+                // Assert TakeMessageHeaders consumed all the input
-            var scan = socketInput.ConsumingStart();
+                var scan = socketInput.ConsumingStart();
-            Assert.True(scan.IsEnd);
+                Assert.True(scan.IsEnd);
+            }
         }
     }
 }

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

@@ -50,6 +50,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                     hit.Seek(ref vectorCh, ref vectorMaxValues, ref vectorMaxValues);
                     Assert.Equal(ch, iterator.GetLength(hit));
                 }
+
+                pool.Return(block);
             }
         }
 
@@ -109,6 +111,10 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                     hit.Seek(ref vectorMaxValues, ref vectorMaxValues, ref vectorCh);
                     Assert.Equal(ch, iterator.GetLength(hit));
                 }
+
+                pool.Return(block1);
+                pool.Return(block2);
+                pool.Return(block3);
             }
         }
 
@@ -192,6 +198,9 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 AssertIterator(iter300a2, block2, block2.End);
                 AssertIterator(iter300b2, block2, block2.End);
                 AssertIterator(iter300c2, block2, block2.End);
+
+                pool.Return(block1);
+                pool.Return(block2);
             }
         }
 
@@ -227,6 +236,9 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
 
                 endIterator.CopyTo(array, 0, 256, out actual);
                 Assert.Equal(0, actual);
+
+                pool.Return(block1);
+                pool.Return(block2);
             }
         }
 
@@ -260,6 +272,15 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 Assert.Equal(-1, start.Take());
                 Assert.Equal(start.Block, end.Block);
                 Assert.Equal(start.Index, end.Index);
+
+                var block = block1;
+                while (block != null)
+                {
+                    var returnBlock = block;
+                    block = block.Next;
+
+                    pool.Return(returnBlock);
+                }
             }
         }
 
@@ -284,6 +305,11 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 Assert.False(iterStart.IsEnd);
                 Assert.False(iterMid.IsEnd);
                 Assert.True(iterEnd.IsEnd);
+
+                pool.Return(block1);
+                pool.Return(block2);
+                pool.Return(block3);
+                pool.Return(block4);
             }
         }
 

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

@@ -116,6 +116,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
 
             Assert.Equal(expectResult, found);
             Assert.Equal(expectIndex, begin.Index - block.Start);
+
+            _pool.Return(block);
         }
 
         [Fact]
@@ -174,6 +176,11 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
 
             // Can't put anything by the end
             Assert.False(head.Put(0xFF));
+
+            for (var i = 0; i < 4; ++i)
+            {
+                _pool.Return(blocks[i]);
+            }
         }
 
         [Fact]
@@ -193,6 +200,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             // Assert
             Assert.Equal(0x0102030405060708, result);
             Assert.Equal(originalIndex, scan.Index);
+
+            _pool.Return(block);
         }
 
         [Theory]
@@ -229,6 +238,9 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             // Assert
             Assert.Equal(0x0102030405060708, result);
             Assert.Equal(originalIndex, scan.Index);
+
+            _pool.Return(block);
+            _pool.Return(nextBlock);
         }
 
         [Theory]
@@ -268,6 +280,9 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             // Assert
             Assert.Equal(0x08, result);
             Assert.NotEqual(originalIndex, scan.Index);
+
+            _pool.Return(block);
+            _pool.Return(nextBlock);
         }
 
         [Theory]
@@ -304,6 +319,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             // Assert
             Assert.Equal(expectedResult, result);
             Assert.Equal(expectedKnownString, knownString);
+
+            _pool.Return(block);
         }
 
         [Theory]
@@ -329,6 +346,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             // Assert
             Assert.Equal(expectedResult, result);
             Assert.Equal(expectedKnownString, knownString);
+
+            _pool.Return(block);
         }
     }
 }

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

@@ -17,66 +17,72 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
         [Fact]
         public void Http10ConnectionClose()
         {
-            var input = new TestInput();
+            using (var input = new TestInput())
-            var body = MessageBody.For("HTTP/1.0", new FrameRequestHeaders(), input.FrameContext);
+            {
-            var stream = new FrameRequestStream().StartAcceptingReads(body);
+                var body = MessageBody.For("HTTP/1.0", new FrameRequestHeaders(), input.FrameContext);
+                var stream = new FrameRequestStream().StartAcceptingReads(body);
 
-            input.Add("Hello", true);
+                input.Add("Hello", true);
 
-            var buffer1 = new byte[1024];
+                var buffer1 = new byte[1024];
-            var count1 = stream.Read(buffer1, 0, 1024);
+                var count1 = stream.Read(buffer1, 0, 1024);
-            AssertASCII("Hello", new ArraySegment<byte>(buffer1, 0, 5));
+                AssertASCII("Hello", new ArraySegment<byte>(buffer1, 0, 5));
 
-            var buffer2 = new byte[1024];
+                var buffer2 = new byte[1024];
-            var count2 = stream.Read(buffer2, 0, 1024);
+                var count2 = stream.Read(buffer2, 0, 1024);
-            Assert.Equal(0, count2);
+                Assert.Equal(0, count2);
+            }
         }
 
         [Fact]
         public async Task Http10ConnectionCloseAsync()
         {
-            var input = new TestInput();
+            using (var input = new TestInput())
-            var body = MessageBody.For("HTTP/1.0", new FrameRequestHeaders(), input.FrameContext);
+            {
-            var stream = new FrameRequestStream().StartAcceptingReads(body);
+                var body = MessageBody.For("HTTP/1.0", new FrameRequestHeaders(), input.FrameContext);
+                var stream = new FrameRequestStream().StartAcceptingReads(body);
 
-            input.Add("Hello", true);
+                input.Add("Hello", true);
 
-            var buffer1 = new byte[1024];
+                var buffer1 = new byte[1024];
-            var count1 = await stream.ReadAsync(buffer1, 0, 1024);
+                var count1 = await stream.ReadAsync(buffer1, 0, 1024);
-            AssertASCII("Hello", new ArraySegment<byte>(buffer1, 0, 5));
+                AssertASCII("Hello", new ArraySegment<byte>(buffer1, 0, 5));
 
-            var buffer2 = new byte[1024];
+                var buffer2 = new byte[1024];
-            var count2 = await stream.ReadAsync(buffer2, 0, 1024);
+                var count2 = await stream.ReadAsync(buffer2, 0, 1024);
-            Assert.Equal(0, count2);
+                Assert.Equal(0, count2);
+            }
         }
 
         [Fact]
         public async Task CanHandleLargeBlocks()
         {
-            var input = new TestInput();
+            using (var input = new TestInput())
-            var body = MessageBody.For("HTTP/1.0", new FrameRequestHeaders(), input.FrameContext);
+            {
-            var stream = new FrameRequestStream().StartAcceptingReads(body);
+                var body = MessageBody.For("HTTP/1.0", new FrameRequestHeaders(), input.FrameContext);
+                var stream = new FrameRequestStream().StartAcceptingReads(body);
 
-            // Input needs to be greater than 4032 bytes to allocate a block not backed by a slab.
+                // Input needs to be greater than 4032 bytes to allocate a block not backed by a slab.
-            var largeInput = new string('a', 8192);
+                var largeInput = new string('a', 8192);
 
-            input.Add(largeInput, true);
+                input.Add(largeInput, true);
-            // Add a smaller block to the end so that SocketInput attempts to return the large
+                // Add a smaller block to the end so that SocketInput attempts to return the large
-            // block to the memory pool.
+                // block to the memory pool.
-            input.Add("Hello", true);
+                input.Add("Hello", true);
 
-            var readBuffer = new byte[8192];
+                var readBuffer = new byte[8192];
 
-            var count1 = await stream.ReadAsync(readBuffer, 0, 8192);
+                var count1 = await stream.ReadAsync(readBuffer, 0, 8192);
-            Assert.Equal(8192, count1);
+                Assert.Equal(8192, count1);
-            AssertASCII(largeInput, new ArraySegment<byte>(readBuffer, 0, 8192));
+                AssertASCII(largeInput, new ArraySegment<byte>(readBuffer, 0, 8192));
 
-            var count2 = await stream.ReadAsync(readBuffer, 0, 8192);
+                var count2 = await stream.ReadAsync(readBuffer, 0, 8192);
-            Assert.Equal(5, count2);
+                Assert.Equal(5, count2);
-            AssertASCII("Hello", new ArraySegment<byte>(readBuffer, 0, 5));
+                AssertASCII("Hello", new ArraySegment<byte>(readBuffer, 0, 5));
 
-            var count3 = await stream.ReadAsync(readBuffer, 0, 8192);
+                var count3 = await stream.ReadAsync(readBuffer, 0, 8192);
-            Assert.Equal(0, count3);
+                Assert.Equal(0, count3);
+            }
         }
 
         private void AssertASCII(string expected, ArraySegment<byte> actual)

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

@@ -19,14 +19,12 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             var trace = new KestrelTrace(new TestKestrelTrace());
             var ltp = new LoggingThreadPool(trace);
             using (var memory2 = new MemoryPool2())
+            using (var socketInput = new SocketInput(memory2, ltp))
             {
-                var socketInput = new SocketInput(memory2, ltp);
-
                 var task0Threw = false;
                 var task1Threw = false;
                 var task2Threw = false;
 
-
                 var task0 = AwaitAsTaskAsync(socketInput);
 
                 Assert.False(task0.IsFaulted);
@@ -84,19 +82,25 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
             var defultIter = new MemoryPoolIterator2();
 
             // Calling ConsumingComplete without a preceding calling to ConsumingStart fails
-            var socketInput = new SocketInput(null, null);
+            using (var socketInput = new SocketInput(null, null))
-            Assert.Throws<InvalidOperationException>(() => socketInput.ConsumingComplete(defultIter, defultIter));
+            {
-
+                Assert.Throws<InvalidOperationException>(() => socketInput.ConsumingComplete(defultIter, defultIter));
-            // Calling ConsumingStart twice in a row fails
+            }
-            socketInput = new SocketInput(null, null);
-            socketInput.ConsumingStart();
-            Assert.Throws<InvalidOperationException>(() => socketInput.ConsumingStart());
 
             // Calling ConsumingComplete twice in a row fails
-            socketInput = new SocketInput(null, null);
+            using (var socketInput = new SocketInput(null, null))
-            socketInput.ConsumingStart();
+            {
-            socketInput.ConsumingComplete(defultIter, defultIter);
+                socketInput.ConsumingStart();
-            Assert.Throws<InvalidOperationException>(() => socketInput.ConsumingComplete(defultIter, defultIter));
+                socketInput.ConsumingComplete(defultIter, defultIter);
+                Assert.Throws<InvalidOperationException>(() => socketInput.ConsumingComplete(defultIter, defultIter));
+            }
+
+            // Calling ConsumingStart twice in a row fails
+            using (var socketInput = new SocketInput(null, null))
+            {
+                socketInput.ConsumingStart();
+                Assert.Throws<InvalidOperationException>(() => socketInput.ConsumingStart());
+            }
         }
 
         private static void TestConcurrentFaultedTask(Task t)
@@ -110,6 +114,5 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
         {
             await socketInput;
         }
-
     }
 }

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

@@ -33,8 +33,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 }
             };
 
-            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             using (var memory = new MemoryPool2())
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             {
                 kestrelEngine.Start(count: 1);
 
@@ -42,7 +42,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 var socket = new MockSocket(kestrelThread.Loop.ThreadId, new TestKestrelTrace());
                 var trace = new KestrelTrace(new TestKestrelTrace());
                 var ltp = new LoggingThreadPool(trace);
-                var socketOutput = new SocketOutput(kestrelThread, socket, memory, null, 0, trace, ltp, new Queue<UvWriteReq>());
+                var socketOutput = new SocketOutput(kestrelThread, socket, memory, new MockConnection(), 0, trace, ltp, new Queue<UvWriteReq>());
 
                 // I doubt _maxBytesPreCompleted will ever be over a MB. If it is, we should change this test.
                 var bufferSize = 1048576;
@@ -60,6 +60,10 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
 
                 // Assert
                 Assert.True(completedWh.Wait(1000));
+
+                // Cleanup
+                var cleanupTask = socketOutput.WriteAsync(
+                    default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
             }
         }
 
@@ -80,8 +84,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 }
             };
 
-            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             using (var memory = new MemoryPool2())
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             {
                 kestrelEngine.Start(count: 1);
 
@@ -89,7 +93,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 var socket = new MockSocket(kestrelThread.Loop.ThreadId, new TestKestrelTrace());
                 var trace = new KestrelTrace(new TestKestrelTrace());
                 var ltp = new LoggingThreadPool(trace);
-                var socketOutput = new SocketOutput(kestrelThread, socket, memory, null, 0, trace, ltp, new Queue<UvWriteReq>());
+                var socketOutput = new SocketOutput(kestrelThread, socket, memory, new MockConnection(), 0, trace, ltp, new Queue<UvWriteReq>());
 
                 var bufferSize = maxBytesPreCompleted;
                 var buffer = new ArraySegment<byte>(new byte[bufferSize], 0, bufferSize);
@@ -117,6 +121,15 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 // Assert
                 // Finishing the first write should allow the second write to pre-complete.
                 Assert.True(completedWh.Wait(1000));
+
+                // Cleanup
+                var cleanupTask = socketOutput.WriteAsync(
+                    default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
+
+                foreach (var triggerCompleted in completeQueue)
+                {
+                    triggerCompleted(0);
+                }
             }
         }
 
@@ -139,8 +152,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 }
             };
 
-            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             using (var memory = new MemoryPool2())
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             {
                 kestrelEngine.Start(count: 1);
 
@@ -148,7 +161,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 var socket = new MockSocket(kestrelThread.Loop.ThreadId, new TestKestrelTrace());
                 var trace = new KestrelTrace(new TestKestrelTrace());
                 var ltp = new LoggingThreadPool(trace);
-                var socketOutput = new SocketOutput(kestrelThread, socket, memory, null, 0, trace, ltp, new Queue<UvWriteReq>());
+                var socketOutput = new SocketOutput(kestrelThread, socket, memory, new MockConnection(), 0, trace, ltp, new Queue<UvWriteReq>());
 
                 var bufferSize = maxBytesPreCompleted / 2;
                 var data = new byte[bufferSize];
@@ -183,6 +196,15 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 // Assert
                 // Finishing the first write should allow the second write to pre-complete.
                 Assert.True(writeTask2.Wait(1000));
+
+                // Cleanup
+                var cleanupTask = socketOutput.WriteAsync(
+                    default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
+
+                foreach (var triggerCompleted in completeQueue)
+                {
+                    triggerCompleted(0);
+                }
             }
         }
 
@@ -203,8 +225,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 }
             };
 
-            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             using (var memory = new MemoryPool2())
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             {
                 kestrelEngine.Start(count: 1);
 
@@ -212,74 +234,87 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 var socket = new MockSocket(kestrelThread.Loop.ThreadId, new TestKestrelTrace());
                 var trace = new KestrelTrace(new TestKestrelTrace());
                 var ltp = new LoggingThreadPool(trace);
-                ISocketOutput socketOutput = new SocketOutput(kestrelThread, socket, memory, new MockConnection(socket), 0, trace, ltp, new Queue<UvWriteReq>());
                 
-                var bufferSize = maxBytesPreCompleted;
+                using (var mockConnection = new MockConnection())
+                {
+                    ISocketOutput socketOutput = new SocketOutput(kestrelThread, socket, memory, mockConnection, 0, trace, ltp, new Queue<UvWriteReq>());
 
-                var data = new byte[bufferSize];
+                    var bufferSize = maxBytesPreCompleted;
-                var fullBuffer = new ArraySegment<byte>(data, 0, bufferSize);
 
-                var cts = new CancellationTokenSource();
+                    var data = new byte[bufferSize];
+                    var fullBuffer = new ArraySegment<byte>(data, 0, bufferSize);
 
-                // Act 
+                    var cts = new CancellationTokenSource();
-                var task1Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: cts.Token);
-                // task1 should complete successfully as < _maxBytesPreCompleted
 
-                // First task is completed and successful
+                    // Act 
-                Assert.True(task1Success.IsCompleted);
+                    var task1Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: cts.Token);
-                Assert.False(task1Success.IsCanceled);
+                    // task1 should complete successfully as < _maxBytesPreCompleted
-                Assert.False(task1Success.IsFaulted);
 
-                // following tasks should wait.
+                    // First task is completed and successful
-                var task2Throw = socketOutput.WriteAsync(fullBuffer, cancellationToken: cts.Token);
+                    Assert.True(task1Success.IsCompleted);
-                var task3Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: default(CancellationToken));
+                    Assert.False(task1Success.IsCanceled);
+                    Assert.False(task1Success.IsFaulted);
 
-                // Give time for tasks to percolate
+                    // following tasks should wait.
-                await Task.Delay(1000);
+                    var task2Throw = socketOutput.WriteAsync(fullBuffer, cancellationToken: cts.Token);
+                    var task3Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: default(CancellationToken));
 
-                // Second task is not completed
+                    // Give time for tasks to percolate
-                Assert.False(task2Throw.IsCompleted);
+                    await Task.Delay(1000);
-                Assert.False(task2Throw.IsCanceled);
-                Assert.False(task2Throw.IsFaulted);
 
-                // Third task is not completed 
+                    // Second task is not completed
-                Assert.False(task3Success.IsCompleted);
+                    Assert.False(task2Throw.IsCompleted);
-                Assert.False(task3Success.IsCanceled);
+                    Assert.False(task2Throw.IsCanceled);
-                Assert.False(task3Success.IsFaulted);
+                    Assert.False(task2Throw.IsFaulted);
 
-                cts.Cancel();
+                    // Third task is not completed 
+                    Assert.False(task3Success.IsCompleted);
+                    Assert.False(task3Success.IsCanceled);
+                    Assert.False(task3Success.IsFaulted);
 
-                // Second task is now canceled
+                    cts.Cancel();
-                await Assert.ThrowsAsync<TaskCanceledException>(() => task2Throw);
-                Assert.True(task2Throw.IsCanceled);
 
-                // Third task is now completed
+                    // Second task is now canceled
-                await task3Success;
+                    await Assert.ThrowsAsync<TaskCanceledException>(() => task2Throw);
+                    Assert.True(task2Throw.IsCanceled);
 
-                // Fourth task immediately cancels as the token is canceled 
+                    // Third task is now completed
-                var task4Throw = socketOutput.WriteAsync(fullBuffer, cancellationToken: cts.Token);
+                    await task3Success;
 
-                Assert.True(task4Throw.IsCompleted);
+                    // Fourth task immediately cancels as the token is canceled 
-                Assert.True(task4Throw.IsCanceled);
+                    var task4Throw = socketOutput.WriteAsync(fullBuffer, cancellationToken: cts.Token);
-                Assert.False(task4Throw.IsFaulted);
 
-                var task5Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: default(CancellationToken));
+                    Assert.True(task4Throw.IsCompleted);
-                // task5 should complete immediately
+                    Assert.True(task4Throw.IsCanceled);
+                    Assert.False(task4Throw.IsFaulted);
 
-                Assert.True(task5Success.IsCompleted);
+                    var task5Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: default(CancellationToken));
-                Assert.False(task5Success.IsCanceled);
+                    // task5 should complete immediately
-                Assert.False(task5Success.IsFaulted);
 
-                cts = new CancellationTokenSource();
+                    Assert.True(task5Success.IsCompleted);
+                    Assert.False(task5Success.IsCanceled);
+                    Assert.False(task5Success.IsFaulted);
 
-                var task6Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: cts.Token);
+                    cts = new CancellationTokenSource();
-                // task6 should complete immediately but not cancel as its cancellation token isn't set
 
-                Assert.True(task6Success.IsCompleted);
+                    var task6Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: cts.Token);
-                Assert.False(task6Success.IsCanceled);
+                    // task6 should complete immediately but not cancel as its cancellation token isn't set
-                Assert.False(task6Success.IsFaulted);
 
-                Assert.True(true);
+                    Assert.True(task6Success.IsCompleted);
+                    Assert.False(task6Success.IsCanceled);
+                    Assert.False(task6Success.IsFaulted);
+
+                    Assert.True(true);
+
+                    // Cleanup
+                    var cleanupTask = ((SocketOutput)socketOutput).WriteAsync(
+                        default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
+
+                    foreach (var triggerCompleted in completeQueue)
+                    {
+                        triggerCompleted(0);
+                    }
+                }
             }
         }
 
@@ -300,9 +335,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 }
             };
 
-            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             using (var memory = new MemoryPool2())
-            using (var abortedSource = new CancellationTokenSource())
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             {
                 kestrelEngine.Start(count: 1);
 
@@ -311,50 +345,61 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 var trace = new KestrelTrace(new TestKestrelTrace());
                 var ltp = new LoggingThreadPool(trace);
 
-                var mockConnection = new MockConnection(socket);
+                using (var mockConnection = new MockConnection())
-                mockConnection.RequestAbortedSource = abortedSource;
+                {
-                ISocketOutput socketOutput = new SocketOutput(kestrelThread, socket, memory, mockConnection, 0, trace, ltp, new Queue<UvWriteReq>());
+                    var abortedSource = mockConnection.RequestAbortedSource;
+                    ISocketOutput socketOutput = new SocketOutput(kestrelThread, socket, memory, mockConnection, 0, trace, ltp, new Queue<UvWriteReq>());
 
-                var bufferSize = maxBytesPreCompleted;
+                    var bufferSize = maxBytesPreCompleted;
 
-                var data = new byte[bufferSize];
+                    var data = new byte[bufferSize];
-                var fullBuffer = new ArraySegment<byte>(data, 0, bufferSize);
+                    var fullBuffer = new ArraySegment<byte>(data, 0, bufferSize);
 
-                // Act 
+                    // Act 
-                var task1Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: abortedSource.Token);
+                    var task1Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: abortedSource.Token);
-                // task1 should complete successfully as < _maxBytesPreCompleted
+                    // task1 should complete successfully as < _maxBytesPreCompleted
 
-                // First task is completed and successful
+                    // First task is completed and successful
-                Assert.True(task1Success.IsCompleted);
+                    Assert.True(task1Success.IsCompleted);
-                Assert.False(task1Success.IsCanceled);
+                    Assert.False(task1Success.IsCanceled);
-                Assert.False(task1Success.IsFaulted);
+                    Assert.False(task1Success.IsFaulted);
 
-                // following tasks should wait.
+                    // following tasks should wait.
-                var task2Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: default(CancellationToken));
+                    var task2Success = socketOutput.WriteAsync(fullBuffer, cancellationToken: default(CancellationToken));
-                var task3Canceled = socketOutput.WriteAsync(fullBuffer, cancellationToken: abortedSource.Token);
+                    var task3Canceled = socketOutput.WriteAsync(fullBuffer, cancellationToken: abortedSource.Token);
 
-                // Give time for tasks to percolate
+                    // Give time for tasks to percolate
-                await Task.Delay(1000);
+                    await Task.Delay(1000);
 
-                // Second task is not completed
+                    // Second task is not completed
-                Assert.False(task2Success.IsCompleted);
+                    Assert.False(task2Success.IsCompleted);
-                Assert.False(task2Success.IsCanceled);
+                    Assert.False(task2Success.IsCanceled);
-                Assert.False(task2Success.IsFaulted);
+                    Assert.False(task2Success.IsFaulted);
 
-                // Third task is not completed 
+                    // Third task is not completed 
-                Assert.False(task3Canceled.IsCompleted);
+                    Assert.False(task3Canceled.IsCompleted);
-                Assert.False(task3Canceled.IsCanceled);
+                    Assert.False(task3Canceled.IsCanceled);
-                Assert.False(task3Canceled.IsFaulted);
+                    Assert.False(task3Canceled.IsFaulted);
 
-                // Cause the first write to fail.
+                    // Cause the first write to fail.
-                completeQueue.Dequeue()(-1);
+                    completeQueue.Dequeue()(-1);
 
-                // Second task is now completed
+                    // Second task is now completed
-                await task2Success;
+                    await task2Success;
 
-                // Third task is now canceled
+                    // Third task is now canceled
-                await Assert.ThrowsAsync<TaskCanceledException>(() => task3Canceled);
+                    await Assert.ThrowsAsync<TaskCanceledException>(() => task3Canceled);
-                Assert.True(task3Canceled.IsCanceled);
+                    Assert.True(task3Canceled.IsCanceled);
+
+                    // Cleanup
+                    var cleanupTask = ((SocketOutput)socketOutput).WriteAsync(
+                        default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
+
+                    foreach (var triggerCompleted in completeQueue)
+                    {
+                        triggerCompleted(0);
+                    }
+                }
             }
         }
 
@@ -378,8 +423,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 }
             };
 
-            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             using (var memory = new MemoryPool2())
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             {
                 kestrelEngine.Start(count: 1);
 
@@ -387,7 +432,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 var socket = new MockSocket(kestrelThread.Loop.ThreadId, new TestKestrelTrace());
                 var trace = new KestrelTrace(new TestKestrelTrace());
                 var ltp = new LoggingThreadPool(trace);
-                var socketOutput = new SocketOutput(kestrelThread, socket, memory, null, 0, trace, ltp, new Queue<UvWriteReq>());
+                var socketOutput = new SocketOutput(kestrelThread, socket, memory, new MockConnection(), 0, trace, ltp, new Queue<UvWriteReq>());
 
                 var bufferSize = maxBytesPreCompleted;
                 var buffer = new ArraySegment<byte>(new byte[bufferSize], 0, bufferSize);
@@ -435,6 +480,15 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 // Assert
                 // Finishing the first write should allow the second write to pre-complete.
                 Assert.True(completedWh2.Wait(1000));
+
+                // Cleanup
+                var cleanupTask = ((SocketOutput)socketOutput).WriteAsync(
+                    default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
+
+                foreach (var triggerCompleted in completeQueue)
+                {
+                    triggerCompleted(0);
+                }
             }
         }
 
@@ -455,8 +509,8 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 }
             };
 
-            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             using (var memory = new MemoryPool2())
+            using (var kestrelEngine = new KestrelEngine(mockLibuv, new TestServiceContext()))
             {
                 kestrelEngine.Start(count: 1);
 
@@ -464,7 +518,7 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
                 var socket = new MockSocket(kestrelThread.Loop.ThreadId, new TestKestrelTrace());
                 var trace = new KestrelTrace(new TestKestrelTrace());
                 var ltp = new LoggingThreadPool(trace);
-                var socketOutput = new SocketOutput(kestrelThread, socket, memory, null, 0, trace, ltp, new Queue<UvWriteReq>());
+                var socketOutput = new SocketOutput(kestrelThread, socket, memory, new MockConnection(), 0, trace, ltp, new Queue<UvWriteReq>());
 
                 // block 1
                 var start = socketOutput.ProducingStart();
@@ -484,6 +538,10 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
 
                 Assert.True(nBufferWh.Wait(1000));
                 Assert.Equal(2, nBuffers);
+
+                // Cleanup
+                var cleanupTask = socketOutput.WriteAsync(
+                    default(ArraySegment<byte>), default(CancellationToken), socketDisconnect: true);
             }
         }
 

test/Microsoft.AspNetCore.Server.KestrelTests/TestHelpers/MockConnection.cs

@@ -1,16 +1,20 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
 using System;
 using System.Threading;
-using Microsoft.AspNetCore.Server.Kestrel.Networking;
+using Microsoft.AspNetCore.Server.Kestrel;
 using Microsoft.AspNetCore.Server.Kestrel.Http;
+using Microsoft.AspNetCore.Server.Kestrel.Infrastructure;
+using Microsoft.AspNetCore.Server.Kestrel.Networking;
 
 namespace Microsoft.AspNetCore.Server.KestrelTests.TestHelpers
 {
-    public class MockConnection : Connection
+    public class MockConnection : Connection, IDisposable
     {
-        public MockConnection(UvStreamHandle socket)
+        public MockConnection()
-            : base (new ListenerContext(), socket)
         {
-
+            RequestAbortedSource = new CancellationTokenSource();
         }
 
         public override void Abort()
@@ -21,6 +25,15 @@ namespace Microsoft.AspNetCore.Server.KestrelTests.TestHelpers
             }
         }
 
-        public CancellationTokenSource RequestAbortedSource { get; set; }
+        public override void OnSocketClosed()
+        {
+        }
+
+        public CancellationTokenSource RequestAbortedSource { get; }
+
+        public void Dispose()
+        {
+            RequestAbortedSource.Dispose();
+        }
     }
 }

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

@@ -10,19 +10,22 @@ using Microsoft.AspNetCore.Server.Kestrel.Infrastructure;
 
 namespace Microsoft.AspNetCore.Server.KestrelTests
 {
-    class TestInput : IConnectionControl, IFrameControl
+    class TestInput : IConnectionControl, IFrameControl, IDisposable
     {
+        private MemoryPool2 _memoryPool;
+
         public TestInput()
         {
             var trace = new KestrelTrace(new TestKestrelTrace());
             var ltp = new LoggingThreadPool(trace);
-            var memory2 = new MemoryPool2();
             FrameContext = new FrameContext
             {
-                SocketInput = new SocketInput(memory2, ltp),
                 ConnectionControl = this,
                 FrameControl = this
             };
+
+            _memoryPool = new MemoryPool2();
+            FrameContext.SocketInput = new SocketInput(_memoryPool, ltp);
         }
 
         public FrameContext FrameContext { get; set; }
@@ -81,6 +84,12 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
         {
             return Task.FromResult(0);
         }
+
+        public void Dispose()
+        {
+            FrameContext.SocketInput.Dispose();
+            _memoryPool.Dispose();
+        }
     }
 }
 

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

@@ -11,7 +11,6 @@ namespace Microsoft.AspNetCore.Server.KestrelTests
 {
     public class UrlPathDecoderTests
     {
-
         [Fact]
         public void Empty()
         {