Improve ConnectionLimitMiddleware and connection pipeline building (#2010)

* Improve ConnectionLimitMiddleware and connection pipeline building
* Add IDecrementConcurrentConnectionCountFeature
* Flow connection features from connection middleware

benchmarks/Kestrel.Performance/FrameFeatureCollection.cs

@@ -85,6 +85,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
             var frameContext = new FrameContext
             {
                 ServiceContext = serviceContext,
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = new PipeFactory()
             };
 

benchmarks/Kestrel.Performance/FrameParsingOverheadBenchmark.cs

@@ -3,6 +3,7 @@
 
 using System.IO.Pipelines;
 using BenchmarkDotNet.Attributes;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
@@ -29,6 +30,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
             var frameContext = new FrameContext
             {
                 ServiceContext = serviceContext,
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = new PipeFactory(),
                 TimeoutControl = new MockTimeoutControl()
             };

benchmarks/Kestrel.Performance/FrameWritingBenchmark.cs

@@ -7,6 +7,7 @@ using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using BenchmarkDotNet.Attributes;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
@@ -107,6 +108,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
             var frame = new TestFrame<object>(application: null, context: new FrameContext
             {
                 ServiceContext = serviceContext,
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = pipeFactory,
                 Application = pair.Application,
                 Transport = pair.Transport

benchmarks/Kestrel.Performance/RequestParsingBenchmark.cs

@@ -3,6 +3,7 @@
 
 using System.IO.Pipelines;
 using BenchmarkDotNet.Attributes;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
@@ -33,6 +34,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
             var frameContext = new FrameContext
             {
                 ServiceContext = serviceContext,
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = PipeFactory,
                 TimeoutControl = new MockTimeoutControl()
             };

benchmarks/Kestrel.Performance/ResponseHeaderCollectionBenchmark.cs

@@ -6,6 +6,7 @@ using System.Runtime.CompilerServices;
 using System.Text;
 using BenchmarkDotNet.Attributes;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Http.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
@@ -177,6 +178,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
             var frameContext = new FrameContext
             {
                 ServiceContext = serviceContext,
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = new PipeFactory()
             };
 

benchmarks/Kestrel.Performance/ResponseHeadersWritingBenchmark.cs

@@ -7,6 +7,7 @@ using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using BenchmarkDotNet.Attributes;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
@@ -123,6 +124,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Performance
             var frame = new TestFrame<object>(application: null, context: new FrameContext
             {
                 ServiceContext = serviceContext,
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = pipeFactory,
                 TimeoutControl = new MockTimeoutControl(),
                 Application = pair.Application,

src/Kestrel.Core/Features/IDecrementConcurrentConnectionCountFeature.cs

@@ -0,0 +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.
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Features
+{
+    /// <summary>
+    /// A connection feature allowing middleware to stop counting connections towards <see cref="KestrelServerLimits.MaxConcurrentConnections"/>.
+    /// This is used by Kestrel internally to stop counting upgraded connections towards this limit.
+    /// </summary>
+    public interface IDecrementConcurrentConnectionCountFeature
+    {
+        /// <summary>
+        /// Idempotent method to stop counting a connection towards <see cref="KestrelServerLimits.MaxConcurrentConnections"/>.
+        /// </summary>
+        void ReleaseConnection();
+    }
+}

src/Kestrel.Core/Internal/ConnectionLimitBuilderExtensions.cs

@@ -1,16 +0,0 @@
-using Microsoft.AspNetCore.Protocols;
-
-namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
-{
-    public static class ConnectionLimitBuilderExtensions
-    {
-        public static IConnectionBuilder UseConnectionLimit(this IConnectionBuilder builder, ServiceContext serviceContext)
-        {
-            return builder.Use(next =>
-            {
-                var middleware = new ConnectionLimitMiddleware(next, serviceContext);
-                return middleware.OnConnectionAsync;
-            });
-        }
-    }
-}

src/Kestrel.Core/Internal/ConnectionLimitMiddleware.cs

@@ -1,32 +1,74 @@
-using System.Threading.Tasks;
+// 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.Threading.Tasks;
 using Microsoft.AspNetCore.Protocols;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
 {
     public class ConnectionLimitMiddleware
     {
-        private readonly ServiceContext _serviceContext;
         private readonly ConnectionDelegate _next;
+        private readonly ResourceCounter _concurrentConnectionCounter;
+        private readonly IKestrelTrace _trace;
 
-        public ConnectionLimitMiddleware(ConnectionDelegate next, ServiceContext serviceContext)
+        public ConnectionLimitMiddleware(ConnectionDelegate next, long connectionLimit, IKestrelTrace trace)
+            : this(next, ResourceCounter.Quota(connectionLimit), trace)
         {
-            _next = next;
-            _serviceContext = serviceContext;
         }
 
-        public Task OnConnectionAsync(ConnectionContext connection)
+        // For Testing
+        internal ConnectionLimitMiddleware(ConnectionDelegate next, ResourceCounter concurrentConnectionCounter, IKestrelTrace trace)
         {
-            if (!_serviceContext.ConnectionManager.NormalConnectionCount.TryLockOne())
+            _next = next;
+            _concurrentConnectionCounter = concurrentConnectionCounter;
+            _trace = trace;
+        }
+
+        public async Task OnConnectionAsync(ConnectionContext connection)
+        {
+            if (!_concurrentConnectionCounter.TryLockOne())
             {
                 KestrelEventSource.Log.ConnectionRejected(connection.ConnectionId);
-                _serviceContext.Log.ConnectionRejected(connection.ConnectionId);
+                _trace.ConnectionRejected(connection.ConnectionId);
                 connection.Transport.Input.Complete();
                 connection.Transport.Output.Complete();
-                return Task.CompletedTask;
+                return;
             }
 
-            return _next(connection);
+            var releasor = new ConnectionReleasor(_concurrentConnectionCounter);
+
+            try
+            {
+                connection.Features.Set<IDecrementConcurrentConnectionCountFeature>(releasor);
+                await _next(connection);
+            }
+            finally
+            {
+                releasor.ReleaseConnection();
+            }
+        }
+
+        private class ConnectionReleasor : IDecrementConcurrentConnectionCountFeature
+        {
+            private readonly ResourceCounter _concurrentConnectionCounter;
+            private bool _connectionReleased;
+
+            public ConnectionReleasor(ResourceCounter normalConnectionCounter)
+            {
+                _concurrentConnectionCounter = normalConnectionCounter;
+            }
+
+            public void ReleaseConnection()
+            {
+                if (!_connectionReleased)
+                {
+                    _connectionReleased = true;
+                    _concurrentConnectionCounter.ReleaseOne();
+                }
+            }
         }
     }
 }

src/Kestrel.Core/Internal/FrameConnection.cs

@@ -141,7 +141,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                     adaptedPipelineTask = adaptedPipeline.RunAsync(stream);
                 }
 
-                if (_frame.ConnectionFeatures?.Get<ITlsApplicationProtocolFeature>()?.ApplicationProtocol == "h2" &&
+                if (_frame.ConnectionFeatures.Get<ITlsApplicationProtocolFeature>()?.ApplicationProtocol == "h2" &&
                     Interlocked.CompareExchange(ref _http2ConnectionState, Http2ConnectionStarted, Http2ConnectionNotStarted) == Http2ConnectionNotStarted)
                 {
                     await _http2Connection.ProcessAsync(httpApplication);
@@ -167,10 +167,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                 {
                     _context.ServiceContext.ConnectionManager.UpgradedConnectionCount.ReleaseOne();
                 }
-                else
-                {
-                    _context.ServiceContext.ConnectionManager.NormalConnectionCount.ReleaseOne();
-                }
 
                 KestrelEventSource.Log.ConnectionStop(this);
             }
@@ -181,6 +177,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
             _frame = new Frame<TContext>(httpApplication, new FrameContext
             {
                 ConnectionId = _context.ConnectionId,
+                ConnectionFeatures = _context.ConnectionFeatures,
                 PipeFactory = PipeFactory,
                 LocalEndPoint = LocalEndPoint,
                 RemoteEndPoint = RemoteEndPoint,
@@ -255,10 +252,9 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
         {
             Debug.Assert(_frame != null, $"{nameof(_frame)} is null");
 
-            var features = new FeatureCollection();
             var connectionAdapters = _context.ConnectionAdapters;
             var stream = new RawStream(_context.Transport.Input, _context.Transport.Output);
-            var adapterContext = new ConnectionAdapterContext(features, stream);
+            var adapterContext = new ConnectionAdapterContext(_frame.ConnectionFeatures, stream);
             _adaptedConnections = new List<IAdaptedConnection>(connectionAdapters.Count);
 
             try
@@ -267,7 +263,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                 {
                     var adaptedConnection = await connectionAdapters[i].OnConnectionAsync(adapterContext);
                     _adaptedConnections.Add(adaptedConnection);
-                    adapterContext = new ConnectionAdapterContext(features, adaptedConnection.ConnectionStream);
+                    adapterContext = new ConnectionAdapterContext(_frame.ConnectionFeatures, adaptedConnection.ConnectionStream);
                 }
             }
             catch (Exception ex)
@@ -276,10 +272,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
 
                 return null;
             }
-            finally
-            {
-                _frame.ConnectionFeatures = features;
-            }
 
             return adapterContext.ConnectionStream;
         }

src/Kestrel.Core/Internal/FrameConnectionContext.cs

@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.IO.Pipelines;
 using System.Net;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Adapter.Internal;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
@@ -13,6 +14,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
         public string ConnectionId { get; set; }
         public long FrameConnectionId { get; set; }
         public ServiceContext ServiceContext { get; set; }
+        public IFeatureCollection ConnectionFeatures { get; set; }
         public IList<IConnectionAdapter> ConnectionAdapters { get; set; }
         public PipeFactory PipeFactory { get; set; }
         public IPEndPoint LocalEndPoint { get; set; }

src/Kestrel.Core/Internal/Http/Frame.FeatureCollection.cs

@@ -251,13 +251,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
         object IFeatureCollection.this[Type key]
         {
-            get => FastFeatureGet(key) ?? ConnectionFeatures?[key];
+            get => FastFeatureGet(key) ?? ConnectionFeatures[key];
             set => FastFeatureSet(key, value);
         }
 
         TFeature IFeatureCollection.Get<TFeature>()
         {
-            return (TFeature)(FastFeatureGet(typeof(TFeature)) ?? ConnectionFeatures?[typeof(TFeature)]);
+            return (TFeature)(FastFeatureGet(typeof(TFeature)) ?? ConnectionFeatures[typeof(TFeature)]);
         }
 
         void IFeatureCollection.Set<TFeature>(TFeature instance)
@@ -294,7 +294,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
             _wasUpgraded = true;
 
-            ServiceContext.ConnectionManager.NormalConnectionCount.ReleaseOne();
+            ConnectionFeatures.Get<IDecrementConcurrentConnectionCountFeature>()?.ReleaseConnection();
 
             StatusCode = StatusCodes.Status101SwitchingProtocols;
             ReasonPhrase = "Switching Protocols";

src/Kestrel.Core/Internal/Http/Frame.cs

@@ -107,7 +107,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         private IPEndPoint LocalEndPoint => _frameContext.LocalEndPoint;
         private IPEndPoint RemoteEndPoint => _frameContext.RemoteEndPoint;
 
-        public IFeatureCollection ConnectionFeatures { get; set; }
+        public IFeatureCollection ConnectionFeatures => _frameContext.ConnectionFeatures;
         public IPipeReader Input => _frameContext.Transport.Input;
         public OutputProducer Output { get; }
         public ITimeoutControl TimeoutControl => _frameContext.TimeoutControl;

src/Kestrel.Core/Internal/Http/FrameContext.cs

@@ -3,6 +3,7 @@
 
 using System.IO.Pipelines;
 using System.Net;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Protocols;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 
@@ -12,6 +13,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
     {
         public string ConnectionId { get; set; }
         public ServiceContext ServiceContext { get; set; }
+        public IFeatureCollection ConnectionFeatures { get; set; }
         public PipeFactory PipeFactory { get; set; }
         public IPEndPoint RemoteEndPoint { get; set; }
         public IPEndPoint LocalEndPoint { get; set; }

src/Kestrel.Core/Internal/HttpConnectionMiddleware.cs

@@ -43,6 +43,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal
                 ConnectionId = connectionContext.ConnectionId,
                 FrameConnectionId = frameConnectionId,
                 ServiceContext = _serviceContext,
+                ConnectionFeatures = connectionContext.Features,
                 PipeFactory = connectionContext.PipeFactory,
                 ConnectionAdapters = _connectionAdapters,
                 Transport = connectionContext.Transport,

src/Kestrel.Core/Internal/Infrastructure/FrameConnectionManager.cs

@@ -11,23 +11,17 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
         private readonly ConcurrentDictionary<long, FrameConnectionReference> _connectionReferences = new ConcurrentDictionary<long, FrameConnectionReference>();
         private readonly IKestrelTrace _trace;
 
-        public FrameConnectionManager(IKestrelTrace trace, long? normalConnectionLimit, long? upgradedConnectionLimit)
+        public FrameConnectionManager(IKestrelTrace trace, long? upgradedConnectionLimit)
-            : this(trace, GetCounter(normalConnectionLimit), GetCounter(upgradedConnectionLimit))
+            : this(trace, GetCounter(upgradedConnectionLimit))
         {
         }
 
-        public FrameConnectionManager(IKestrelTrace trace, ResourceCounter normalConnections, ResourceCounter upgradedConnections)
+        public FrameConnectionManager(IKestrelTrace trace, ResourceCounter upgradedConnections)
         {
-            NormalConnectionCount = normalConnections;
             UpgradedConnectionCount = upgradedConnections;
             _trace = trace;
         }
 
-        /// <summary>
-        /// TCP connections processed by Kestrel.
-        /// </summary>
-        public ResourceCounter NormalConnectionCount { get; }
-
         /// <summary>
         /// Connections that have been switched to a different protocol.
         /// </summary>

src/Kestrel.Core/KestrelServer.cs

@@ -70,7 +70,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
             var trace = new KestrelTrace(logger);
             var connectionManager = new FrameConnectionManager(
                 trace,
-                serverOptions.Limits.MaxConcurrentConnections,
                 serverOptions.Limits.MaxConcurrentUpgradedConnections);
 
             var systemClock = new SystemClock();
@@ -135,16 +134,18 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
 
                 async Task OnBind(ListenOptions endpoint)
                 {
-                    // Add the connection limit middleware
-                    endpoint.UseConnectionLimit(ServiceContext);
-
-                    // Configure the user delegate
-                    endpoint.Configure(endpoint);
-
                     // Add the HTTP middleware as the terminal connection middleware
                     endpoint.UseHttpServer(endpoint.ConnectionAdapters, ServiceContext, application);
 
-                    var connectionHandler = new ConnectionHandler(ServiceContext, endpoint.Build()); 
+                    var connectionDelegate = endpoint.Build();
+
+                    // Add the connection limit middleware
+                    if (Options.Limits.MaxConcurrentConnections.HasValue)
+                    {
+                        connectionDelegate = new ConnectionLimitMiddleware(connectionDelegate, Options.Limits.MaxConcurrentConnections.Value, Trace).OnConnectionAsync;
+                    }
+
+                    var connectionHandler = new ConnectionHandler(ServiceContext, connectionDelegate);
                     var transport = _transportFactory.Create(endpoint, connectionHandler);
                     _transports.Add(transport);
 

src/Kestrel.Core/KestrelServerLimits.cs

@@ -201,7 +201,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
         }
 
         /// <summary>
-        /// Gets or sets the maximum number of open HTTP/S connections. When set to null, the number of connections is unlimited.
+        /// Gets or sets the maximum number of open connections. When set to null, the number of connections is unlimited.
         /// </summary>
         /// <remarks>
         /// <para>

src/Kestrel.Core/KestrelServerOptions.cs

@@ -100,7 +100,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
                 throw new ArgumentNullException(nameof(configure));
             }
 
-            var listenOptions = new ListenOptions(endPoint) { KestrelServerOptions = this, Configure = configure };
+            var listenOptions = new ListenOptions(endPoint) { KestrelServerOptions = this };
+            configure(listenOptions);
             ListenOptions.Add(listenOptions);
         }
 
@@ -131,7 +132,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
                 throw new ArgumentNullException(nameof(configure));
             }
 
-            var listenOptions = new ListenOptions(socketPath) { KestrelServerOptions = this, Configure = configure };
+            var listenOptions = new ListenOptions(socketPath) { KestrelServerOptions = this };
+            configure(listenOptions);
             ListenOptions.Add(listenOptions);
         }
 
@@ -154,7 +156,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
                 throw new ArgumentNullException(nameof(configure));
             }
 
-            var listenOptions = new ListenOptions(handle) { KestrelServerOptions = this, Configure = configure };
+            var listenOptions = new ListenOptions(handle) { KestrelServerOptions = this };
+            configure(listenOptions);
             ListenOptions.Add(listenOptions);
         }
     }

