Skip to content

Commit 9177683

Browse files
jkotalikjkotalik
authored
Fix OnCompleted and implement CompleteAsync (#24686)
1 parent 9b9d46d commit 9177683

File tree

12 files changed

+367
-65
lines changed

12 files changed

+367
-65
lines changed

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

Lines changed: 43 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -195,18 +195,56 @@ Task IHttpResponseBodyFeature.StartAsync(CancellationToken cancellationToken)
195195
Task IHttpResponseBodyFeature.SendFileAsync(string path, long offset, long? count, CancellationToken cancellation)
196196
=> SendFileFallback.SendFileAsync(ResponseBody, path, offset, count, cancellation);
197197

198-
Task IHttpResponseBodyFeature.CompleteAsync() => CompleteResponseBodyAsync();
199-
200198
// TODO: In the future this could complete the body all the way down to the server. For now it just ensures
201199
// any unflushed data gets flushed.
202-
protected Task CompleteResponseBodyAsync()
200+
Task IHttpResponseBodyFeature.CompleteAsync()
203201
{
204202
if (ResponsePipeWrapper != null)
205203
{
206-
return ResponsePipeWrapper.CompleteAsync().AsTask();
204+
var completeAsyncValueTask = ResponsePipeWrapper.CompleteAsync();
205+
if (!completeAsyncValueTask.IsCompletedSuccessfully)
206+
{
207+
return CompleteResponseBodyAwaited(completeAsyncValueTask);
208+
}
209+
completeAsyncValueTask.GetAwaiter().GetResult();
207210
}
208211

209-
return Task.CompletedTask;
212+
if (!HasResponseStarted)
213+
{
214+
var initializeTask = InitializeResponse(flushHeaders: false);
215+
if (!initializeTask.IsCompletedSuccessfully)
216+
{
217+
return CompleteInitializeResponseAwaited(initializeTask);
218+
}
219+
}
220+
221+
// Completing the body output will trigger a final flush to IIS.
222+
// We'd rather not bypass the bodyoutput to flush, to guarantee we avoid
223+
// calling flush twice at the same time.
224+
// awaiting the writeBodyTask guarantees the response has finished the final flush.
225+
_bodyOutput.Complete();
226+
return _writeBodyTask;
227+
}
228+
229+
private async Task CompleteResponseBodyAwaited(ValueTask completeAsyncTask)
230+
{
231+
await completeAsyncTask;
232+
233+
if (!HasResponseStarted)
234+
{
235+
await InitializeResponse(flushHeaders: false);
236+
}
237+
238+
_bodyOutput.Complete();
239+
await _writeBodyTask;
240+
}
241+
242+
private async Task CompleteInitializeResponseAwaited(Task initializeTask)
243+
{
244+
await initializeTask;
245+
246+
_bodyOutput.Complete();
247+
await _writeBodyTask;
210248
}
211249

212250
bool IHttpUpgradeFeature.IsUpgradableRequest => true;

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,9 @@ private async Task WriteBody(bool flush = false)
175175
SetResponseTrailers();
176176
}
177177

178+
// Done with response, say there is no more data after writing trailers.
179+
await AsyncIO.FlushAsync(moreData: false);
180+
178181
break;
179182
}
180183

@@ -228,7 +231,7 @@ internal void AbortIO(bool clientDisconnect)
228231
Log.ConnectionDisconnect(_logger, ((IHttpConnectionFeature)this).ConnectionId);
229232
}
230233

231-
_bodyOutput.Dispose();
234+
_bodyOutput.Complete();
232235

