Lazy create ExceptionHandlerMiddleware statemachine (#9549)

src/Middleware/Diagnostics/ref/Microsoft.AspNetCore.Diagnostics.netcoreapp3.0.cs

@@ -71,8 +71,7 @@ namespace Microsoft.AspNetCore.Diagnostics
     }
     public partial class ExceptionHandlerMiddleware
     {
-        public ExceptionHandlerMiddleware(Microsoft.AspNetCore.Http.RequestDelegate next, Microsoft.Extensions.Logging.ILoggerFactory loggerFactory, Microsoft.Extensions.Options.IOptions<Microsoft.AspNetCore.Builder.ExceptionHandlerOptions> options, System.Diagnostics.DiagnosticSource diagnosticSource) { }
-        [System.Diagnostics.DebuggerStepThroughAttribute]
+        public ExceptionHandlerMiddleware(Microsoft.AspNetCore.Http.RequestDelegate next, Microsoft.Extensions.Logging.ILoggerFactory loggerFactory, Microsoft.Extensions.Options.IOptions<Microsoft.AspNetCore.Builder.ExceptionHandlerOptions> options, System.Diagnostics.DiagnosticListener diagnosticListener) { }
         public System.Threading.Tasks.Task Invoke(Microsoft.AspNetCore.Http.HttpContext context) { throw null; }
     }
     public partial class StatusCodeContext

src/Middleware/Diagnostics/src/ExceptionHandler/ExceptionHandlerMiddleware.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Diagnostics;
+using System.Runtime.ExceptionServices;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Diagnostics.Internal;
@@ -19,19 +20,19 @@ namespace Microsoft.AspNetCore.Diagnostics
         private readonly ExceptionHandlerOptions _options;
         private readonly ILogger _logger;
         private readonly Func<object, Task> _clearCacheHeadersDelegate;
-        private readonly DiagnosticSource _diagnosticSource;
+        private readonly DiagnosticListener _diagnosticListener;
 
         public ExceptionHandlerMiddleware(
             RequestDelegate next,
             ILoggerFactory loggerFactory,
             IOptions<ExceptionHandlerOptions> options,
-            DiagnosticSource diagnosticSource)
+            DiagnosticListener diagnosticListener)
         {
             _next = next;
             _options = options.Value;
             _logger = loggerFactory.CreateLogger<ExceptionHandlerMiddleware>();
             _clearCacheHeadersDelegate = ClearCacheHeaders;
-            _diagnosticSource = diagnosticSource;
+            _diagnosticListener = diagnosticListener;
             if (_options.ExceptionHandler == null)
             {
                 if (_options.ExceptionHandlingPath == null)
@@ -45,64 +46,99 @@ namespace Microsoft.AspNetCore.Diagnostics
             }
         }
 