src/Kestrel.Core/ListenOptions.cs

@@ -131,8 +131,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core
 
         public IServiceProvider ApplicationServices => KestrelServerOptions?.ApplicationServices;
 
-        internal Action<ListenOptions> Configure { get; set; } = _ => { };
-
         /// <summary>
         /// Gets the name of this endpoint to display on command-line when the web server starts.
         /// </summary>

test/Kestrel.Core.Tests/FrameConnectionManagerTests.cs

@@ -18,7 +18,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
         {
             var connectionId = "0";
             var trace = new Mock<IKestrelTrace>();
-            var frameConnectionManager = new FrameConnectionManager(trace.Object, ResourceCounter.Unlimited, ResourceCounter.Unlimited);
+            var frameConnectionManager = new FrameConnectionManager(trace.Object, ResourceCounter.Unlimited);
 
             // Create FrameConnection in inner scope so it doesn't get rooted by the current frame.
             UnrootedConnectionsGetRemovedFromHeartbeatInnerScope(connectionId, frameConnectionManager, trace);

test/Kestrel.Core.Tests/FrameConnectionTests.cs

@@ -2,11 +2,10 @@
 // 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 System.Collections.Generic;
-using System.Linq;
+using System.IO.Pipelines;
 using System.Threading;
-using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Adapter.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
@@ -31,6 +30,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             {
                 ConnectionId = "0123456789",
                 ConnectionAdapters = new List<IConnectionAdapter>(),
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = _pipeFactory,
                 FrameConnectionId = long.MinValue,
                 Application = pair.Application,

test/Kestrel.Core.Tests/FrameResponseHeadersTests.cs

@@ -2,10 +2,11 @@
 // 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 System.Collections.Generic;
 using System.Globalization;
+using System.IO.Pipelines;
 using Microsoft.AspNetCore.Http;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.Extensions.Primitives;
@@ -23,6 +24,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             var frameContext = new FrameContext
             {
                 ServiceContext = new TestServiceContext(),
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = factory,
                 Application = pair.Application,
                 Transport = pair.Transport,

test/Kestrel.Core.Tests/FrameTests.cs

@@ -63,6 +63,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             _frameContext = new FrameContext
             {
                 ServiceContext = _serviceContext,
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = _pipelineFactory,
                 TimeoutControl = _timeoutControl.Object,
                 Application = pair.Application,

test/Kestrel.Core.Tests/KestrelServerOptionsTests.cs

@@ -18,9 +18,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
                 d.NoDelay = false;
             });
 
-            // Execute the callback
-            o1.ListenOptions[1].Configure(o1.ListenOptions[1]);
-
             Assert.True(o1.ListenOptions[0].NoDelay);
             Assert.False(o1.ListenOptions[1].NoDelay);
         }

test/Kestrel.Core.Tests/TestInput.cs

@@ -4,6 +4,7 @@
 using System;
 using System.IO.Pipelines;
 using System.Text;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.AspNetCore.Testing;
@@ -27,6 +28,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             FrameContext = new FrameContext
             {
                 ServiceContext = new TestServiceContext(),
+                ConnectionFeatures = new FeatureCollection(),
                 Application = Application,
                 Transport = Transport,
                 PipeFactory = _pipelineFactory,

test/Kestrel.FunctionalTests/ConnectionLimitTests.cs

@@ -2,12 +2,13 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
-using System.IO;
+using System.Net;
-using System.Net.Sockets;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
+using Microsoft.AspNetCore.Server.Kestrel.Core;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.AspNetCore.Server.Kestrel.Tests;
 using Microsoft.AspNetCore.Testing;
@@ -23,15 +24,15 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             var requestTcs = new TaskCompletionSource<object>();
             var releasedTcs = new TaskCompletionSource<object>();
             var lockedTcs = new TaskCompletionSource<bool>();
-            var (serviceContext, counter) = SetupMaxConnections(max: 1);
+            var counter = new EventRaisingResourceCounter(ResourceCounter.Quota(1));
             counter.OnLock += (s, e) => lockedTcs.TrySetResult(e);
             counter.OnRelease += (s, e) => releasedTcs.TrySetResult(null);
 
-            using (var server = new TestServer(async context =>
+            using (var server = CreateServerWithMaxConnections(async context =>
             {
                 await context.Response.WriteAsync("Hello");
                 await requestTcs.Task;
-            }, serviceContext))
+            }, counter))
             using (var connection = server.CreateConnection())
             {
                 await connection.SendEmptyGetAsKeepAlive(); ;
@@ -46,8 +47,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         [Fact]
         public async Task UpgradedConnectionsCountsAgainstDifferentLimit()
         {
-            var (serviceContext, _) = SetupMaxConnections(max: 1);
+            using (var server = CreateServerWithMaxConnections(async context =>
-            using (var server = new TestServer(async context =>
             {
                 var feature = context.Features.Get<IHttpUpgradeFeature>();
                 if (feature.IsUpgradableRequest)
@@ -59,7 +59,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
                         await Task.Delay(100);
                     }
                 }
-            }, serviceContext))
+            }, max: 1))
             using (var disposables = new DisposableStack<TestConnection>())
             {
                 var upgraded = server.CreateConnection();
@@ -81,7 +81,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
                     {
                         // this may throw IOException, depending on how fast Kestrel closes the socket
                         await rejected.SendEmptyGetAsKeepAlive();
-                    } catch { }
+                    }
+                    catch { }
 
                     // connection should close without sending any data
                     await rejected.WaitForConnectionClose().TimeoutAfter(TimeSpan.FromSeconds(15));
