Wait to dispose CTS for IIS

src/Servers/IIS/IIS/src/Core/IISHttpContext.IHttpRequestLifetimeFeature.cs

@@ -1,7 +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;
 using System.Threading;
 using Microsoft.AspNetCore.Connections;
 using Microsoft.AspNetCore.Http.Features;
@@ -12,6 +11,8 @@ internal partial class IISHttpContext : IHttpRequestLifetimeFeature
     {
         private CancellationTokenSource _abortedCts;
         private CancellationToken? _manuallySetRequestAbortToken;
+        private object _abortLock = new object();
+        protected volatile bool _requestAborted;
 
         CancellationToken IHttpRequestLifetimeFeature.RequestAborted
         {
@@ -22,16 +23,21 @@ CancellationToken IHttpRequestLifetimeFeature.RequestAborted
                 {
                     return _manuallySetRequestAbortToken.Value;
                 }
-                // Otherwise, get the abort CTS.  If we have one, which would mean that someone previously
-                // asked for the RequestAborted token, simply return its token.  If we don't,
-                // check to see whether we've already aborted, in which case just return an
-                // already canceled token.  Finally, force a source into existence if we still
-                // don't have one, and return its token.
-                var cts = _abortedCts;
-                return
-                    cts != null ? cts.Token :
-                    (_requestAborted == 1) ? new CancellationToken(true) :
-                    RequestAbortedSource.Token;
+
+                lock (_abortLock)
+                {
+                    if (_requestAborted)
+                    {
+                        return new CancellationToken(true);
+                    }
+
+                    if (_abortedCts == null)
+                    {
+                        _abortedCts = new CancellationTokenSource();
+                    }
+
+                    return _abortedCts.Token;
+                }
             }
             set
             {
@@ -41,28 +47,6 @@ CancellationToken IHttpRequestLifetimeFeature.RequestAborted
             }
         }
 
-        private CancellationTokenSource RequestAbortedSource
-        {
-            get
-            {
-                // Get the abort token, lazily-initializing it if necessary.
-                // Make sure it's canceled if an abort request already came in.
-
-                // EnsureInitialized can return null since _abortedCts is reset to null
-                // after it's already been initialized to a non-null value.
-                // If EnsureInitialized does return null, this property was accessed between
-                // requests so it's safe to return an ephemeral CancellationTokenSource.
-                var cts = LazyInitializer.EnsureInitialized(ref _abortedCts, () => new CancellationTokenSource())
-                            ?? new CancellationTokenSource();
-
-                if (_requestAborted == 1)
-                {
-                    cts.Cancel();
-                }
-                return cts;
-            }
-        }
-
         void IHttpRequestLifetimeFeature.Abort()
         {
             Abort(new ConnectionAbortedException(CoreStrings.ConnectionAbortedByApplication));

src/Servers/IIS/IIS/src/Core/IISHttpContext.IO.cs

@@ -3,9 +3,11 @@
 
 using System;
 using System.Buffers;
+using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 
 namespace Microsoft.AspNetCore.Server.IIS.Core
@@ -140,7 +142,6 @@ private async Task WriteBody(bool flush = false)
                     var result = await _bodyOutput.Reader.ReadAsync();
 
                     var buffer = result.Buffer;
-
                     try
                     {
                         if (!buffer.IsEmpty)
@@ -186,9 +187,17 @@ private async Task WriteBody(bool flush = false)
 
         internal void AbortIO(bool clientDisconnect)
         {
-            if (Interlocked.CompareExchange(ref _requestAborted, 1, 0) != 0)
+            var shouldScheduleCancellation = false;
+
+            lock (_abortLock)
             {
-                return;
+                if (_requestAborted)
+                {
+                    return;
+                }
+
+                shouldScheduleCancellation = _abortedCts != null;
+                _requestAborted = true;
             }
 
             if (clientDisconnect)
@@ -198,21 +207,39 @@ internal void AbortIO(bool clientDisconnect)
 
             _bodyOutput.Dispose();
 
-            var cts = _abortedCts;
-            if (cts != null)
+            if (shouldScheduleCancellation)
             {
-                ThreadPool.QueueUserWorkItem(t =>
+                // Potentially calling user code. CancelRequestAbortedToken logs any exceptions.
+                CancelRequestAbortedToken();
+            }
+        }
+
+        private void CancelRequestAbortedToken()
+        {
+            ThreadPool.UnsafeQueueUserWorkItem(ctx =>
                 {
                     try
                     {
-                        cts.Cancel();
+                        CancellationTokenSource localAbortCts = null;
+
+                        lock (ctx._abortLock)
+                        {
+                            if (ctx._abortedCts != null)
+                            {
+                                localAbortCts = ctx._abortedCts;
+                                ctx._abortedCts = null;
+                            }
+                        }
+
+                        // 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.
+                        localAbortCts?.Cancel();
                     }
                     catch (Exception ex)
                     {
                         Log.ApplicationError(_logger, ((IHttpConnectionFeature)this).ConnectionId, TraceIdentifier, ex);
                     }
-                });
-            }
+                }, this, preferLocal: false);
         }
 
         public void Abort(Exception reason)