-        public async Task Invoke(HttpContext context)
+        public Task Invoke(HttpContext context)
         {
+            ExceptionDispatchInfo edi;
             try
             {
-                await _next(context);
-            }
-            catch (Exception ex)
-            {
-                _logger.UnhandledException(ex);
-                // We can't do anything if the response has already started, just abort.
-                if (context.Response.HasStarted)
+                var task = _next(context);
+                if (!task.IsCompletedSuccessfully)
                 {
-                    _logger.ResponseStartedErrorHandler();
-                    throw;
+                    return Awaited(this, context, task);
                 }
 
-                PathString originalPath = context.Request.Path;
-                if (_options.ExceptionHandlingPath.HasValue)
-                {
-                    context.Request.Path = _options.ExceptionHandlingPath;
-                }
+                return Task.CompletedTask;
+            }
+            catch (Exception exception)
+            {
+                // Get the Exception, but don't continue processing in the catch block as its bad for stack usage.
+                edi = ExceptionDispatchInfo.Capture(exception);
+            }
+
+            return HandleException(context, edi);
+
+            static async Task Awaited(ExceptionHandlerMiddleware middleware, HttpContext context, Task task)
+            {
+                ExceptionDispatchInfo edi = null;
                 try
                 {
-                    context.Response.Clear();
-                    var exceptionHandlerFeature = new ExceptionHandlerFeature()
-                    {
-                        Error = ex,
-                        Path = originalPath.Value,
-                    };
-                    context.Features.Set<IExceptionHandlerFeature>(exceptionHandlerFeature);
-                    context.Features.Set<IExceptionHandlerPathFeature>(exceptionHandlerFeature);
-                    context.Response.StatusCode = 500;
-                    context.Response.OnStarting(_clearCacheHeadersDelegate, context.Response);
-
-                    await _options.ExceptionHandler(context);
-
-                    if (_diagnosticSource.IsEnabled("Microsoft.AspNetCore.Diagnostics.HandledException"))
-                    {
-                        _diagnosticSource.Write("Microsoft.AspNetCore.Diagnostics.HandledException", new { httpContext = context, exception = ex });
-                    }
-
-                    // TODO: Optional re-throw? We'll re-throw the original exception by default if the error handler throws.
-                    return;
+                    await task;
                 }
-                catch (Exception ex2)
+                catch (Exception exception)
                 {
-                    // Suppress secondary exceptions, re-throw the original.
-                    _logger.ErrorHandlerException(ex2);
+                    // Get the Exception, but don't continue processing in the catch block as its bad for stack usage.
+                    edi = ExceptionDispatchInfo.Capture(exception);
                 }
-                finally
+
+                if (edi != null)
                 {
-                    context.Request.Path = originalPath;
+                    await middleware.HandleException(context, edi);
                 }
-                throw; // Re-throw the original if we couldn't handle it
             }
         }
 
-        private Task ClearCacheHeaders(object state)
+        private async Task HandleException(HttpContext context, ExceptionDispatchInfo edi)
+        {
+            _logger.UnhandledException(edi.SourceException);
+            // We can't do anything if the response has already started, just abort.
+            if (context.Response.HasStarted)
+            {
+                _logger.ResponseStartedErrorHandler();
+                edi.Throw();
+            }
+
+            PathString originalPath = context.Request.Path;
+            if (_options.ExceptionHandlingPath.HasValue)
+            {
+                context.Request.Path = _options.ExceptionHandlingPath;
+            }
+            try
+            {
+                context.Response.Clear();
+                var exceptionHandlerFeature = new ExceptionHandlerFeature()
+                {
+                    Error = edi.SourceException,
+                    Path = originalPath.Value,
+                };
+                context.Features.Set<IExceptionHandlerFeature>(exceptionHandlerFeature);
+                context.Features.Set<IExceptionHandlerPathFeature>(exceptionHandlerFeature);
+                context.Response.StatusCode = 500;
+                context.Response.OnStarting(_clearCacheHeadersDelegate, context.Response);
+
+                await _options.ExceptionHandler(context);
+
+                if (_diagnosticListener.IsEnabled() && _diagnosticListener.IsEnabled("Microsoft.AspNetCore.Diagnostics.HandledException"))
+                {
+                    _diagnosticListener.Write("Microsoft.AspNetCore.Diagnostics.HandledException", new { httpContext = context, exception = edi.SourceException });
+                }
+
+                // TODO: Optional re-throw? We'll re-throw the original exception by default if the error handler throws.
+                return;
+            }
+            catch (Exception ex2)
+            {
+                // Suppress secondary exceptions, re-throw the original.
+                _logger.ErrorHandlerException(ex2);
+            }
+            finally
+            {
+                context.Request.Path = originalPath;
+            }
+
+            edi.Throw(); // Re-throw the original if we couldn't handle it
+        }
+
+        private static Task ClearCacheHeaders(object state)
         {
             var headers = ((HttpResponse)state).Headers;
             headers[HeaderNames.CacheControl] = "no-cache";

src/Middleware/WebSockets/samples/EchoApp/Startup.cs

@@ -9,6 +9,7 @@ using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 
 namespace EchoApp