@@ -93,14 +94,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
         public async Task RejectsConnectionsWhenLimitReached()
         {
             const int max = 10;
-            var (serviceContext, _) = SetupMaxConnections(max);
             var requestTcs = new TaskCompletionSource<object>();
 
-            using (var server = new TestServer(async context =>
+            using (var server = CreateServerWithMaxConnections(async context =>
             {
                 await context.Response.WriteAsync("Hello");
                 await requestTcs.Task;
-            }, serviceContext))
+            }, max))
             using (var disposables = new DisposableStack<TestConnection>())
             {
                 for (var i = 0; i < max; i++)
@@ -141,7 +141,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             var openedTcs = new TaskCompletionSource<object>();
             var closedTcs = new TaskCompletionSource<object>();
 
-            var (serviceContext, counter) = SetupMaxConnections(uint.MaxValue);
+            var counter = new EventRaisingResourceCounter(ResourceCounter.Quota(uint.MaxValue));
 
             counter.OnLock += (o, e) =>
             {
@@ -159,7 +159,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
                 }
             };
 
-            using (var server = new TestServer(_ => Task.CompletedTask, serviceContext))
+            using (var server = CreateServerWithMaxConnections(_ => Task.CompletedTask, counter))
             {
                 // open a bunch of connections in parallel
                 Parallel.For(0, count, async i =>
@@ -187,12 +187,24 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             }
         }
 
