Skip to content

Commit 102dd03

Browse files
jkotalikjkotalik
authored
Wait to dispose CTS for IIS (#9389)
1 parent e247770 commit 102dd03

File tree

6 files changed

+96
-53
lines changed

6 files changed

+96
-53
lines changed

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

Lines changed: 17 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
// Copyright (c) .NET Foundation. All rights reserved.
22
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
33

4-
using System;
54
using System.Threading;
65
using Microsoft.AspNetCore.Connections;
76
using Microsoft.AspNetCore.Http.Features;
@@ -12,6 +11,8 @@ internal partial class IISHttpContext : IHttpRequestLifetimeFeature
1211
{
1312
private CancellationTokenSource _abortedCts;
1413
private CancellationToken? _manuallySetRequestAbortToken;
14+
private object _abortLock = new object();
15+
protected volatile bool _requestAborted;
1516

1617
CancellationToken IHttpRequestLifetimeFeature.RequestAborted
1718
{
@@ -22,16 +23,21 @@ CancellationToken IHttpRequestLifetimeFeature.RequestAborted
2223
{
2324
return _manuallySetRequestAbortToken.Value;
2425
}
25-
// Otherwise, get the abort CTS. If we have one, which would mean that someone previously
26-
// asked for the RequestAborted token, simply return its token. If we don't,
27-
// check to see whether we've already aborted, in which case just return an
28-
// already canceled token. Finally, force a source into existence if we still
29-
// don't have one, and return its token.
30-
var cts = _abortedCts;
31-
return
32-
cts != null ? cts.Token :
33-
(_requestAborted == 1) ? new CancellationToken(true) :
34-
RequestAbortedSource.Token;
26+
27+
lock (_abortLock)
28+
{
29+
if (_requestAborted)
30+
{
31+
return new CancellationToken(true);
32+
}
33+
34+
if (_abortedCts == null)
35+
{
36+
_abortedCts = new CancellationTokenSource();
37+
}
38+
39+
return _abortedCts.Token;
40+
}
3541
}
3642
set
3743
{
@@ -41,28 +47,6 @@ CancellationToken IHttpRequestLifetimeFeature.RequestAborted
4147
}
4248
}
4349

44-
private CancellationTokenSource RequestAbortedSource
45-
{
46-
get
47-
{
48-
// Get the abort token, lazily-initializing it if necessary.
49-
// Make sure it's canceled if an abort request already came in.
50-
51-
// EnsureInitialized can return null since _abortedCts is reset to null
52-
// after it's already been initialized to a non-null value.
53-
// If EnsureInitialized does return null, this property was accessed between
54-
// requests so it's safe to return an ephemeral CancellationTokenSource.
55-
var cts = LazyInitializer.EnsureInitialized(ref _abortedCts, () => new CancellationTokenSource())
56-
?? new CancellationTokenSource();
57-
58-
if (_requestAborted == 1)
59-
{
60-
cts.Cancel();
61-
}
62-
return cts;
63-
}
64-
}
65-
6650
void IHttpRequestLifetimeFeature.Abort()
6751
{
6852
Abort(new ConnectionAbortedException(CoreStrings.ConnectionAbortedByApplication));

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

Lines changed: 36 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,11 @@
33

44
using System;
55
using System.Buffers;
6+
using System.Net.Http;
67
using System.Threading;
78
using System.Threading.Tasks;
89
using Microsoft.AspNetCore.Connections;
10+
using Microsoft.AspNetCore.Http;
911
using Microsoft.AspNetCore.Http.Features;
1012

1113
namespace Microsoft.AspNetCore.Server.IIS.Core
@@ -140,7 +142,6 @@ private async Task WriteBody(bool flush = false)
140142
var result = await _bodyOutput.Reader.ReadAsync();
141143

142144
var buffer = result.Buffer;
143-
144145
try
145146
{
146147
if (!buffer.IsEmpty)
@@ -186,9 +187,17 @@ private async Task WriteBody(bool flush = false)
186187

187188
internal void AbortIO(bool clientDisconnect)
188189
{
189-
if (Interlocked.CompareExchange(ref _requestAborted, 1, 0) != 0)
190+
var shouldScheduleCancellation = false;
191+
192+
lock (_abortLock)
190193
{
191-
return;
194+
if (_requestAborted)
195+
{
196+
return;
197+
}
198+
199+
shouldScheduleCancellation = _abortedCts != null;
200+
_requestAborted = true;
192201
}
193202

194203
if (clientDisconnect)
@@ -198,21 +207,39 @@ internal void AbortIO(bool clientDisconnect)
198207

199208
_bodyOutput.Dispose();
200209

201-
var cts = _abortedCts;
202-
if (cts != null)
210+
if (shouldScheduleCancellation)
203211
{
204-
ThreadPool.QueueUserWorkItem(t =>
212+
// Potentially calling user code. CancelRequestAbortedToken logs any exceptions.
213+
CancelRequestAbortedToken();
214+
}
215+
}
216+
217+
private void CancelRequestAbortedToken()
218+
{
219+
ThreadPool.UnsafeQueueUserWorkItem(ctx =>
205220
{
206221
try
207222
{
208-
cts.Cancel();
223+
CancellationTokenSource localAbortCts = null;
224+
225+
lock (ctx._abortLock)
226+
{
227+
if (ctx._abortedCts != null)
228+
{
229+
localAbortCts = ctx._abortedCts;
230+
ctx._abortedCts = null;
231+
}
232+
}
233+
234+
// If we cancel the cts, we don't dispose as people may still be using
235+
// the cts. It also isn't necessary to dispose a canceled cts.
236+
localAbortCts?.Cancel();
209237
}
210238
catch (Exception ex)
211239
{
212240
Log.ApplicationError(_logger, ((IHttpConnectionFeature)this).ConnectionId, TraceIdentifier, ex);
213241
}
214-
});
215-
}
242+
}, this, preferLocal: false);
216243
}
217244

218245
public void Abort(Exception reason)

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
using System.IO;
99
using System.IO.Pipelines;
1010
using System.Net;
11+
using System.Net.Http;
1112
using System.Runtime.InteropServices;
1213
using System.Security.Claims;
1314
using System.Security.Principal;
@@ -59,7 +60,6 @@ internal abstract partial class IISHttpContext : NativeRequestContext, IThreadPo
5960
protected Task _writeBodyTask;
6061

6162
private bool _wasUpgraded;
62-
protected int _requestAborted;
6363

6464
protected Pipe _bodyInputPipe;
6565
protected OutputProducer _bodyOutput;
@@ -68,7 +68,6 @@ internal abstract partial class IISHttpContext : NativeRequestContext, IThreadPo
6868
private const string NegotiateString = "Negotiate";
6969
private const string BasicString = "Basic";
7070

71-
7271
internal unsafe IISHttpContext(
7372
MemoryPool<byte> memoryPool,
7473
IntPtr pInProcessHandler,
@@ -509,7 +508,16 @@ protected virtual void Dispose(bool disposing)
509508
wi.Dispose();
510509
}
511510

512-
_abortedCts?.Dispose();
511+
// Lock to prevent CancelRequestAbortedToken from attempting to cancel a disposed CTS.
512+
CancellationTokenSource localAbortCts = null;
513+
514+
lock (_abortLock)
515+
{
516+
localAbortCts = _abortedCts;
517+
_abortedCts = null;
518+
}
519+
520+
localAbortCts?.Dispose();
513521

514522
disposedValue = true;
515523
}

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

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,12 +31,8 @@ public override async Task<bool> ProcessRequestAsync()
3131
try
3232
{
3333
context = _application.CreateContext(this);
34+
3435
await _application.ProcessRequestAsync(context);
35-
// TODO Verification of Response
36-
//if (Volatile.Read(ref _requestAborted) == 0)
37-
//{
38-
// VerifyResponseContentLength();
39-
//}
4036
}
4137
catch (Exception ex)
4238
{
@@ -59,7 +55,7 @@ public override async Task<bool> ProcessRequestAsync()
5955
}
6056
}
6157

