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) { }
+        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) { }
-        [System.Diagnostics.DebuggerStepThroughAttribute]
         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,20 +46,55 @@ namespace Microsoft.AspNetCore.Diagnostics
             }
         }
 
-        public async Task Invoke(HttpContext context)
+        public Task Invoke(HttpContext context)
         {
+            ExceptionDispatchInfo edi;
             try
             {
-                await _next(context);
+                var task = _next(context);
-            }
+                if (!task.IsCompletedSuccessfully)
-            catch (Exception ex)
                 {
-                _logger.UnhandledException(ex);
+                    return Awaited(this, context, task);
+                }
+
+                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
+                {
+                    await task;
+                }
+                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);
+                }
+
+                if (edi != null)
+                {
+                    await middleware.HandleException(context, edi);
+                }
+            }
+        }
+
+        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();
-                    throw;
+                edi.Throw();
             }
 
             PathString originalPath = context.Request.Path;
@@ -71,7 +107,7 @@ namespace Microsoft.AspNetCore.Diagnostics
                 context.Response.Clear();
                 var exceptionHandlerFeature = new ExceptionHandlerFeature()
                 {
-                        Error = ex,
+                    Error = edi.SourceException,
                     Path = originalPath.Value,
                 };
                 context.Features.Set<IExceptionHandlerFeature>(exceptionHandlerFeature);
@@ -81,9 +117,9 @@ namespace Microsoft.AspNetCore.Diagnostics
 
                 await _options.ExceptionHandler(context);
 
-                    if (_diagnosticSource.IsEnabled("Microsoft.AspNetCore.Diagnostics.HandledException"))
+                if (_diagnosticListener.IsEnabled() && _diagnosticListener.IsEnabled("Microsoft.AspNetCore.Diagnostics.HandledException"))
                 {
-                        _diagnosticSource.Write("Microsoft.AspNetCore.Diagnostics.HandledException", new { httpContext = context, exception = ex });
+                    _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.
@@ -98,11 +134,11 @@ namespace Microsoft.AspNetCore.Diagnostics
             {
                 context.Request.Path = originalPath;
             }
-                throw; // Re-throw the original if we couldn't handle it
+
-            }
+            edi.Throw(); // Re-throw the original if we couldn't handle it
         }
 
-        private Task ClearCacheHeaders(object state)
+        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