-        private (TestServiceContext serviceContext, EventRaisingResourceCounter counter) SetupMaxConnections(long max)
+        private TestServer CreateServerWithMaxConnections(RequestDelegate app, long max)
         {
-            var counter = new EventRaisingResourceCounter(ResourceCounter.Quota(max));
             var serviceContext = new TestServiceContext();
-            serviceContext.ConnectionManager = new FrameConnectionManager(serviceContext.Log, counter, ResourceCounter.Unlimited);
+            serviceContext.ServerOptions.Limits.MaxConcurrentConnections = max;
-            return (serviceContext, counter);
+            return new TestServer(app, serviceContext);
+        }
+
+        private TestServer CreateServerWithMaxConnections(RequestDelegate app, ResourceCounter concurrentConnectionCounter)
+        {
+            var serviceContext = new TestServiceContext();
+            var listenOptions = new ListenOptions(new IPEndPoint(IPAddress.Loopback, 0));
+            listenOptions.Use(next =>
+            {
+                var middleware = new ConnectionLimitMiddleware(next, concurrentConnectionCounter, serviceContext.Log);
+                return middleware.OnConnectionAsync;
+            });
+
+            return new TestServer(app, serviceContext, listenOptions);
         }
     }
 }

test/Kestrel.FunctionalTests/TestHelpers/TestServer.cs