src/Servers/IIS/IIS/src/Core/IISHttpContext.cs

@@ -8,6 +8,7 @@
 using System.IO;
 using System.IO.Pipelines;
 using System.Net;
+using System.Net.Http;
 using System.Runtime.InteropServices;
 using System.Security.Claims;
 using System.Security.Principal;
@@ -59,7 +60,6 @@ internal abstract partial class IISHttpContext : NativeRequestContext, IThreadPo
         protected Task _writeBodyTask;
 
         private bool _wasUpgraded;
-        protected int _requestAborted;
 
         protected Pipe _bodyInputPipe;
         protected OutputProducer _bodyOutput;
@@ -68,7 +68,6 @@ internal abstract partial class IISHttpContext : NativeRequestContext, IThreadPo
         private const string NegotiateString = "Negotiate";
         private const string BasicString = "Basic";
 
-
         internal unsafe IISHttpContext(
             MemoryPool<byte> memoryPool,
             IntPtr pInProcessHandler,
@@ -509,7 +508,16 @@ protected virtual void Dispose(bool disposing)
                     wi.Dispose();
                 }
 
-                _abortedCts?.Dispose();
+                // Lock to prevent CancelRequestAbortedToken from attempting to cancel a disposed CTS.
+                CancellationTokenSource localAbortCts = null;
+
+                lock (_abortLock)
+                {
+                    localAbortCts = _abortedCts;
+                    _abortedCts = null;
+                }
+
+                localAbortCts?.Dispose();
 
                 disposedValue = true;
             }

src/Servers/IIS/IIS/src/Core/IISHttpContextOfT.cs

@@ -31,12 +31,8 @@ public override async Task<bool> ProcessRequestAsync()
             try
             {
                 context = _application.CreateContext(this);
+
                 await _application.ProcessRequestAsync(context);
-                // TODO Verification of Response
-                //if (Volatile.Read(ref _requestAborted) == 0)
-                //{
-                //    VerifyResponseContentLength();
-                //}
             }
             catch (Exception ex)
             {
@@ -59,7 +55,7 @@ public override async Task<bool> ProcessRequestAsync()
                 }
             }
 
-            if (Volatile.Read(ref _requestAborted) == 0)
+            if (!_requestAborted)
             {
                 await ProduceEnd();
             }

src/Servers/IIS/IIS/test/IIS.Tests/ClientDisconnectTests.cs

@@ -5,10 +5,12 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Server.IntegrationTesting;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.AspNetCore.Testing.xunit;
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Testing;
 using Xunit;
 
 namespace Microsoft.AspNetCore.Server.IISIntegration.FunctionalTests
@@ -52,8 +54,8 @@ public async Task WritesSucceedAfterClientDisconnect()
         }
 
         [ConditionalFact]
-        [Flaky("https://github.com/aspnet/AspNetCore-Internal/issues/1831", FlakyOn.All)]
-        public async Task WritesCancelledWhenUsingAbortedToken()
+        [Repeat(20)]
+        public async Task WritesCanceledWhenUsingAbortedToken()
         {
             var requestStartedCompletionSource = CreateTaskCompletionSource();
             var requestCompletedCompletionSource = CreateTaskCompletionSource();
@@ -69,6 +71,7 @@ public async Task WritesCancelledWhenUsingAbortedToken()
                     while (true)
                     {
                         await ctx.Response.Body.WriteAsync(data, ctx.RequestAborted);
+                        await Task.Delay(10); // Small delay to not constantly call WriteAsync.
                     }
                 }
                 catch (Exception e)

src/Servers/IIS/IIS/test/IIS.Tests/ResponseAbortTests.cs

@@ -138,6 +138,31 @@ public async Task RequestAbortedIsTrippedAfterAbort()
             Assert.True(tokenAborted);
         }
 
+        [ConditionalFact]
+        public async Task CancellationTokenIsUsableAfterAbortingRequest()
+        {
+            using (var testServer = await TestServer.Create(async ctx =>
+            {
+                var token = ctx.RequestAborted;
+                var originalRegistration = token.Register(() => { });
+
+                ctx.Abort();
+
+                Assert.True(token.WaitHandle.WaitOne(10000));
+                Assert.True(ctx.RequestAborted.WaitHandle.WaitOne(10000));
+                Assert.Equal(token, originalRegistration.Token);
+
+                await Task.CompletedTask;
+            }, LoggerFactory))
+            {
+                using (var connection = testServer.CreateConnection())
+                {
+                    await SendContentLength1Post(connection);
+                    await connection.WaitForConnectionClose();
+                }
+            }
+        }
+
         private static async Task SendContentLength1Post(TestConnection connection)
         {
             await connection.Send(