Suppress execution context flow when starting host (#11498)

* Suppress execution context flow when starting host

* Update release_notes.md

* Fix thread switch on AsyncFlowControl disposal

* Add Task.Yield

* Fix unit tests

Author: jviau

release_notes.md

@@ -7,4 +7,6 @@
 - Add JitTrace Files for v4.1045
 - Throw exception instead of timing out when worker channel exits before initializing gRPC (#10937)
 - Adding empty remote message check in the SystemLogger (#11473)
-- Fix `webPubSubTrigger`'s for Flex consumption sku (#11489)
+- Fix `webPubSubTrigger`'s for Flex consumption sku (#11489)
+- Suppress execution context flow into script host start (#11498)
+- Add AzureWebJobsStorage health check (#11471)

src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs

@@ -326,8 +326,27 @@ private async Task StartHostAsync(CancellationToken cancellationToken, int attem
         /// before calling this method. Host starts and restarts must be synchronous to prevent orphaned
         /// hosts or an incorrect ActiveHost.
         /// </summary>
-        private async Task UnsynchronizedStartHostAsync(ScriptHostStartupOperation activeOperation, int attemptCount = 0, JobHostStartupMode startupMode = JobHostStartupMode.Normal)
+        private Task UnsynchronizedStartHostAsync(ScriptHostStartupOperation activeOperation, int attemptCount = 0, JobHostStartupMode startupMode = JobHostStartupMode.Normal)
         {
+            // We may be started from a variety of contexts, some of which may carry AsyncLocal state (such as Activity). We do not want any of this
+            // to flow into the host startup logic, so we suppress the flow of the ExecutionContext.
+            Task start;
+            using (System.Threading.ExecutionContext.SuppressFlow())
+            {
+                start = UnsynchronizedStartHostCoreAsync(activeOperation, attemptCount, startupMode);
+            }
+
+            return start;
+        }
+
+        /// <summary>
+        /// Starts the host without taking a lock. Callers must take a lock on _hostStartSemaphore
+        /// before calling this method. Host starts and restarts must be synchronous to prevent orphaned
+        /// hosts or an incorrect ActiveHost.
+        /// </summary>
+        private async Task UnsynchronizedStartHostCoreAsync(ScriptHostStartupOperation activeOperation, int attemptCount, JobHostStartupMode startupMode)
+        {
+            await Task.Yield(); // ensure any async context is properly suppressed.
             await CheckFileSystemAsync();
             if (ShutdownRequested)
             {

test/WebJobs.Script.Tests.Shared/TaskExtensions.cs

@@ -0,0 +1,39 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Diagnostics;
+using System.Threading.Tasks;
+
+namespace Microsoft.WebJobs.Script.Tests
+{
+    public static class TaskExtensions
+    {
+        /// <summary>
+        /// Waits for a task to complete with a test-appropriate timeout. If a debugger is attached, waits indefinitely.
+        /// </summary>
+        /// <param name="task">The task to wait on.</param>
+        /// <returns>A <see cref="Task"/> representing the asynchronous wait.</returns>
+        public static Task TestWaitAsync(this Task task)
+        {
+            return task.WaitAsync(TimeSpan.FromSeconds(10));
+        }
+
+        /// <summary>
+        /// Waits for a task to complete with a test-appropriate timeout. If a debugger is attached, waits indefinitely.
+        /// </summary>
+        /// <param name="task">The task to wait on.</param>
+        /// <param name="timeout">The timeout to use if no debugger is attached.</param>
+        /// <returns>A <see cref="Task"/> representing the asynchronous wait.</returns>
+        public static Task TestWaitAsync(this Task task, TimeSpan timeout)
+        {
+            ArgumentNullException.ThrowIfNull(task);
+            if (Debugger.IsAttached)
+            {
+                return task;
+            }
+
+            return task.WaitAsync(timeout);
+        }
+    }
+}

test/WebJobs.Script.Tests/WebJobsScriptHostServiceTests.cs

@@ -105,6 +105,15 @@ private Mock<IHost> CreateMockHost(SemaphoreSlim disposedSemaphore = null)
         public async Task StartAsync_Succeeds()
         {
             var hostBuilder = new Mock<IScriptHostBuilder>();
+            _host.Setup(h => h.StartAsync(It.IsAny<CancellationToken>()))
+                .Callback(() =>
+                {
+                    if (Activity.Current != null)
+                    {
+                        throw new InvalidOperationException("Activity flowed into host start.");
+                    }
+                })
+                .Returns(Task.CompletedTask);
             hostBuilder.Setup(b => b.BuildHost(It.IsAny<bool>(), It.IsAny<bool>())).Returns(_host.Object);
 
             _webHostLoggerProvider = new TestLoggerProvider();
@@ -124,6 +133,8 @@ public async Task StartAsync_Succeeds()
                 _hostBuiltChangeTokenSource,
                 _hostBuiltChangeTokenSourceResolverOptions);
 
+            using Activity activity = new("TestActivity_ShouldNotFlowIntoHostStart");
+            activity.Start();
             await _hostService.StartAsync(CancellationToken.None);
 
             // add general post startup validations here
@@ -190,12 +201,9 @@ public async Task HostRestart_Specialization_Succeeds()
             await _hostService.StartAsync(CancellationToken.None);
             Assert.True(AreRequiredMetricsGenerated(metricsLogger));
 
-            Thread restartHostThread = new Thread(new ThreadStart(RestartHost));
-            Thread specializeHostThread = new Thread(new ThreadStart(SpecializeHost));
-            restartHostThread.Start();
-            specializeHostThread.Start();
-            restartHostThread.Join();
-            specializeHostThread.Join();
+            Task restartHost = _hostService.RestartHostAsync("test", default);
+            Task specializeHost = SpecializeHostAsync();
+            await Task.WhenAll(restartHost, specializeHost).TestWaitAsync(TimeSpan.FromSeconds(60)); // this is a slow test for some reason.
 
             var logMessages = _webHostLoggerProvider.GetAllLogMessages().Where(m => m.FormattedMessage.Contains("Restarting host."));
             Assert.Equal(2, logMessages.Count());
@@ -204,9 +212,8 @@ public async Task HostRestart_Specialization_Succeeds()
         [Fact]
         public async Task HostRestart_DuringInitializationWithError_Recovers()
         {
-            SemaphoreSlim semaphore = new SemaphoreSlim(1, 1);
-            await semaphore.WaitAsync();
-            bool waitingInHostBuilder = false;
+            TaskCompletionSource startBlock = new();
+            TaskCompletionSource initialStart = new();
             var metricsLogger = new TestMetricsLogger();
 
             // Have the first host start, but pause. We'll issue a restart, then let
@@ -216,8 +223,8 @@ public async Task HostRestart_DuringInitializationWithError_Recovers()
                 .Setup(h => h.StartAsync(It.IsAny<CancellationToken>()))
                 .Returns(async () =>
                 {
-                    waitingInHostBuilder = true;
-                    await semaphore.WaitAsync();
+                    initialStart.SetResult();
+                    await startBlock.Task;
                     throw new InvalidOperationException("Something happened at startup!");
                 });
 
@@ -226,17 +233,14 @@ public async Task HostRestart_DuringInitializationWithError_Recovers()
                 .Setup(h => h.StartAsync(It.IsAny<CancellationToken>()))
                 .Returns(() => Task.CompletedTask);
 
-            var hostBuilder = new Mock<IScriptHostBuilder>();
-            hostBuilder.SetupSequence(b => b.BuildHost(It.IsAny<bool>(), It.IsAny<bool>()))
-                .Returns(hostA.Object)
-                .Returns(hostB.Object);
+            OrderedScriptHostBuilder hostBuilder = new(hostA.Object, hostB.Object);
 
             _webHostLoggerProvider = new TestLoggerProvider();
             _loggerFactory = new LoggerFactory();
             _loggerFactory.AddProvider(_webHostLoggerProvider);
 
             _hostService = new WebJobsScriptHostService(
-                _monitor, hostBuilder.Object, _loggerFactory,
+                _monitor, hostBuilder, _loggerFactory,
                 _mockScriptWebHostEnvironment.Object, _mockEnvironment.Object,
                 _hostPerformanceManager, _healthMonitorOptions,
                 metricsLogger, new Mock<IApplicationLifetime>().Object,
@@ -245,20 +249,17 @@ public async Task HostRestart_DuringInitializationWithError_Recovers()
             TestLoggerProvider hostALogger = hostA.Object.GetTestLoggerProvider();
             TestLoggerProvider hostBLogger = hostB.Object.GetTestLoggerProvider();
 
-            Task initialStart = _hostService.StartAsync(CancellationToken.None);
-
-            Thread restartHostThread = new Thread(new ThreadStart(RestartHost));
-            restartHostThread.Start();
-            await TestHelpers.Await(() => waitingInHostBuilder);
+            Task firstStart = _hostService.StartAsync(CancellationToken.None);
+            await initialStart.Task.TestWaitAsync();
+            Task restartTask = _hostService.RestartHostAsync("test", default);
 
             // Now let the first host continue. It will throw, which should be correctly handled.
-            semaphore.Release();
+            startBlock.SetResult();
 
-            await TestHelpers.Await(() => _hostService.State == ScriptHostState.Running);
-
-            restartHostThread.Join();
+            await restartTask;
+            Assert.Equal(ScriptHostState.Running, _hostService.State);
 
-            await initialStart;
+            await firstStart;
             Assert.True(AreRequiredMetricsGenerated(metricsLogger));
 
             // Make sure the error was logged to the correct logger
@@ -280,61 +281,48 @@ public async Task HostRestart_DuringInitializationWithError_Recovers()
         [Fact]
         public async Task HostRestart_DuringInitialization_Cancels()
         {
-            SemaphoreSlim semaphore = new SemaphoreSlim(1, 1);
-            await semaphore.WaitAsync();
-            bool waitingInHostBuilder = false;
+            TaskCompletionSource startBlock = new();
+            TaskCompletionSource initialStart = new();
             var metricsLogger = new TestMetricsLogger();
 
             // Have the first host start, but pause. We'll issue a restart, then
-            // let this first host continue and see the cancelation token has fired.
+            // let this first host continue and see the cancellation token has fired.
             var hostA = CreateMockHost();
             hostA
                 .Setup(h => h.StartAsync(It.IsAny<CancellationToken>()))
                 .Returns<CancellationToken>(async ct =>
                 {
-                    waitingInHostBuilder = true;
-                    await semaphore.WaitAsync();
-                    ct.ThrowIfCancellationRequested();
+                    initialStart.SetResult();
+                    await startBlock.Task.WaitAsync(ct);
                 });
 
             var hostB = CreateMockHost();
             hostB
                 .Setup(h => h.StartAsync(It.IsAny<CancellationToken>()))
                 .Returns(() => Task.CompletedTask);
 
-            var hostBuilder = new Mock<IScriptHostBuilder>();
-            hostBuilder.SetupSequence(b => b.BuildHost(It.IsAny<bool>(), It.IsAny<bool>()))
-                .Returns(hostA.Object)
-                .Returns(hostB.Object);
+            OrderedScriptHostBuilder hostBuilder = new(hostA.Object, hostB.Object);
 
             _webHostLoggerProvider = new TestLoggerProvider();
             _loggerFactory = new LoggerFactory();
             _loggerFactory.AddProvider(_webHostLoggerProvider);
 
             _hostService = new WebJobsScriptHostService(
-                _monitor, hostBuilder.Object, _loggerFactory,
+                _monitor, hostBuilder, _loggerFactory,
                 _mockScriptWebHostEnvironment.Object, _mockEnvironment.Object,
                 _hostPerformanceManager, _healthMonitorOptions,
                 metricsLogger, new Mock<IApplicationLifetime>().Object,
                 _mockConfig, new TestScriptEventManager(), _hostMetrics, _functionsHostingConfigOptions, _hostBuiltChangeTokenSource, _hostBuiltChangeTokenSourceResolverOptions);
 
             TestLoggerProvider hostALogger = hostA.Object.GetTestLoggerProvider();
+            Task firstStart = _hostService.StartAsync(default);
+            await initialStart.Task.TestWaitAsync();
+            Task restartTask = _hostService.RestartHostAsync("test", default);
 
-            Task initialStart = _hostService.StartAsync(CancellationToken.None);
-
-            Thread restartHostThread = new Thread(new ThreadStart(RestartHost));
-            restartHostThread.Start();
-            await TestHelpers.Await(() => waitingInHostBuilder);
-            await TestHelpers.Await(() => _webHostLoggerProvider.GetLog().Contains("Restart requested."));
-
-            // Now let the first host continue. It will see that startup is canceled.
-            semaphore.Release();
-
-            await TestHelpers.Await(() => _hostService.State == ScriptHostState.Running);
-
-            restartHostThread.Join();
+            await restartTask;
+            Assert.Equal(ScriptHostState.Running, _hostService.State);
 
-            await initialStart;
+            await firstStart;
             Assert.True(AreRequiredMetricsGenerated(metricsLogger));
 
             // Make sure the error was logged to the correct logger
@@ -582,12 +570,7 @@ public async Task DependencyTrackingTelemetryModule_Race()
             await listenerTask;
         }
 
-        public void RestartHost()
-        {
-            _hostService.RestartHostAsync("test", CancellationToken.None).Wait();
-        }
-
-        public void SpecializeHost()
+        public Task SpecializeHostAsync()
         {
             var mockScriptWebHostEnvironment = new Mock<IScriptWebHostEnvironment>();
             Mock<IConfigurationRoot> mockConfiguration = new Mock<IConfigurationRoot>();
@@ -600,7 +583,7 @@ public void SpecializeHost()
             var hostNameProvider = new HostNameProvider(mockEnvironment.Object);
             var mockApplicationLifetime = new Mock<IApplicationLifetime>(MockBehavior.Strict);
             var manager = new StandbyManager(_hostService, mockLanguageWorkerChannelManager.Object, mockConfiguration.Object, mockScriptWebHostEnvironment.Object, mockEnvironment.Object, _monitor, testLogger, hostNameProvider, mockApplicationLifetime.Object, new TestMetricsLogger());
-            manager.SpecializeHostAsync().Wait();
+            return manager.SpecializeHostAsync();
         }
 
         private bool AreRequiredMetricsGenerated(TestMetricsLogger testMetricsLogger)
@@ -646,5 +629,16 @@ public IHost BuildHost(bool skipHostStartup, bool skipHostConfigurationParsing)
                 return null;
             }
         }
+
+        // Moq SetupSequence doesn't work with execution context suppression.
+        private class OrderedScriptHostBuilder(params IHost[] hosts) : IScriptHostBuilder
+        {
+            private readonly Queue<IHost> _hosts = new(hosts);
+
+            public IHost BuildHost(bool skipHostStartup, bool skipHostConfigurationParsing)
+            {
+                return _hosts.Dequeue();
+            }
+        }
     }
 }