React to pool cleanup (#2526)

Directory.Build.props

@@ -7,6 +7,10 @@
   <Import Project="build\dependencies.props" />
   <Import Project="build\sources.props" />
 
+  <PropertyGroup>
+    <DefineConstants Condition="'$(InnerLoop)' != ''">$(DefineConstants);INNER_LOOP</DefineConstants>
+  </PropertyGroup>
+
   <PropertyGroup>
     <Product>Microsoft ASP.NET Core</Product>
     <RepositoryUrl>https://github.com/aspnet/KestrelHttpServer</RepositoryUrl>

benchmarks/Kestrel.Performance/ResponseHeadersWritingBenchmark.cs

@@ -26,6 +26,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
 
         private MemoryPool<byte> _memoryPool;
 
+        private DuplexPipe.DuplexPipePair _pair;
+
         [Params(
             BenchmarkTypes.TechEmpowerPlaintext,
             BenchmarkTypes.PlaintextChunked,
@@ -115,7 +117,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
         {
             _memoryPool = KestrelMemoryPool.Create();
             var options = new PipeOptions(_memoryPool, readerScheduler: PipeScheduler.Inline, writerScheduler: PipeScheduler.Inline, useSynchronizationContext: false);
-            var pair = DuplexPipe.CreateConnectionPair(options, options);
+            _pair = DuplexPipe.CreateConnectionPair(options, options);
 
             var serviceContext = new ServiceContext
             {
@@ -131,8 +133,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
                 ConnectionFeatures = new FeatureCollection(),
                 MemoryPool = _memoryPool,
                 TimeoutControl = new MockTimeoutControl(),
-                Application = pair.Application,
-                Transport = pair.Transport
+                Application = _pair.Application,
+                Transport = _pair.Transport
             });
 
             http1Connection.Reset();
@@ -143,6 +145,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
         [IterationCleanup]
         public void Cleanup()
         {
+            _pair.Application.Input.Complete();
+            _pair.Application.Output.Complete();
+            _pair.Transport.Input.Complete();
+            _pair.Transport.Output.Complete();
             _memoryPool.Dispose();
         }
 

build/dependencies.props

@@ -18,7 +18,7 @@
     <MicrosoftAspNetCoreTestingPackageVersion>2.2.0-preview1-34066</MicrosoftAspNetCoreTestingPackageVersion>
     <MicrosoftAspNetCoreWebUtilitiesPackageVersion>2.2.0-preview1-34066</MicrosoftAspNetCoreWebUtilitiesPackageVersion>
     <MicrosoftExtensionsActivatorUtilitiesSourcesPackageVersion>2.2.0-preview1-34066</MicrosoftExtensionsActivatorUtilitiesSourcesPackageVersion>
-    <MicrosoftExtensionsBuffersMemoryPoolSourcesPackageVersion>2.2.0-preview1-34066</MicrosoftExtensionsBuffersMemoryPoolSourcesPackageVersion>
+    <MicrosoftExtensionsBuffersMemoryPoolSourcesPackageVersion>2.2.0-a-preview1-pool-cleanup-16484</MicrosoftExtensionsBuffersMemoryPoolSourcesPackageVersion>
     <MicrosoftExtensionsBuffersSourcesPackageVersion>2.2.0-preview1-34066</MicrosoftExtensionsBuffersSourcesPackageVersion>
     <MicrosoftExtensionsBuffersTestingSourcesPackageVersion>2.2.0-preview1-34066</MicrosoftExtensionsBuffersTestingSourcesPackageVersion>
     <MicrosoftExtensionsConfigurationBinderPackageVersion>2.2.0-preview1-34066</MicrosoftExtensionsConfigurationBinderPackageVersion>

src/Directory.Build.props

@@ -1,8 +1,5 @@
 <Project>
   <Import Project="..\Directory.Build.props" />
-  <PropertyGroup>
-    <DefineConstants Condition="'$(InnerLoop)' != ''">$(DefineConstants);INNER_LOOP</DefineConstants>
-  </PropertyGroup>
   <ItemGroup>
     <PackageReference Include="Internal.AspNetCore.Sdk" PrivateAssets="All" Version="$(InternalAspNetCoreSdkPackageVersion)" />
   </ItemGroup>

src/Kestrel.Core/Adapter/Internal/LoggingStream.cs

@@ -148,6 +148,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Adapter.Internal
 
         private void Log(string method, ReadOnlySpan<byte> buffer)
         {
+            if (!_logger.IsEnabled(LogLevel.Debug))
+            {
+                return;
+            }
+
             var builder = new StringBuilder($"{method}[{buffer.Length}] ");
 
             // Write the hex

src/Kestrel.Transport.Abstractions/Internal/KestrelMemoryPool.cs

@@ -1,13 +1,26 @@
 // 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.Buffers;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal
 {
     public static class KestrelMemoryPool
     {
-        public static MemoryPool<byte> Create() => new SlabMemoryPool();
+        public static MemoryPool<byte> Create()
+        {
+#if DEBUG
+            return new DiagnosticMemoryPool(CreateSlabMemoryPool());
+#else
+            return CreateSlabMemoryPool();
+#endif
+        }
+
+        public static MemoryPool<byte> CreateSlabMemoryPool()
+        {
+            return new SlabMemoryPool();
+        }
 
         public static readonly int MinimumSegmentSize = 4096;
     }

src/Kestrel.Transport.Libuv/AssemblyInfo.cs

@@ -0,0 +1,6 @@
+// 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.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("Libuv.FunctionalTests, PublicKey=0024000004800000940000000602000000240000525341310004000001000100f33a29044fa9d740c9b3213a93e57c84b472c84e0b8a0e1ae48e67a9f8f6de9d5f7f3d52ac23e48ac51801f1dc950abe901da34d2a9e3baadb141a17c77ef3c565dd5ee5054b91cf63bb3c6ab83f72ab3aafe93d0fc3c2348b764fafb0b1c0733de51459aeab46580384bf9d74c4e28164b7cde247f891ba07891c9d872ad2bb")]

src/Kestrel.Transport.Libuv/Internal/LibuvThread.cs

@@ -38,7 +38,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
         private readonly object _startSync = new object();
         private bool _stopImmediate = false;
         private bool _initCompleted = false;
-        private ExceptionDispatchInfo _closeError;
+        private Exception _closeError;
         private readonly ILibuvTrace _log;
 
         public LibuvThread(LibuvTransport transport)
@@ -61,7 +61,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
 #endif
             QueueCloseHandle = PostCloseHandle;
             QueueCloseAsyncHandle = EnqueueCloseHandle;
-            MemoryPool = KestrelMemoryPool.Create();
+            MemoryPool = transport.TransportOptions.MemoryPoolFactory();
             WriteReqPool = new WriteReqPool(this, _log);
         }
 
@@ -82,7 +82,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
         public List<WeakReference> Requests { get; } = new List<WeakReference>();
 #endif
 
-        public ExceptionDispatchInfo FatalError { get { return _closeError; } }
+        public Exception FatalError => _closeError;
 
         public Action<Action<IntPtr>, IntPtr> QueueCloseHandle { get; }
 
@@ -133,7 +133,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
                 }
             }
 
-            _closeError?.Throw();
+            if (_closeError != null)
+            {
+                ExceptionDispatchInfo.Capture(_closeError).Throw();
+            }
         }
 
 #if DEBUG && !INNER_LOOP
@@ -323,14 +326,21 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Internal
             }
             catch (Exception ex)
             {
-                _closeError = ExceptionDispatchInfo.Capture(ex);
+                _closeError = ex;
                 // Request shutdown so we can rethrow this exception
                 // in Stop which should be observable.
                 _appLifetime.StopApplication();
             }
             finally
             {
-                MemoryPool.Dispose();
+                try
+                {
+                    MemoryPool.Dispose();
+                }
+                catch (Exception ex)
+                {
+                    _closeError = _closeError == null ? ex : new AggregateException(_closeError, ex);
+                }
                 WriteReqPool.Dispose();
                 _threadTcs.SetResult(null);
 

src/Kestrel.Transport.Libuv/LibuvTransportOptions.cs

@@ -2,6 +2,8 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Buffers;
+using Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv
 {
@@ -18,6 +20,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv
         /// </remarks>
         public int ThreadCount { get; set; } = ProcessorThreadCount;
 
+        internal Func<MemoryPool<byte>> MemoryPoolFactory { get; set; } = () => KestrelMemoryPool.Create();
+
         private static int ProcessorThreadCount
         {
             get

src/Kestrel.Transport.Sockets/AssemblyInfo.cs

@@ -0,0 +1,6 @@
+// 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.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("Sockets.FunctionalTests, PublicKey=0024000004800000940000000602000000240000525341310004000001000100f33a29044fa9d740c9b3213a93e57c84b472c84e0b8a0e1ae48e67a9f8f6de9d5f7f3d52ac23e48ac51801f1dc950abe901da34d2a9e3baadb141a17c77ef3c565dd5ee5054b91cf63bb3c6ab83f72ab3aafe93d0fc3c2348b764fafb0b1c0733de51459aeab46580384bf9d74c4e28164b7cde247f891ba07891c9d872ad2bb")]

src/Kestrel.Transport.Sockets/SocketTransport.cs

@@ -22,7 +22,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
     {
         private static readonly PipeScheduler[] ThreadPoolSchedulerArray = new PipeScheduler[] { PipeScheduler.ThreadPool };
 
-        private readonly MemoryPool<byte> _memoryPool = KestrelMemoryPool.Create();
+        private readonly MemoryPool<byte> _memoryPool;
         private readonly IEndPointInformation _endPointInformation;
         private readonly IConnectionDispatcher _dispatcher;
         private readonly IApplicationLifetime _appLifetime;
@@ -39,7 +39,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
             IConnectionDispatcher dispatcher,
             IApplicationLifetime applicationLifetime,
             int ioQueueCount,
-            ISocketsTrace trace)
+            ISocketsTrace trace,
+            MemoryPool<byte> memoryPool)
         {
             Debug.Assert(endPointInformation != null);
             Debug.Assert(endPointInformation.Type == ListenType.IPEndPoint);
@@ -51,6 +52,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
             _dispatcher = dispatcher;
             _appLifetime = applicationLifetime;
             _trace = trace;
+            _memoryPool = memoryPool;
 
             if (ioQueueCount > 0)
             {

src/Kestrel.Transport.Sockets/SocketTransportFactory.cs

@@ -57,7 +57,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
                 throw new ArgumentNullException(nameof(dispatcher));
             }
 
-            return new SocketTransport(endPointInformation, dispatcher, _appLifetime, _options.IOQueueCount, _trace);
+            return new SocketTransport(endPointInformation, dispatcher, _appLifetime, _options.IOQueueCount, _trace, _options.MemoryPoolFactory());
         }
     }
 }