233236
if (shouldScheduleCancellation)
234237
{

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -288,7 +288,7 @@ private async Task ProduceStart(bool flushHeaders)
288288

289289
if (!canHaveNonEmptyBody)
290290
{
291-
_bodyOutput.Dispose();
291+
_bodyOutput.Complete();
292292
}
293293
else
294294
{
@@ -664,7 +664,6 @@ private async Task HandleRequest()
664664
// Post completion after completing the request to resume the state machine
665665
PostCompletion(ConvertRequestCompletionResults(successfulRequest));
666666

667-
668667
// Dispose the context
669668
Dispose();
670669
}

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

Lines changed: 56 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -26,31 +26,36 @@ public IISHttpContextOfT(MemoryPool<byte> memoryPool, IHttpApplication<TContext>
2626

2727
public override async Task<bool> ProcessRequestAsync()
2828
{
29-
InitializeContext();
30-
3129
var context = default(TContext);
3230
var success = true;
3331

3432
try
3533
{
36-
context = _application.CreateContext(this);
34+
InitializeContext();
35+
36+
try
37+
{
38+
context = _application.CreateContext(this);
39+
40+
await _application.ProcessRequestAsync(context);
41+
}
42+
catch (BadHttpRequestException ex)
43+
{
44+
SetBadRequestState(ex);
45+
ReportApplicationError(ex);
46+
success = false;
47+
}
48+
catch (Exception ex)
49+
{
50+
ReportApplicationError(ex);
51+
success = false;
52+
}
53+
54+
if (ResponsePipeWrapper != null)
55+
{
56+
await ResponsePipeWrapper.CompleteAsync();
57+
}
3758

38-
await _application.ProcessRequestAsync(context);
39-
}
40-
catch (BadHttpRequestException ex)
41-
{
42-
SetBadRequestState(ex);
43-
ReportApplicationError(ex);
44-
success = false;
45-
}
46-
catch (Exception ex)
47-
{
48-
ReportApplicationError(ex);
49-
success = false;
50-
}
51-
finally
52-
{
53-
await CompleteResponseBodyAsync();
5459
_streams.Stop();
5560

5661
if (!HasResponseStarted && _applicationException == null && _onStarting != null)
@@ -66,38 +71,20 @@ public override async Task<bool> ProcessRequestAsync()
6671
SetResetCode(2);
6772
}
6873

69-
if (_onCompleted != null)
74+
if (!_requestAborted)
7075
{
71-
await FireOnCompleted();
76+
await ProduceEnd();
77+
}
78+
else if (!HasResponseStarted && _requestRejectedException == null)
79+
{
80+
// If the request was aborted and no response was sent, there's no
81+
// meaningful status code to log.
82+
StatusCode = 0;
83+
success = false;
7284
}
73-
}
74-
75-
if (!_requestAborted)
76-
{
77-
await ProduceEnd();
78-
}
79-
else if (!HasResponseStarted && _requestRejectedException == null)
80-
{
81-
// If the request was aborted and no response was sent, there's no
82-
// meaningful status code to log.
83-
StatusCode = 0;
84-
success = false;
85-
}
8685

87-
try
88-
{
89-
_application.DisposeContext(context, _applicationException);
90-
}
91-
catch (Exception ex)
92-
{
93-
// TODO Log this
94-
_applicationException = _applicationException ?? ex;
95-
success = false;
96-
}
97-
finally
98-
{
9986
// Complete response writer and request reader pipe sides
100-
_bodyOutput.Dispose();
87+
_bodyOutput.Complete();
10188
_bodyInputPipe?.Reader.Complete();
10289

10390
// Allow writes to drain
@@ -107,13 +94,34 @@ public override async Task<bool> ProcessRequestAsync()
10794
}
10895

10996
// Cancel all remaining IO, there might be reads pending if not entire request body was sent by client
110-
AsyncIO?.Dispose();
97+
AsyncIO?.Complete();
11198

11299
if (_readBodyTask != null)
113100
{
114101
await _readBodyTask;
115102
}
116103
}
104+
catch (Exception ex)
105+
{
106+
success = false;
107+
ReportApplicationError(ex);
108+
}
109+
finally
110+
{
111+
if (_onCompleted != null)
112+
{
113+
await FireOnCompleted();
114+
}
115+
116+
try
117+
{
118+
_application.DisposeContext(context, _applicationException);
119+
}
120+
catch (Exception ex)
121+
{
122+
ReportApplicationError(ex);
123+
}
124+
}
117125
return success;
118126
}
119127
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,6 @@ private void Run(AsyncIOOperation ioOperation)
8989
}
9090
}
9191

92-
9392
public ValueTask FlushAsync(bool moreData)
9493
{
9594
var flush = GetFlushOperation();
@@ -137,7 +136,7 @@ public void NotifyCompletion(int hr, int bytes)
137136
nextContinuation?.Invoke();
138137
}
139138

