Start each request on fresh ExecutionContext (#14146)

src/Servers/Kestrel/Core/src/Internal/Http/HttpProtocol.cs

@@ -547,6 +547,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         {
             try
             {
+                // We run the request processing loop in a seperate async method so per connection
+                // exception handling doesn't complicate the generated asm for the loop.
                 await ProcessRequests(application);
             }
             catch (BadHttpRequestException ex)
@@ -624,6 +626,26 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
 
                 InitializeBodyControl(messageBody);
 
+                // We run user controlled request processing in a seperate async method
+                // so any changes made to ExecutionContext are undone when it returns and
+                // each request starts with a fresh ExecutionContext state.
+                await ProcessRequest(application);
+
+                // Even for non-keep-alive requests, try to consume the entire body to avoid RSTs.
+                if (!_connectionAborted && _requestRejectedException == null && !messageBody.IsEmpty)
+                {
+                    await messageBody.ConsumeAsync();
+                }
+
+                if (HasStartedConsumingRequestBody)
+                {
+                    await messageBody.StopAsync();
+                }
+            }
+        }
+
+        private async ValueTask ProcessRequest<TContext>(IHttpApplication<TContext> application)
+        {
             var context = application.CreateContext(this);
 
             try
@@ -709,18 +731,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
             }
 
             application.DisposeContext(context, _applicationException);
-
-                // Even for non-keep-alive requests, try to consume the entire body to avoid RSTs.
-                if (!_connectionAborted && _requestRejectedException == null && !messageBody.IsEmpty)
-                {
-                    await messageBody.ConsumeAsync();
-                }
-
-                if (HasStartedConsumingRequestBody)
-                {
-                    await messageBody.StopAsync();
-                }
-            }
         }
 
         public void OnStarting(Func<object, Task> callback, object state)
@@ -749,100 +759,46 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
         protected Task FireOnStarting()
         {
             var onStarting = _onStarting;
-
+            if (onStarting?.Count > 0)
-            if (onStarting == null || onStarting.Count == 0)
             {
-                return Task.CompletedTask;
+                return ProcessEvents(this, onStarting);
-            }
-            else
-            {
-                return FireOnStartingMayAwait(onStarting);
-            }
-        }
-
-        private Task FireOnStartingMayAwait(Stack<KeyValuePair<Func<object, Task>, object>> onStarting)
-        {
-            try
-            {
-                while (onStarting.TryPop(out var entry))
-                {
-                    var task = entry.Key.Invoke(entry.Value);
-                    if (!task.IsCompletedSuccessfully)
-                    {
-                        return FireOnStartingAwaited(task, onStarting);
-                    }
-                }
-            }
-            catch (Exception ex)
-            {
-                ReportApplicationError(ex);
             }
 
             return Task.CompletedTask;
-        }
 
-        private async Task FireOnStartingAwaited(Task currentTask, Stack<KeyValuePair<Func<object, Task>, object>> onStarting)
+            static async Task ProcessEvents(HttpProtocol protocol, Stack<KeyValuePair<Func<object, Task>, object>> events)
             {
+                // Try/Catch is outside the loop as any error that occurs is before the request starts.
+                // So we want to report it as an ApplicationError to fail the request and not process more events.
                 try
                 {
-                await currentTask;
+                    while (events.TryPop(out var entry))
-
-                while (onStarting.TryPop(out var entry))
                     {
                         await entry.Key.Invoke(entry.Value);
                     }
                 }
                 catch (Exception ex)
                 {
-                ReportApplicationError(ex);
+                    protocol.ReportApplicationError(ex);
+                }
             }
         }
 
         protected Task FireOnCompleted()
         {
             var onCompleted = _onCompleted;
-
+            if (onCompleted?.Count > 0)
-            if (onCompleted == null || onCompleted.Count == 0)
             {
-                return Task.CompletedTask;
+                return ProcessEvents(this, onCompleted);
-            }
-
-            return FireOnCompletedMayAwait(onCompleted);
-        }
-
-        private Task FireOnCompletedMayAwait(Stack<KeyValuePair<Func<object, Task>, object>> onCompleted)
-        {
-            while (onCompleted.TryPop(out var entry))
-            {
-                try
-                {
-                    var task = entry.Key.Invoke(entry.Value);
-                    if (!task.IsCompletedSuccessfully)
-                    {
-                        return FireOnCompletedAwaited(task, onCompleted);
-                    }
-                }
-                catch (Exception ex)
-                {
-                    ReportApplicationError(ex);
-                }
             }
 
             return Task.CompletedTask;
-        }
 
