Handle IIS OnCompleted callbacks later #17268 (#17756)

Author: Tratcher

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

@@ -460,7 +460,7 @@ public unsafe void SetResponseTrailers()
             }
         }
 
-        public abstract Task<bool> ProcessRequestAsync();
+        public abstract Task ProcessRequestAsync();
 
         public void OnStarting(Func<object, Task> callback, object state)
         {
@@ -650,30 +650,19 @@ public void Execute()
 
         private async Task HandleRequest()
         {
-            bool successfulRequest = false;
             try
             {
-                successfulRequest = await ProcessRequestAsync();
+                await ProcessRequestAsync();
             }
             catch (Exception ex)
             {
                 _logger.LogError(0, ex, $"Unexpected exception in {nameof(IISHttpContext)}.{nameof(HandleRequest)}.");
             }
             finally
             {
-                // Post completion after completing the request to resume the state machine
-                PostCompletion(ConvertRequestCompletionResults(successfulRequest));
-
-
                 // Dispose the context
                 Dispose();
             }
         }
-
-        private static NativeMethods.REQUEST_NOTIFICATION_STATUS ConvertRequestCompletionResults(bool success)
-        {
-            return success ? NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_CONTINUE
-                           : NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_FINISH_REQUEST;
-        }
     }
 }

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

@@ -24,32 +24,33 @@ public IISHttpContextOfT(MemoryPool<byte> memoryPool, IHttpApplication<TContext>
             _application = application;
         }
 
-        public override async Task<bool> ProcessRequestAsync()
+        public override async Task ProcessRequestAsync()
         {
-            InitializeContext();
-
             var context = default(TContext);
             var success = true;
 
             try
             {
-                context = _application.CreateContext(this);
+                InitializeContext();
+
+                try
+                {
+                    context = _application.CreateContext(this);
+
+                    await _application.ProcessRequestAsync(context);
+                }
+                catch (BadHttpRequestException ex)
+                {
+                    SetBadRequestState(ex);
+                    ReportApplicationError(ex);
+                    success = false;
+                }
+                catch (Exception ex)
+                {
+                    ReportApplicationError(ex);
+                    success = false;
+                }
 
-                await _application.ProcessRequestAsync(context);
-            }
-            catch (BadHttpRequestException ex)
-            {
-                SetBadRequestState(ex);
-                ReportApplicationError(ex);
-                success = false;
-            }
-            catch (Exception ex)
-            {
-                ReportApplicationError(ex);
-                success = false;
-            }
-            finally
-            {
                 await CompleteResponseBodyAsync();
                 _streams.Stop();
 
@@ -66,36 +67,18 @@ public override async Task<bool> ProcessRequestAsync()
                     SetResetCode(2);
                 }
 
-                if (_onCompleted != null)
+                if (!_requestAborted)
                 {
-                    await FireOnCompleted();
+                    await ProduceEnd();
+                }
+                else if (!HasResponseStarted && _requestRejectedException == null)
+                {
+                    // If the request was aborted and no response was sent, there's no
+                    // meaningful status code to log.
+                    StatusCode = 0;
+                    success = false;
                 }
-            }
-
-            if (!_requestAborted)
-            {
-                await ProduceEnd();
-            }
-            else if (!HasResponseStarted && _requestRejectedException == null)
-            {
-                // If the request was aborted and no response was sent, there's no
-                // meaningful status code to log.
-                StatusCode = 0;
-                success = false;
-            }
 
-            try
-            {
-                _application.DisposeContext(context, _applicationException);
-            }
-            catch (Exception ex)
-            {
-                // TODO Log this
-                _applicationException = _applicationException ?? ex;
-                success = false;
-            }
-            finally
-            {
                 // Complete response writer and request reader pipe sides
                 _bodyOutput.Dispose();
                 _bodyInputPipe?.Reader.Complete();
@@ -114,7 +97,36 @@ public override async Task<bool> ProcessRequestAsync()
                     await _readBodyTask;
                 }
             }
-            return success;
+            catch (Exception ex)
+            {
+                success = false;
+                ReportApplicationError(ex);
+            }
+            finally
+            {
+                // We're done with anything that touches the request or response, unblock the client.
+                PostCompletion(ConvertRequestCompletionResults(success));
+
+                if (_onCompleted != null)
+                {
+                    await FireOnCompleted();
+                }
+
+                try
+                {
+                    _application.DisposeContext(context, _applicationException);
+                }
+                catch (Exception ex)
+                {
+                    ReportApplicationError(ex);
+                }
+            }
+        }
+
+        private static NativeMethods.REQUEST_NOTIFICATION_STATUS ConvertRequestCompletionResults(bool success)
+        {
+            return success ? NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_CONTINUE
+                           : NativeMethods.REQUEST_NOTIFICATION_STATUS.RQ_NOTIFICATION_FINISH_REQUEST;
         }
     }
 }

src/Servers/IIS/IIS/test/Common.FunctionalTests/Inprocess/ResponseBodyTests.cs

@@ -30,5 +30,12 @@ public async Task ResponseBodyTest_FlushedPipeAndThenUnflushedPipe_AutoFlushed()
         {
             Assert.Equal(20, (await _fixture.Client.GetByteArrayAsync($"/FlushedPipeAndThenUnflushedPipe")).Length);
         }
+
+        [ConditionalFact]
+        [RequiresNewHandler]
+        public async Task ResponseBodyTest_BodyCompletionNotBlockedByOnCompleted()
+        {
+            Assert.Equal("SlowOnCompleted", await _fixture.Client.GetStringAsync($"/SlowOnCompleted"));
+        }
     }
 }

src/Servers/IIS/IIS/test/testassets/InProcessWebSite/Startup.cs

@@ -1339,6 +1339,13 @@ public async Task Reset_DuringRequestBody_Resets_Complete(HttpContext httpContex
             await _resetDuringRequestBodyResetsCts.Task;
         }
 
+        public async Task SlowOnCompleted(HttpContext context)
+        {
+            // This shouldn't block the response or the server from shutting down.
+            context.Response.OnCompleted(() => Task.Delay(TimeSpan.FromMinutes(5)));
+            await context.Response.WriteAsync("SlowOnCompleted");
+        }
+
         internal static readonly HashSet<(string, StringValues, StringValues)> NullTrailers = new HashSet<(string, StringValues, StringValues)>()
         {
             ("NullString", (string)null, (string)null),