140-
public void Dispose()
139+
public void Complete()
141140
{
142141
lock (_contextSync)
143142
{

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,12 @@
77

88
namespace Microsoft.AspNetCore.Server.IIS.Core.IO
99
{
10-
internal interface IAsyncIOEngine: IDisposable
10+
internal interface IAsyncIOEngine
1111
{
1212
ValueTask<int> ReadAsync(Memory<byte> memory);
1313
ValueTask<int> WriteAsync(ReadOnlySequence<byte> data);
1414
ValueTask FlushAsync(bool moreData);
1515
void NotifyCompletion(int hr, int bytes);
16+
void Complete();
1617
}
1718
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ private void ThrowIfNotInitialized()
110110
}
111111
}
112112

113-
public void Dispose()
113+
public void Complete()
114114
{
115115
lock (_contextLock)
116116
{

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ public Task FlushAsync(CancellationToken cancellationToken)
4040
return FlushAsync(_pipe.Writer, cancellationToken);
4141
}
4242

43-
public void Dispose()
43+
public void Complete()
4444
{
4545
lock (_contextLock)
4646
{

‎src/Servers/IIS/IIS/test/Common.FunctionalTests/Inprocess/ResetTests.cs

Lines changed: 73 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,79 @@ public async Task Reset_DuringRequestBody_Resets()
354354
.Build().RunAsync();
355355
}
356356

357+
[ConditionalFact]
358+
[MinimumOSVersion(OperatingSystems.Windows, "10.0.19529", SkipReason = "Reset support was added in Win10_20H2.")]
359+
public async Task Reset_AfterCompleteAsync_NoReset()
360+
{
361+
var deploymentParameters = GetHttpsDeploymentParameters();
362+
var deploymentResult = await DeployAsync(deploymentParameters);
363+
364+
await new HostBuilder()
365+
.UseHttp2Cat(deploymentResult.ApplicationBaseUri, async h2Connection =>
366+
{
367+
await h2Connection.InitializeConnectionAsync();
368+
369+
h2Connection.Logger.LogInformation("Initialized http2 connection. Starting stream 1.");
370+
371+
await h2Connection.StartStreamAsync(1, GetHeaders("/Reset_AfterCompleteAsync_NoReset"), endStream: true);
372+
373+
await h2Connection.ReceiveHeadersAsync(1, decodedHeaders =>
374+
{
375+
Assert.Equal("200", decodedHeaders[HeaderNames.Status]);
376+
});
377+
378+
var dataFrame = await h2Connection.ReceiveFrameAsync();
379+
Http2Utilities.VerifyDataFrame(dataFrame, 1, endOfStream: false, length: 11);
380+
381+
dataFrame = await h2Connection.ReceiveFrameAsync();
382+
Http2Utilities.VerifyDataFrame(dataFrame, 1, endOfStream: true, length: 0);
383+
384+
h2Connection.Logger.LogInformation("Connection stopped.");
385+
})
386+
.Build().RunAsync();
387+
}
388+
389+
[ConditionalFact]
390+
[MinimumOSVersion(OperatingSystems.Windows, "10.0.19529", SkipReason = "Reset support was added in Win10_20H2.")]
391+
public async Task Reset_CompleteAsyncDuringRequestBody_Resets()
392+
{
393+
var deploymentParameters = GetHttpsDeploymentParameters();
394+
var deploymentResult = await DeployAsync(deploymentParameters);
395+
396+
await new HostBuilder()
397+
.UseHttp2Cat(deploymentResult.ApplicationBaseUri, async h2Connection =>
398+
{
399+
await h2Connection.InitializeConnectionAsync();
400+
401+
h2Connection.Logger.LogInformation("Initialized http2 connection. Starting stream 1.");
402+
403+
await h2Connection.StartStreamAsync(1, Http2Utilities.PostRequestHeaders, endStream: false);
404+
await h2Connection.SendDataAsync(1, new byte[10], endStream: false);
405+
406+
await h2Connection.ReceiveHeadersAsync(1, decodedHeaders =>
407+
{
408+
Assert.Equal("200", decodedHeaders[HeaderNames.Status]);
409+
});
410+
411+
var dataFrame = await h2Connection.ReceiveFrameAsync();
412+
if (Environment.OSVersion.Version >= Win10_Regressed_DataFrame)
413+
{
414+
// TODO: Remove when the regression is fixed.
415+
// https://github.com/dotnet/aspnetcore/issues/23164#issuecomment-652646163
416+
Http2Utilities.VerifyDataFrame(dataFrame, 1, endOfStream: false, length: 0);
417+
418+
dataFrame = await h2Connection.ReceiveFrameAsync();
419+
}
420+
Http2Utilities.VerifyDataFrame(dataFrame, 1, endOfStream: true, length: 0);
421+
422+
var resetFrame = await h2Connection.ReceiveFrameAsync();
423+
Http2Utilities.VerifyResetFrame(resetFrame, expectedStreamId: 1, expectedErrorCode: Http2ErrorCode.NO_ERROR);
424+
425+
h2Connection.Logger.LogInformation("Connection stopped.");
426+
})
427+
.Build().RunAsync();
428+
}
429+
357430
private static List<KeyValuePair<string, string>> GetHeaders(string path)
358431
{
359432
var headers = Headers.ToList();
@@ -372,7 +445,6 @@ private static List<KeyValuePair<string, string>> GetPostHeaders(string path)
372445
return headers;
373446
}
374447

375-
376448
private IISDeploymentParameters GetHttpsDeploymentParameters()
377449
{
378450
var port = TestPortHelper.GetNextSSLPort();

0 commit comments

Comments
 (0)