@@ -14,7 +14,6 @@ using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Transport.Abstractions.Internal;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.Extensions.DependencyInjection;
-using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 
 namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests

test/Kestrel.FunctionalTests/UpgradeTests.cs

@@ -260,7 +260,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.FunctionalTests
             const int limit = 10;
             var upgradeTcs = new TaskCompletionSource<object>();
             var serviceContext = new TestServiceContext();
-            serviceContext.ConnectionManager = new FrameConnectionManager(serviceContext.Log, ResourceCounter.Unlimited, ResourceCounter.Quota(limit));
+            serviceContext.ConnectionManager = new FrameConnectionManager(serviceContext.Log, ResourceCounter.Quota(limit));
 
             using (var server = new TestServer(async context =>
             {

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

@@ -6,6 +6,7 @@ using System.Collections.Concurrent;
 using System.IO.Pipelines;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
@@ -699,6 +700,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Libuv.Tests
             var frame = new Frame<object>(null, new FrameContext
             {
                 ServiceContext = serviceContext,
+                ConnectionFeatures = new FeatureCollection(),
                 PipeFactory = _pipeFactory,
                 TimeoutControl = Mock.Of<ITimeoutControl>(),
                 Application = pair.Application,

test/shared/TestServiceContext.cs

@@ -38,7 +38,7 @@ namespace Microsoft.AspNetCore.Testing
             ThreadPool = new LoggingThreadPool(Log);
             SystemClock = new MockSystemClock();
             DateHeaderValueManager = new DateHeaderValueManager(SystemClock);
-            ConnectionManager = new FrameConnectionManager(Log, ResourceCounter.Unlimited, ResourceCounter.Unlimited);
+            ConnectionManager = new FrameConnectionManager(Log, ResourceCounter.Unlimited);
             HttpParserFactory = frameAdapter => new HttpParser<FrameAdapter>(frameAdapter.Frame.ServiceContext.Log.IsEnabled(LogLevel.Information));
             ServerOptions = new KestrelServerOptions
             {