src/Kestrel.Transport.Sockets/SocketTransportOptions.cs

@@ -2,6 +2,8 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Buffers;
+using Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
 {
@@ -14,5 +16,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets
         /// Defaults to <see cref="Environment.ProcessorCount" /> rounded down and clamped between 1 and 16.
         /// </remarks>
         public int IOQueueCount { get; set; } = Math.Min(Environment.ProcessorCount, 16);
+
+        internal Func<MemoryPool<byte>> MemoryPoolFactory { get; set; } = () => KestrelMemoryPool.Create();
     }
 }

test/Kestrel.Core.Tests/Http2ConnectionTests.cs

@@ -287,6 +287,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
         public void Dispose()
         {
+            _pair.Application.Input.Complete();
+            _pair.Application.Output.Complete();
+            _pair.Transport.Input.Complete();
+            _pair.Transport.Output.Complete();
             _memoryPool.Dispose();
         }
 
@@ -1213,7 +1217,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
                 withLength: 0,
                 withFlags: (byte)Http2DataFrameFlags.END_STREAM,
                 withStreamId: 1);
-            
+
             await StopConnectionAsync(expectedLastStreamId: 1, ignoreNonGoAwayFrames: false);
 
             _hpackDecoder.Decode(headersFrame.HeadersPayload, endHeaders: false, handler: this);

