handling RestartAsync() being called before StartAsync() in WebJobsScriptHostService

Author: brettsam

src/WebJobs.Script.WebHost/Diagnostics/Extensions/ScriptHostServiceLoggerExtension.cs

@@ -130,6 +130,12 @@ public static class ScriptHostServiceLoggerExtension
                 new EventId(519, nameof(EnteringRestart)),
                 "Restart requested. Cancelling any active host startup.");
 
+        private static readonly Action<ILogger, Exception> _restartBeforeStart =
+            LoggerMessage.Define(
+                LogLevel.Debug,
+                new EventId(520, nameof(RestartBeforeStart)),
+                "RestartAsync was called before StartAsync. Delaying restart until StartAsync has been called.");
+
         public static void ScriptHostServiceInitCanceledByRuntime(this ILogger logger)
         {
             _scriptHostServiceInitCanceledByRuntime(logger, null);
@@ -229,5 +235,10 @@ public static void EnteringRestart(this ILogger logger)
         {
             _enteringRestart(logger, null);
         }
+
+        public static void RestartBeforeStart(this ILogger logger)
+        {
+            _restartBeforeStart(logger, null);
+        }
     }
 }

src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs

@@ -32,6 +32,8 @@ public class WebJobsScriptHostService : IHostedService, IScriptHostManager, IDis
         private readonly SlidingWindow<bool> _healthCheckWindow;
         private readonly Timer _hostHealthCheckTimer;
         private readonly SemaphoreSlim _hostStartSemaphore = new SemaphoreSlim(1, 1);
+        private readonly TaskCompletionSource<bool> _hostStartedSource = new TaskCompletionSource<bool>();
+        private readonly Task _hostStarted;
 
         private IHost _host;
         private CancellationTokenSource _startupLoopTokenSource;
@@ -61,6 +63,8 @@ public WebJobsScriptHostService(IOptionsMonitor<ScriptApplicationHostOptions> ap
             _healthMonitorOptions = healthMonitorOptions ?? throw new ArgumentNullException(nameof(healthMonitorOptions));
             _logger = loggerFactory.CreateLogger(ScriptConstants.LogCategoryHostGeneral);
 
+            _hostStarted = _hostStartedSource.Task;
+
             State = ScriptHostState.Default;
 
             if (ShouldMonitorHostHealth)
@@ -140,6 +144,12 @@ private async Task StartHostAsync(CancellationToken cancellationToken, int attem
             try
             {
                 await _hostStartSemaphore.WaitAsync();
+
+                // Now that we're inside the semaphore, set this task as completed. This prevents
+                // restarts from being invoked (via the PlaceholderSpecializationMiddleware) before
+                // the IHostedService has ever started.
+                _hostStartedSource.TrySetResult(true);
+
                 await UnsynchronizedStartHostAsync(cancellationToken, attemptCount, startupMode);
             }
             finally
@@ -155,11 +165,12 @@ private async Task StartHostAsync(CancellationToken cancellationToken, int attem
         /// </summary>
         private async Task UnsynchronizedStartHostAsync(CancellationToken cancellationToken, int attemptCount = 0, JobHostStartupMode startupMode = JobHostStartupMode.Normal)
         {
-            cancellationToken.ThrowIfCancellationRequested();
             IHost localHost = null;
 
             try
             {
+                cancellationToken.ThrowIfCancellationRequested();
+
                 // if we were in an error state retain that,
                 // otherwise move to default
                 if (State != ScriptHostState.Error)
@@ -332,6 +343,14 @@ public async Task StopAsync(CancellationToken cancellationToken)
 
         public async Task RestartHostAsync(CancellationToken cancellationToken)
         {
+            // Do not invoke a restart if the host has not yet been started. This can lead
+            // to invalid state.
+            if (!_hostStarted.IsCompleted)
+            {
+                _logger.RestartBeforeStart();
+                await _hostStarted;
+            }
+
             _logger.EnteringRestart();
 
             // If anything is mid-startup, cancel it.

test/WebJobs.Script.Tests/WebJobsScriptHostServiceTests.cs

@@ -327,6 +327,41 @@ public async Task DisposedHost_ServicesNotExposed()
             await startTask;
         }
 
+        [Fact]
+        public async Task HostRestart_BeforeStart_WaitsForStartToContinue()
+        {
+            _host = CreateMockHost();
+
+            var hostBuilder = new Mock<IScriptHostBuilder>();
+            hostBuilder.SetupSequence(b => b.BuildHost(It.IsAny<bool>(), It.IsAny<bool>()))
+                .Returns(_host.Object)
+                .Returns(_host.Object);
+
+            _webHostLoggerProvider = new TestLoggerProvider();
+            _loggerFactory = new LoggerFactory();
+            _loggerFactory.AddProvider(_webHostLoggerProvider);
+
+            _hostService = new WebJobsScriptHostService(
+               _monitor, hostBuilder.Object, _loggerFactory, _mockRootServiceProvider.Object, _mockRootScopeFactory.Object,
+               _mockScriptWebHostEnvironment.Object, _mockEnvironment.Object, _hostPerformanceManager, _healthMonitorOptions);
+
+            // Simulate a call to specialize coming from the PlaceholderSpecializationMiddleware. This
+            // can happen before we ever start the service, which could create invalid state.
+            Task restartTask = _hostService.RestartHostAsync(CancellationToken.None);
+
+            await _hostService.StartAsync(CancellationToken.None);
+            await restartTask;
+
+            var messages = _webHostLoggerProvider.GetAllLogMessages();
+
+            // The CancellationToken is canceled quick enough that we never build the initial host, so the
+            // only one we start is the restarted/specialized one.
+            Assert.NotNull(messages.Single(p => p.EventId.Id == 513)); // "Building" message
+            Assert.NotNull(messages.Single(p => p.EventId.Id == 514)); // "StartupWasCanceled" message
+            Assert.NotNull(messages.Single(p => p.EventId.Id == 520)); // "RestartBeforeStart" message
+            _host.Verify(p => p.StartAsync(It.IsAny<CancellationToken>()), Times.Once);
+        }
+
         public void RestartHost()
         {
             _hostService.RestartHostAsync(CancellationToken.None).Wait();