Wait to dispose RequestAborted CTS

src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs

@@ -35,7 +35,7 @@ internal abstract partial class HttpProtocol : IDefaultHttpContextContainer, IHt
         private Stack<KeyValuePair<Func<object, Task>, object>> _onCompleted;
 
         private object _abortLock = new object();
-        private volatile bool _requestAborted;
+        private volatile bool _connectionAborted;
         private bool _preventRequestAbortedCancellation;
         private CancellationTokenSource _abortedCts;
         private CancellationToken? _manuallySetRequestAbortToken;
@@ -257,7 +257,7 @@ public CancellationToken RequestAborted
                         return new CancellationToken(false);
                     }
 
-                    if (_requestAborted)
+                    if (_connectionAborted)
                     {
                         return new CancellationToken(true);
                     }
@@ -375,18 +375,21 @@ public void Reset()
             Scheme = _scheme;
 
             _manuallySetRequestAbortToken = null;
-            _preventRequestAbortedCancellation = false;
 
-            // Lock to prevent CancelRequestAbortedToken from attempting to cancel an disposed CTS.
+            // Lock to prevent CancelRequestAbortedToken from attempting to cancel a disposed CTS.
+            CancellationTokenSource localAbortCts = null;
             lock (_abortLock)
             {
-                if (!_requestAborted)
+                _preventRequestAbortedCancellation = false;
+                if (_abortedCts != null)
                 {
-                    _abortedCts?.Dispose();
+                    localAbortCts = _abortedCts;
                     _abortedCts = null;
                 }
             }
 
+            localAbortCts?.Dispose();
+
             if (_wrapperObjectsToDispose != null)
             {
                 foreach (var disposable in _wrapperObjectsToDispose)
@@ -442,9 +445,37 @@ private void CancelRequestAbortedToken()
         {
             try
             {
-                _abortedCts.Cancel();
-                _abortedCts.Dispose();
-                _abortedCts = null;
+                CancellationTokenSource localAbortCts = null;
+                var shouldDispose = false;
+                var shouldCancel = false;
+                lock (_abortLock)
+                {
+                    if (_abortedCts != null)
+                    {
+                        localAbortCts = _abortedCts;
+                        _abortedCts = null;
+
+                        if (_preventRequestAbortedCancellation)
+                        {
+                            shouldDispose = true;
+                        }
+                        else
+                        {
+                            shouldCancel = true;
+                        }
+                    }
+                }
+
+                // If we cancel the cts, we don't dispose as people may still be using
+                // the cts. It also isn't necessary to dispose a canceled cts.
+                if (shouldDispose)
+                {
+                    localAbortCts.Dispose();
+                }
+                else if (shouldCancel)
+                {
+                    localAbortCts.Cancel();
+                }
             }
             catch (Exception ex)
             {
@@ -454,17 +485,19 @@ private void CancelRequestAbortedToken()
 
         protected void AbortRequest()
         {
+            var shouldScheduleCancellation = false;
             lock (_abortLock)
             {
-                if (_requestAborted)
+                if (_connectionAborted)
                 {
                     return;
                 }
 
-                _requestAborted = true;
+                shouldScheduleCancellation = _abortedCts != null && !_preventRequestAbortedCancellation;
+                _connectionAborted = true;
             }
 
-            if (_abortedCts != null)
+            if (shouldScheduleCancellation)
             {
                 // Potentially calling user code. CancelRequestAbortedToken logs any exceptions.
                 ServiceContext.Scheduler.Schedule(state => ((HttpProtocol)state).CancelRequestAbortedToken(), this);
@@ -481,14 +514,12 @@ private void PreventRequestAbortedCancellation()
         {
             lock (_abortLock)
             {
-                if (_requestAborted)
+                if (_connectionAborted)
                 {
                     return;
                 }
 
                 _preventRequestAbortedCancellation = true;
-                _abortedCts?.Dispose();
-                _abortedCts = null;
             }
         }
 
@@ -594,7 +625,7 @@ private async Task ProcessRequests<TContext>(IHttpApplication<TContext> applicat
                     // Run the application code for this request
                     await application.ProcessRequestAsync(context);
 
-                    if (!_requestAborted)
+                    if (!_connectionAborted)
                     {
                         VerifyResponseContentLength();
                     }
@@ -630,7 +661,7 @@ private async Task ProcessRequests<TContext>(IHttpApplication<TContext> applicat
                 // 4XX responses are written by TryProduceInvalidRequestResponse during connection tear down.
                 if (_requestRejectedException == null)
                 {
-                    if (!_requestAborted)
+                    if (!_connectionAborted)
                     {
                         // Call ProduceEnd() before consuming the rest of the request body to prevent
                         // delaying clients waiting for the chunk terminator:
@@ -662,7 +693,7 @@ private async Task ProcessRequests<TContext>(IHttpApplication<TContext> applicat
                 application.DisposeContext(context, _applicationException);
 
                 // Even for non-keep-alive requests, try to consume the entire body to avoid RSTs.
-                if (!_requestAborted && _requestRejectedException == null && !messageBody.IsEmpty)
+                if (!_connectionAborted && _requestRejectedException == null && !messageBody.IsEmpty)
                 {
                     await messageBody.ConsumeAsync();
                 }
@@ -964,7 +995,7 @@ private void VerifyInitializeState(int firstWriteByteCount)
         protected Task TryProduceInvalidRequestResponse()
         {
             // If _requestAborted is set, the connection has already been closed.
-            if (_requestRejectedException != null && !_requestAborted)
+            if (_requestRejectedException != null && !_connectionAborted)
             {
                 return ProduceEnd();
             }

src/Servers/Kestrel/Core/test/Http1ConnectionTests.cs

@@ -729,6 +729,20 @@ public async Task RequestAbortedTokenIsResetBeforeLastWriteWithChunkedEncoding()
             Assert.False(_http1Connection.RequestAborted.IsCancellationRequested);
         }
 
+        [Fact]
+        public void RequestAbortedTokenIsFullyUsableAfterCancellation()
+        {
+            var originalToken = _http1Connection.RequestAborted;
+            var originalRegistration = originalToken.Register(() => { });
+
+            _http1Connection.Abort(new ConnectionAbortedException());
+
+            Assert.True(originalToken.WaitHandle.WaitOne(TestConstants.DefaultTimeout));
+            Assert.True(_http1Connection.RequestAborted.WaitHandle.WaitOne(TestConstants.DefaultTimeout));
+
+            Assert.Equal(originalToken, originalRegistration.Token);
+        }
+
         [Fact]
         public void RequestAbortedTokenIsUsableAfterCancellation()
         {

src/Servers/Kestrel/test/InMemory.FunctionalTests/HttpsTests.cs

@@ -217,6 +217,7 @@ await sslStream.AuthenticateAsClientAsync("127.0.0.1", clientCertificates: null,
         }
 
         [Fact]
+        [Repeat(20)]
         public async Task DoesNotThrowObjectDisposedExceptionFromWriteAsyncAfterConnectionIsAborted()
         {
             var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);