test/Kestrel.Core.Tests/PipelineExtensionTests.cs

@@ -26,6 +26,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
         public void Dispose()
         {
+            _pipe.Reader.Complete();
+            _pipe.Writer.Complete();
             _memoryPool.Dispose();
         }
 

test/Kestrel.Core.Tests/TestInput.cs

@@ -68,6 +68,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
         public void Dispose()
         {
+            Application.Input.Complete();
+            Application.Output.Complete();
+            Transport.Input.Complete();
+            Transport.Output.Complete();
             _memoryPool.Dispose();
         }
     }

test/Kestrel.FunctionalTests/ChunkedRequestTests.cs

@@ -12,11 +12,16 @@ using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.Extensions.Logging.Testing;
 using Xunit;
+using Xunit.Abstractions;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
 {
     public class ChunkedRequestTests : LoggedTest
     {
+        public ChunkedRequestTests(ITestOutputHelper output) : base(output)
+        {
+        }
+
         public static TheoryData<ListenOptions> ConnectionAdapterData => new TheoryData<ListenOptions>
         {
             new ListenOptions(new IPEndPoint(IPAddress.Loopback, 0)),
@@ -88,7 +93,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         [MemberData(nameof(ConnectionAdapterData))]
         public async Task Http10KeepAliveTransferEncoding(ListenOptions listenOptions)
         {
-            var testContext = new TestServiceContext(LoggerFactory);
+            var testContext = new TestServiceContext();
 
             using (var server = new TestServer(AppChunked, testContext, listenOptions))
             {

test/Kestrel.FunctionalTests/DiagnosticMemoryPoolFactory.cs

@@ -0,0 +1,46 @@
+// 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.Buffers;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
+{
+    public class DiagnosticMemoryPoolFactory
+    {
+        private readonly bool _allowLateReturn;
+
+        private readonly bool _rentTracking;
+
+        private readonly List<DiagnosticMemoryPool> _pools;
+
+        public DiagnosticMemoryPoolFactory(bool allowLateReturn = false, bool rentTracking = false)
+        {
+            _allowLateReturn = allowLateReturn;
+            _rentTracking = rentTracking;
+            _pools = new List<DiagnosticMemoryPool>();
+        }
+
+        public MemoryPool<byte> Create()
+        {
+            lock (_pools)
+            {
+                var pool = new DiagnosticMemoryPool(KestrelMemoryPool.CreateSlabMemoryPool(),_allowLateReturn, _rentTracking);
+                _pools.Add(pool);
+                return pool;
+            }
+        }
+
+        public Task WhenAllBlocksReturned(TimeSpan span)
+        {
+            lock (_pools)
+            {
+                return Task.WhenAll(_pools.Select(p=>p.WhenAllBlocksReturnedAsync(span)));
+            }
+        }
+    }
+}

test/Kestrel.FunctionalTests/MaxRequestBufferSizeTests.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Buffers;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -90,7 +91,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             var clientFinishedSendingRequestBody = new TaskCompletionSource<object>();
             var lastBytesWritten = DateTime.MaxValue;
 
-            using (var host = StartWebHost(maxRequestBufferSize, data, connectionAdapter, startReadingRequestBody, clientFinishedSendingRequestBody))
+            var memoryPoolFactory = new DiagnosticMemoryPoolFactory(allowLateReturn: true);
+
+            using (var host = StartWebHost(maxRequestBufferSize, data, connectionAdapter, startReadingRequestBody, clientFinishedSendingRequestBody, memoryPoolFactory.Create))
             {
                 var port = host.GetPort();
                 using (var socket = CreateSocket(port))
@@ -167,6 +170,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
                     }
                 }
             }
+
+            await memoryPoolFactory.WhenAllBlocksReturned(TestConstants.DefaultTimeout);
         }
 
         [Fact]
@@ -182,7 +187,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             var clientFinishedSendingRequestBody = new TaskCompletionSource<object>();
             var lastBytesWritten = DateTime.MaxValue;
 
-            using (var host = StartWebHost(16 * 1024, data, false, startReadingRequestBody, clientFinishedSendingRequestBody))
+            var memoryPoolFactory = new DiagnosticMemoryPoolFactory(allowLateReturn: true);
+
+            using (var host = StartWebHost(16 * 1024, data, false, startReadingRequestBody, clientFinishedSendingRequestBody, memoryPoolFactory.Create))
             {
                 var port = host.GetPort();
                 using (var socket = CreateSocket(port))
@@ -239,15 +246,20 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
                     host.Dispose();
                 }
             }
+            // Allow appfunc to unblock
+            startReadingRequestBody.SetResult(null);
+            clientFinishedSendingRequestBody.SetResult(null);
+            await memoryPoolFactory.WhenAllBlocksReturned(TestConstants.DefaultTimeout);
         }
 
         private IWebHost StartWebHost(long? maxRequestBufferSize,
             byte[] expectedBody,
             bool useConnectionAdapter,
             TaskCompletionSource<object> startReadingRequestBody,
-            TaskCompletionSource<object> clientFinishedSendingRequestBody)
+            TaskCompletionSource<object> clientFinishedSendingRequestBody,
+            Func<MemoryPool<byte>> memoryPoolFactory = null)
         {
-            var host = TransportSelector.GetWebHostBuilder()
+            var host = TransportSelector.GetWebHostBuilder(memoryPoolFactory)
                 .ConfigureServices(AddTestLogging)
                 .UseKestrel(options =>
                 {

test/Kestrel.FunctionalTests/RequestTests.cs

@@ -590,9 +590,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         [InlineData("http://localhost/path%20with%20space?q=abc%20123", "/path with space", "abc 123")]
         public async Task CanHandleRequestsWithUrlInAbsoluteForm(string requestUrl, string expectedPath, string queryValue)
         {
-            var pathTcs = new TaskCompletionSource<PathString>();
-            var rawTargetTcs = new TaskCompletionSource<string>();
-            var queryTcs = new TaskCompletionSource<IQueryCollection>();
+            var pathTcs = new TaskCompletionSource<PathString>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var rawTargetTcs = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var queryTcs = new TaskCompletionSource<IQueryCollection>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             using (var server = new TestServer(async context =>
                  {
@@ -1031,8 +1031,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         {
             var testContext = new TestServiceContext(LoggerFactory);
 
-            var readTcs = new TaskCompletionSource<object>();
-            var registrationTcs = new TaskCompletionSource<int>();
+            var readTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var registrationTcs = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             var requestId = 0;
 
             using (var server = new TestServer(async httpContext =>

test/Kestrel.FunctionalTests/ResponseTests.cs

@@ -33,6 +33,7 @@ using Microsoft.Extensions.Logging.Testing;
 using Microsoft.Extensions.Primitives;
 using Moq;
 using Xunit;
+using Xunit.Abstractions;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
 {
@@ -145,7 +146,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         public async Task OnCompleteCalledEvenWhenOnStartingNotCalled()
         {
             var onStartingCalled = false;
-            var onCompletedTcs = new TaskCompletionSource<object>();
+            var onCompletedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             var hostBuilder = TransportSelector.GetWebHostBuilder()
                 .UseKestrel()
@@ -341,7 +342,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         [Fact]
         public async Task OnCompletedShouldNotBlockAResponse()
         {
-            var delayTcs = new TaskCompletionSource<object>();
+            var delayTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var hostBuilder = TransportSelector.GetWebHostBuilder()
                 .UseKestrel()
                 .UseUrls("http://127.0.0.1:0/")
@@ -375,7 +376,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         [Fact]
         public async Task InvalidChunkedEncodingInRequestShouldNotBlockOnCompleted()
         {
-            var onCompletedTcs = new TaskCompletionSource<object>();
+            var onCompletedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             using (var server = new TestServer(httpContext =>
             {
@@ -418,7 +419,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             mockHttpContextFactory.Setup(f => f.Create(It.IsAny<IFeatureCollection>()))
                 .Returns<IFeatureCollection>(fc => new DefaultHttpContext(fc));
 
-            var disposedTcs = new TaskCompletionSource<int>();
+            var disposedTcs = new TaskCompletionSource<int>(TaskCreationOptions.RunContinuationsAsynchronously);
             mockHttpContextFactory.Setup(f => f.Dispose(It.IsAny<HttpContext>()))
                 .Callback<HttpContext>(c =>
                 {
@@ -619,7 +620,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         {
             const string response = "hello, world";
 
-            var logTcs = new TaskCompletionSource<object>();
+            var logTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var mockKestrelTrace = new Mock<IKestrelTrace>();
             mockKestrelTrace
                 .Setup(trace => trace.ConnectionHeadResponseBodyWrite(It.IsAny<string>(), response.Length))
@@ -808,7 +809,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         [Fact]
         public async Task WhenAppWritesLessThanContentLengthErrorLogged()
         {
-            var logTcs = new TaskCompletionSource<object>();
+            var logTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var mockTrace = new Mock<IKestrelTrace>();
             mockTrace
                 .Setup(trace => trace.ApplicationError(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<InvalidOperationException>()))
@@ -1182,7 +1183,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         {
             var connectionClosed = new ManualResetEventSlim();
             var requestStarted = new ManualResetEventSlim();
-            var tcs = new TaskCompletionSource<object>();
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             using (var server = new TestServer(async httpContext =>
             {
@@ -2235,7 +2236,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             // Ensure string is long enough to disable write-behind buffering
             var largeString = new string('a', maxBytesPreCompleted + 1);
 
-            var writeTcs = new TaskCompletionSource<object>();
+            var writeTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
             var requestAbortedWh = new ManualResetEventSlim();
             var requestStartWh = new ManualResetEventSlim();
 
@@ -2400,7 +2401,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             var testContext= new TestServiceContext(LoggerFactory);
 
             var callOrder = new Stack<int>();
-            var onStartingTcs = new TaskCompletionSource<object>();
+            var onStartingTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             using (var server = new TestServer(async context =>
             {
@@ -2452,7 +2453,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             var testContext= new TestServiceContext(LoggerFactory);
 
             var callOrder = new Stack<int>();
-            var onCompletedTcs = new TaskCompletionSource<object>();
+            var onCompletedTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             using (var server = new TestServer(async context =>
             {

test/Kestrel.Transport.Libuv.FunctionalTests/TransportSelector.cs

@@ -1,15 +1,17 @@
 // 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.Buffers;
 using Microsoft.AspNetCore.Hosting;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
 {
     public static class TransportSelector
     {
-        public static IWebHostBuilder GetWebHostBuilder()
+        public static IWebHostBuilder GetWebHostBuilder(Func<MemoryPool<byte>> memoryPoolFactory = null)
         {
-            return new WebHostBuilder().UseLibuv().ConfigureServices(TestServer.RemoveDevCert);
+            return new WebHostBuilder().UseLibuv(options => { options.MemoryPoolFactory = memoryPoolFactory ?? options.MemoryPoolFactory; }).ConfigureServices(TestServer.RemoveDevCert);
         }
     }
 }

test/Kestrel.Transport.Libuv.Tests/LibuvConnectionTests.cs

@@ -24,7 +24,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
             var transportContext = new TestLibuvTransportContext() { ConnectionDispatcher = mockConnectionDispatcher };
             var transport = new LibuvTransport(mockLibuv, transportContext, null);
             var thread = new LibuvThread(transport);
-
+            Task connectionTask = null;
             try
             {
                 await thread.StartAsync();
@@ -36,7 +36,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
                     };
                     var socket = new MockSocket(mockLibuv, Thread.CurrentThread.ManagedThreadId, transportContext.Log);
                     var connection = new LibuvConnection(listenerContext, socket);
-                    _ = connection.Start();
+                    connectionTask = connection.Start();
 
                     mockLibuv.AllocCallback(socket.InternalGetHandle(), 2048, out var ignored);
                     mockLibuv.ReadCallback(socket.InternalGetHandle(), 0, ref ignored);
@@ -47,6 +47,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
             }
             finally
             {
+                mockConnectionDispatcher.Input.Reader.Complete();
+                mockConnectionDispatcher.Output.Writer.Complete();
+                await connectionTask;
+
                 await thread.StopAsync(TimeSpan.FromSeconds(5));
             }
         }
@@ -108,6 +112,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
             }
             finally
             {
+                mockConnectionDispatcher.Input.Reader.Complete();
+                mockConnectionDispatcher.Output.Writer.Complete();
+
                 await thread.StopAsync(TimeSpan.FromSeconds(5));
             }
         }
@@ -189,6 +196,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
             }
             finally
             {
+                mockConnectionDispatcher.Input.Reader.Complete();
+                mockConnectionDispatcher.Output.Writer.Complete();
+
                 await thread.StopAsync(TimeSpan.FromSeconds(5));
             }
         }
@@ -202,6 +212,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
             var transport = new LibuvTransport(mockLibuv, transportContext, null);
             var thread = new LibuvThread(transport);
 
+            Task connectionTask = null;
             try
             {
                 await thread.StartAsync();
@@ -213,7 +224,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
                     };
                     var socket = new MockSocket(mockLibuv, Thread.CurrentThread.ManagedThreadId, transportContext.Log);
                     var connection = new LibuvConnection(listenerContext, socket);
-                    _ = connection.Start();
+                    connectionTask = connection.Start();
 
                     var ignored = new LibuvFunctions.uv_buf_t();
                     mockLibuv.ReadCallback(socket.InternalGetHandle(), TestConstants.EOF, ref ignored);
@@ -224,6 +235,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
             }
             finally
             {
+                mockConnectionDispatcher.Input.Reader.Complete();
+                mockConnectionDispatcher.Output.Writer.Complete();
+                await connectionTask;
+
                 await thread.StopAsync(TimeSpan.FromSeconds(5));
             }
         }

test/Kestrel.Transport.Libuv.Tests/LibuvTransportTests.cs

@@ -1,12 +1,14 @@
 // 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.Collections.Generic;
 using System.Linq;
 using System.Net;
 using System.Net.Http;
 using System.Net.Sockets;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
@@ -85,6 +87,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
                 }
             }
 
+            Assert.True(await serviceContext.ConnectionManager.CloseAllConnectionsAsync(new CancellationTokenSource(TestConstants.DefaultTimeout).Token));
             await transport.UnbindAsync();
             await transport.StopAsync();
         }

test/Kestrel.Transport.Sockets.FunctionalTests/TransportSelector.cs

@@ -1,15 +1,17 @@
 // 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.Buffers;
 using Microsoft.AspNetCore.Hosting;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
 {
     public static class TransportSelector
     {
-        public static IWebHostBuilder GetWebHostBuilder()
+        public static IWebHostBuilder GetWebHostBuilder(Func<MemoryPool<byte>> memoryPoolFactory = null)
         {
-            return new WebHostBuilder().UseSockets().ConfigureServices(TestServer.RemoveDevCert);
+            return new WebHostBuilder().UseSockets(options => { options.MemoryPoolFactory = memoryPoolFactory ?? options.MemoryPoolFactory; }).ConfigureServices(TestServer.RemoveDevCert);
         }
     }
 }

test/shared/CompositeKestrelTrace.cs

@@ -0,0 +1,187 @@
+// 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.IO.Pipelines;
+using Microsoft.AspNetCore.Server.Kestrel.Core;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.HPack;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.AspNetCore.Testing
+{
+    public class CompositeKestrelTrace: IKestrelTrace
+    {
+        private readonly IKestrelTrace _trace1;
+        private readonly IKestrelTrace _trace2;
+
+        public CompositeKestrelTrace(IKestrelTrace kestrelTrace, KestrelTrace kestrelTrace1)
+        {
+            _trace1 = kestrelTrace;
+            _trace2 = kestrelTrace1;
+        }
+
+        public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception exception, Func<TState, Exception, string> formatter)
+        {
+            _trace1.Log(logLevel, eventId, state, exception, formatter);
+            _trace2.Log(logLevel, eventId, state, exception, formatter);
+        }
+
+        public bool IsEnabled(LogLevel logLevel)
+        {
+            return _trace1.IsEnabled(logLevel) || _trace2.IsEnabled(logLevel);
+        }
+
+        public IDisposable BeginScope<TState>(TState state)
+        {
+            return _trace1.BeginScope(state);
+        }
+
+        public void ConnectionStart(string connectionId)
+        {
+            _trace1.ConnectionStart(connectionId);
+            _trace2.ConnectionStart(connectionId);
+        }
+
+        public void ConnectionStop(string connectionId)
+        {
+            _trace1.ConnectionStop(connectionId);
+            _trace2.ConnectionStop(connectionId);
+        }
+
+        public void ConnectionPause(string connectionId)
+        {
+            _trace1.ConnectionPause(connectionId);
+            _trace2.ConnectionPause(connectionId);
+        }
+
+        public void ConnectionResume(string connectionId)
+        {
+            _trace1.ConnectionResume(connectionId);
+            _trace2.ConnectionResume(connectionId);
+        }
+
+        public void ConnectionRejected(string connectionId)
+        {
+            _trace1.ConnectionRejected(connectionId);
+            _trace2.ConnectionRejected(connectionId);
+        }
+
+        public void ConnectionKeepAlive(string connectionId)
+        {
+            _trace1.ConnectionKeepAlive(connectionId);
+            _trace2.ConnectionKeepAlive(connectionId);
+        }
+
+        public void ConnectionDisconnect(string connectionId)
+        {
+            _trace1.ConnectionDisconnect(connectionId);
+            _trace2.ConnectionDisconnect(connectionId);
+        }
+
+        public void RequestProcessingError(string connectionId, Exception ex)
+        {
+            _trace1.RequestProcessingError(connectionId, ex);
+            _trace2.RequestProcessingError(connectionId, ex);
+        }
+
+        public void ConnectionHeadResponseBodyWrite(string connectionId, long count)
+        {
+            _trace1.ConnectionHeadResponseBodyWrite(connectionId, count);
+            _trace2.ConnectionHeadResponseBodyWrite(connectionId, count);
+        }
+
+        public void NotAllConnectionsClosedGracefully()
+        {
+            _trace1.NotAllConnectionsClosedGracefully();
+            _trace2.NotAllConnectionsClosedGracefully();
+        }
+
+        public void ConnectionBadRequest(string connectionId, BadHttpRequestException ex)
+        {
+            _trace1.ConnectionBadRequest(connectionId, ex);
+            _trace2.ConnectionBadRequest(connectionId, ex);
+        }
+
+        public void ApplicationError(string connectionId, string traceIdentifier, Exception ex)
+        {
+            _trace1.ApplicationError(connectionId, traceIdentifier, ex);
+            _trace2.ApplicationError(connectionId, traceIdentifier, ex);
+        }
+
+        public void NotAllConnectionsAborted()
+        {
+            _trace1.NotAllConnectionsAborted();
+            _trace2.NotAllConnectionsAborted();
+        }
+
+        public void HeartbeatSlow(TimeSpan interval, DateTimeOffset now)
+        {
+            _trace1.HeartbeatSlow(interval, now);
+            _trace2.HeartbeatSlow(interval, now);
+        }
+
+        public void ApplicationNeverCompleted(string connectionId)
+        {
+            _trace1.ApplicationNeverCompleted(connectionId);
+            _trace2.ApplicationNeverCompleted(connectionId);
+        }
+
+        public void RequestBodyStart(string connectionId, string traceIdentifier)
+        {
+            _trace1.RequestBodyStart(connectionId, traceIdentifier);
+            _trace2.RequestBodyStart(connectionId, traceIdentifier);
+        }
+
+        public void RequestBodyDone(string connectionId, string traceIdentifier)
+        {
+            _trace1.RequestBodyDone(connectionId, traceIdentifier);
+            _trace2.RequestBodyDone(connectionId, traceIdentifier);
+        }
+
+        public void RequestBodyNotEntirelyRead(string connectionId, string traceIdentifier)
+        {
+            _trace1.RequestBodyNotEntirelyRead(connectionId, traceIdentifier);
+            _trace2.RequestBodyNotEntirelyRead(connectionId, traceIdentifier);
+        }
+
+        public void RequestBodyDrainTimedOut(string connectionId, string traceIdentifier)
+        {
+            _trace1.RequestBodyDrainTimedOut(connectionId, traceIdentifier);
+            _trace2.RequestBodyDrainTimedOut(connectionId, traceIdentifier);
+        }
+
+        public void RequestBodyMininumDataRateNotSatisfied(string connectionId, string traceIdentifier, double rate)
+        {
+            _trace1.RequestBodyMininumDataRateNotSatisfied(connectionId, traceIdentifier, rate);
+            _trace2.RequestBodyMininumDataRateNotSatisfied(connectionId, traceIdentifier, rate);
+        }
+
+        public void ResponseMininumDataRateNotSatisfied(string connectionId, string traceIdentifier)
+        {
+            _trace1.ResponseMininumDataRateNotSatisfied(connectionId, traceIdentifier);
+            _trace2.ResponseMininumDataRateNotSatisfied(connectionId, traceIdentifier);
+        }
+
+        public void Http2ConnectionError(string connectionId, Http2ConnectionErrorException ex)
+        {
+            _trace1.Http2ConnectionError(connectionId, ex);
+            _trace2.Http2ConnectionError(connectionId, ex);
+        }
+
+        public void Http2StreamError(string connectionId, Http2StreamErrorException ex)
+        {
+            _trace1.Http2StreamError(connectionId, ex);
+            _trace2.Http2StreamError(connectionId, ex);
+        }
+
+        public void HPackDecodingError(string connectionId, int streamId, HPackDecodingException ex)
+        {
+            _trace1.HPackDecodingError(connectionId, streamId, ex);
+            _trace2.HPackDecodingError(connectionId, streamId, ex);
+        }
+    }
+}

test/shared/TestApplicationErrorLogger.cs

@@ -44,15 +44,16 @@ namespace Microsoft.AspNetCore.Testing
 
         public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception exception, Func<TState, Exception, string> formatter)
         {
+            var exceptionIsIgnored = IgnoredExceptions.Contains(exception?.GetType());
 #if true
-            if (logLevel == LogLevel.Critical && ThrowOnCriticalErrors && !IgnoredExceptions.Contains(exception.GetType()))
+            if (logLevel == LogLevel.Critical && ThrowOnCriticalErrors && !exceptionIsIgnored)
 #endif
             {
                 var log = $"Log {logLevel}[{eventId}]: {formatter(state, exception)} {exception}";
 
                 Console.WriteLine(log);
 
-                if (logLevel == LogLevel.Critical && ThrowOnCriticalErrors && !IgnoredExceptions.Contains(exception.GetType()))
+                if (logLevel == LogLevel.Critical && ThrowOnCriticalErrors && !exceptionIsIgnored)
                 {
                     throw new Exception($"Unexpected critical error. {log}", exception);
                 }

test/shared/TestServiceContext.cs

@@ -1,10 +1,13 @@
 // 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.IO.Pipelines;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2.HPack;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.Extensions.Logging;
 
@@ -22,13 +25,18 @@ namespace Microsoft.AspNetCore.Testing
         }
 
         public TestServiceContext(ILoggerFactory loggerFactory)
-            : this(loggerFactory, new KestrelTrace(loggerFactory.CreateLogger("Microsoft.AspNetCore.Server.Kestrel")))
         {
+            Initialize(loggerFactory, CreateLoggingTrace(loggerFactory));
         }
 
         public TestServiceContext(ILoggerFactory loggerFactory, IKestrelTrace kestrelTrace)
         {
-            Initialize(loggerFactory, kestrelTrace);
+            Initialize(loggerFactory, new CompositeKestrelTrace(kestrelTrace, CreateLoggingTrace(loggerFactory)));
+        }
+
+        private static KestrelTrace CreateLoggingTrace(ILoggerFactory loggerFactory)
+        {
+            return new KestrelTrace(loggerFactory.CreateLogger("Microsoft.AspNetCore.Server.Kestrel"));
         }
 
         private void Initialize(ILoggerFactory loggerFactory, IKestrelTrace kestrelTrace)

test/xunit.runner.json

@@ -2,5 +2,6 @@
   "$schema": "http://json.schemastore.org/xunit.runner.schema",
   "appDomain": "denied",
   "methodDisplay": "method",
-  "longRunningTestSeconds": 60
+  "longRunningTestSeconds": 60,
+  "maxParallelThreads": -1
 }