-        private async Task FireOnCompletedAwaited(Task currentTask, Stack<KeyValuePair<Func<object, Task>, object>> onCompleted)
+            static async Task ProcessEvents(HttpProtocol protocol, Stack<KeyValuePair<Func<object, Task>, object>> events)
             {
-            try
+                // Try/Catch is inside the loop as any error that occurs is after the request has finished.
-            {
+                // So we will just log it and keep processing the events, as the completion has already happened.
-                await currentTask;
+                while (events.TryPop(out var entry))
-            }
-            catch (Exception ex)
-            {
-                Log.ApplicationError(ConnectionId, TraceIdentifier, ex);
-            }
-
-            while (onCompleted.TryPop(out var entry))
                 {
                     try
                     {
@@ -850,7 +806,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http
                     }
                     catch (Exception ex)
                     {
-                    Log.ApplicationError(ConnectionId, TraceIdentifier, ex);
+                        protocol.Log.ApplicationError(protocol.ConnectionId, protocol.TraceIdentifier, ex);
+                    }
                 }
             }
         }

src/Servers/Kestrel/test/InMemory.FunctionalTests/RequestTests.cs

@@ -285,6 +285,264 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
             }
         }
 
+        [Fact]
+        public async Task ExecutionContextMutationsOfValueTypeDoNotLeakAcrossRequestsOnSameConnection()
+        {
+            var local = new AsyncLocal<int>();
+
+            // It's important this method isn't async as that will revert the ExecutionContext
+            Task ExecuteApplication(HttpContext context)
+            {
+                var value = local.Value;
+                Assert.Equal(0, value);
+
+                context.Response.OnStarting(() =>
+                {
+                    local.Value++;
+                    return Task.CompletedTask;
+                });
+
+                context.Response.OnCompleted(() =>
+                {
+                    local.Value++;
+                    return Task.CompletedTask;
+                });
+
+                local.Value++;
+                context.Response.ContentLength = 1;
+                return context.Response.WriteAsync($"{value}");
+            }
+
+            var testContext = new TestServiceContext(LoggerFactory);
+
+            await using var server = new TestServer(ExecuteApplication, testContext);
+            await TestAsyncLocalValues(testContext, server);
+        }
+
+        [Fact]
+        public async Task ExecutionContextMutationsOfReferenceTypeDoNotLeakAcrossRequestsOnSameConnection()
+        {
+            var local = new AsyncLocal<IntAsClass>();
+
+            // It's important this method isn't async as that will revert the ExecutionContext
+            Task ExecuteApplication(HttpContext context)
+            {
+                Assert.Null(local.Value);
+                local.Value = new IntAsClass();
+
+                var value = local.Value.Value;
+                Assert.Equal(0, value);
+
+                context.Response.OnStarting(() =>
+                {
+                    local.Value.Value++;
+                    return Task.CompletedTask;
+                });
+
+                context.Response.OnCompleted(() =>
+                {
+                    local.Value.Value++;
+                    return Task.CompletedTask;
+                });
+
+                local.Value.Value++;
+                context.Response.ContentLength = 1;
+                return context.Response.WriteAsync($"{value}");
+            }
+
+            var testContext = new TestServiceContext(LoggerFactory);
+
+            await using var server = new TestServer(ExecuteApplication, testContext);
+            await TestAsyncLocalValues(testContext, server);
+        }
+
+#pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously
+        [Fact]
+        public async Task ExecutionContextMutationsDoNotLeakAcrossAwaits()
+        {
+            var local = new AsyncLocal<int>();
+
+            // It's important this method isn't async as that will revert the ExecutionContext
+            Task ExecuteApplication(HttpContext context)
+            {
+                var value = local.Value;
+                Assert.Equal(0, value);
+
+                context.Response.OnStarting(async () =>
+                {
+                    local.Value++;
+                    Assert.Equal(1, local.Value);
+                });
+
+                context.Response.OnCompleted(async () =>
+                {
+                    local.Value++;
+                    Assert.Equal(1, local.Value);
+                });
+
+                context.Response.ContentLength = 1;
+                return context.Response.WriteAsync($"{value}");
+            }
+
+            var testContext = new TestServiceContext(LoggerFactory);
+
+            await using var server = new TestServer(ExecuteApplication, testContext);
+            await TestAsyncLocalValues(testContext, server);
+        }
+
+        [Fact]
+        public async Task ExecutionContextMutationsOfValueTypeFlowIntoButNotOutOfAsyncEvents()
+        {
+            var local = new AsyncLocal<int>();
+
+            async Task ExecuteApplication(HttpContext context)
+            {
+                var value = local.Value;
+                Assert.Equal(0, value);
+
+                context.Response.OnStarting(async () =>
+                {
+                    local.Value++;
+                    Assert.Equal(2, local.Value);
+                });
+
+                context.Response.OnCompleted(async () =>
+                {
+                    local.Value++;
+                    Assert.Equal(2, local.Value);
+                });
+
+                local.Value++;
+                Assert.Equal(1, local.Value);
+
+                context.Response.ContentLength = 1;
+                await context.Response.WriteAsync($"{value}");
+
+                local.Value++;
+                Assert.Equal(2, local.Value);
+            }
+
+            var testContext = new TestServiceContext(LoggerFactory);
+
+            await using var server = new TestServer(ExecuteApplication, testContext);
+            await TestAsyncLocalValues(testContext, server);
+        }
+
+        [Fact]
+        public async Task ExecutionContextMutationsOfReferenceTypeFlowThroughAsyncEvents()
+        {
+            var local = new AsyncLocal<IntAsClass>();
+
+            async Task ExecuteApplication(HttpContext context)
+            {
+                Assert.Null(local.Value);
+                local.Value = new IntAsClass();
+
+                var value = local.Value.Value;
+                Assert.Equal(0, value); // Start
+
+                context.Response.OnStarting(async () =>
+                {
+                    local.Value.Value++;
+                    Assert.Equal(2, local.Value.Value); // Second
+                });
+
+                context.Response.OnCompleted(async () =>
+                {
+                    local.Value.Value++;
+                    Assert.Equal(4, local.Value.Value); // Fourth
+                });
+
+                local.Value.Value++;
+                Assert.Equal(1, local.Value.Value); // First
+
+                context.Response.ContentLength = 1;
+                await context.Response.WriteAsync($"{value}");
+
+                local.Value.Value++;
+                Assert.Equal(3, local.Value.Value); // Third
+            }
+
+            var testContext = new TestServiceContext(LoggerFactory);
+
+            await using var server = new TestServer(ExecuteApplication, testContext);
+            await TestAsyncLocalValues(testContext, server);
+        }
+#pragma warning restore CS1998 // Async method lacks 'await' operators and will run synchronously
+
+        [Fact]
+        public async Task ExecutionContextMutationsOfValueTypeFlowIntoButNotOutOfNonAsyncEvents()
+        {
+            var local = new AsyncLocal<int>();
+
+            async Task ExecuteApplication(HttpContext context)
+            {
+                var value = local.Value;
+                Assert.Equal(0, value);
+
+                context.Response.OnStarting(() =>
+                {
+                    local.Value++;
+                    Assert.Equal(2, local.Value);
+
+                    return Task.CompletedTask;
+                });
+
+                context.Response.OnCompleted(() =>
+                {
+                    local.Value++;
+                    Assert.Equal(2, local.Value);
+
+                    return Task.CompletedTask;
+                });
+
+                local.Value++;
+                Assert.Equal(1, local.Value);
+
+                context.Response.ContentLength = 1;
+                await context.Response.WriteAsync($"{value}");
+
+                local.Value++;
+                Assert.Equal(2, local.Value);
+            }
+
+            var testContext = new TestServiceContext(LoggerFactory);
+
+            await using var server = new TestServer(ExecuteApplication, testContext);
+            await TestAsyncLocalValues(testContext, server);
+        }
+
+        private static async Task TestAsyncLocalValues(TestServiceContext testContext, TestServer server)
+        {
+            using var connection = server.CreateConnection();
+
+            await connection.Send(
+                "GET / HTTP/1.1",
+                "Host:",
+                "",
+                "");
+
+            await connection.Receive(
+                "HTTP/1.1 200 OK",
+                $"Date: {testContext.DateHeaderValue}",
+                "Content-Length: 1",
+                "",
+                "0");
+
+            await connection.Send(
+                "GET / HTTP/1.1",
+                "Host:",
+                "",
+                "");
+
+            await connection.Receive(
+                "HTTP/1.1 200 OK",
+                $"Date: {testContext.DateHeaderValue}",
+                "Content-Length: 1",
+                "",
+                "0");
+        }
+
         [Fact]
         public async Task AppCanSetTraceIdentifier()
         {
@@ -1803,5 +2061,10 @@ namespace Microsoft.AspNetCore.Server.Kestrel.InMemory.FunctionalTests
         }
 
         public static TheoryData<string, string> HostHeaderData => HttpParsingData.HostHeaderData;
+
+        private class IntAsClass
+        {
+            public int Value;
+        }
     }
 }