62-
if (Volatile.Read(ref _requestAborted) == 0)
58+
if (!_requestAborted)
6359
{
6460
await ProduceEnd();
6561
}

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,12 @@
55
using System.Threading;
66
using System.Threading.Tasks;
77
using Microsoft.AspNetCore.Connections;
8+
using Microsoft.AspNetCore.Http;
89
using Microsoft.AspNetCore.Server.IntegrationTesting;
910
using Microsoft.AspNetCore.Testing;
1011
using Microsoft.AspNetCore.Testing.xunit;
1112
using Microsoft.Extensions.Logging;
13+
using Microsoft.Extensions.Logging.Testing;
1214
using Xunit;
1315

1416
namespace Microsoft.AspNetCore.Server.IISIntegration.FunctionalTests
@@ -52,8 +54,8 @@ public async Task WritesSucceedAfterClientDisconnect()
5254
}
5355

5456
[ConditionalFact]
55-
[Flaky("https://github.com/aspnet/AspNetCore-Internal/issues/1831", FlakyOn.All)]
56-
public async Task WritesCancelledWhenUsingAbortedToken()
57+
[Repeat(20)]
58+
public async Task WritesCanceledWhenUsingAbortedToken()
5759
{
5860
var requestStartedCompletionSource = CreateTaskCompletionSource();
5961
var requestCompletedCompletionSource = CreateTaskCompletionSource();
@@ -69,6 +71,7 @@ public async Task WritesCancelledWhenUsingAbortedToken()
6971
while (true)
7072
{
7173
await ctx.Response.Body.WriteAsync(data, ctx.RequestAborted);
74+
await Task.Delay(10); // Small delay to not constantly call WriteAsync.
7275
}
7376
}
7477
catch (Exception e)

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

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,31 @@ public async Task RequestAbortedIsTrippedAfterAbort()
138138
Assert.True(tokenAborted);
139139
}
140140

141+
[ConditionalFact]
142+
public async Task CancellationTokenIsUsableAfterAbortingRequest()
143+
{
144+
using (var testServer = await TestServer.Create(async ctx =>
145+
{
146+
var token = ctx.RequestAborted;
147+
var originalRegistration = token.Register(() => { });
148+
149+
ctx.Abort();
150+
151+
Assert.True(token.WaitHandle.WaitOne(10000));
152+
Assert.True(ctx.RequestAborted.WaitHandle.WaitOne(10000));
153+
Assert.Equal(token, originalRegistration.Token);
154+
155+
await Task.CompletedTask;
156+
}, LoggerFactory))
157+
{
158+
using (var connection = testServer.CreateConnection())
159+
{
160+
await SendContentLength1Post(connection);
161+
await connection.WaitForConnectionClose();
162+
}
163+
}
164+
}
165+
141166
private static async Task SendContentLength1Post(TestConnection connection)
142167
{
143168
await connection.Send(

0 commit comments

Comments
 (0)