Ensure graceful host shutdown when shutdown file is created (Azure#2088…

…) (Azure#2730)

src/Microsoft.Azure.WebJobs.Host/Constants.cs

@@ -10,6 +10,7 @@ internal static class Constants
         public const string DevelopmentEnvironmentValue = "Development";
         public const string DynamicSku = "Dynamic";
         public const string AzureWebsiteSku = "WEBSITE_SKU";
+        public const string AzureWebJobsShutdownFile = "WEBJOBS_SHUTDOWN_FILE";
         public const string DateTimeFormatString = "yyyy'-'MM'-'dd'T'HH':'mm':'ss'.'fffK";
     }
 }

src/Microsoft.Azure.WebJobs.Host/Executors/JobHostContextFactory.cs

@@ -21,6 +21,7 @@
 using Microsoft.Azure.WebJobs.Host.Timers;
 using Microsoft.Azure.WebJobs.Host.Triggers;
 using Microsoft.Azure.WebJobs.Logging;
+using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 
@@ -49,6 +50,7 @@ internal class JobHostContextFactory : IJobHostContextFactory
         private readonly IDashboardLoggingSetup _dashboardLoggingSetup;
         private readonly IScaleMonitorManager _monitorManager;
         private readonly IDrainModeManager _drainModeManager;
+        private readonly IApplicationLifetime _applicationLifetime;
 
         public JobHostContextFactory(
             IDashboardLoggingSetup dashboardLoggingSetup,
@@ -70,7 +72,8 @@ public JobHostContextFactory(
             IConverterManager converterManager,
             IAsyncCollector<FunctionInstanceLogEntry> eventCollector,
             IScaleMonitorManager monitorManager,
-            IDrainModeManager drainModeManager)
+            IDrainModeManager drainModeManager,
+            IApplicationLifetime applicationLifetime)
         {
             _dashboardLoggingSetup = dashboardLoggingSetup;
             _functionExecutor = functionExecutor;
@@ -92,10 +95,18 @@ public JobHostContextFactory(
             _eventCollector = eventCollector;
             _monitorManager = monitorManager;
             _drainModeManager = drainModeManager;
+            _applicationLifetime = applicationLifetime;
         }
 
         public async Task<JobHostContext> Create(JobHost host, CancellationToken shutdownToken, CancellationToken cancellationToken)
         {
+            shutdownToken.Register(() =>
+            {
+                // when a shutdown is triggered we want to stop the application, to ensure the host
+                // shuts down gracefully
+                _applicationLifetime.StopApplication();
+            });
+
             using (CancellationTokenSource combinedCancellationSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, shutdownToken))
             {
                 CancellationToken combinedCancellationToken = combinedCancellationSource.Token;

src/Microsoft.Azure.WebJobs.Host/WebjobsShutdownWatcher.cs

@@ -5,6 +5,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
 using System.Threading;
+using Microsoft.Azure.WebJobs.Host;
 
 namespace Microsoft.Azure.WebJobs
 {
@@ -32,7 +33,7 @@ private WebJobsShutdownWatcher(CancellationTokenSource cancellationTokenSource,
         {
             // http://blog.amitapple.com/post/2014/05/webjobs-graceful-shutdown/#.U3aIXRFOVaQ
             // Antares will set this file to signify shutdown
-            _shutdownFile = Environment.GetEnvironmentVariable("WEBJOBS_SHUTDOWN_FILE");
+            _shutdownFile = Environment.GetEnvironmentVariable(Constants.AzureWebJobsShutdownFile);
             if (_shutdownFile == null)
             {
                 // If env var is not set, then no shutdown support

test/Microsoft.Azure.WebJobs.Host.FunctionalTests/JobHostTests.cs

@@ -10,16 +10,17 @@
 using System.Runtime.ExceptionServices;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Azure.Storage;
+using Microsoft.Azure.Storage.Blob;
+using Microsoft.Azure.Storage.Queue;
 using Microsoft.Azure.WebJobs.Host.Executors;
 using Microsoft.Azure.WebJobs.Host.Indexers;
 using Microsoft.Azure.WebJobs.Host.Listeners;
 using Microsoft.Azure.WebJobs.Host.TestCommon;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection.Extensions;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
-using Microsoft.Azure.Storage;
-using Microsoft.Azure.Storage.Blob;
-using Microsoft.Azure.Storage.Queue;
 using Moq;
 using Xunit;
 
@@ -314,6 +315,62 @@ public async Task StopAsync_WhenAlreadyStopping_ReturnsSameTask()
             }
         }
 
+        [Fact]
+        public async Task ShutdownFileCreated_HostShutsDownGracefully()
+        {
+            // Setup the directory where we'll be writing the shutdown file
+            string shutdownPath = Path.Combine(Path.GetTempPath(), "WebJobs", "shutdown");
+            Environment.SetEnvironmentVariable(Constants.AzureWebJobsShutdownFile, shutdownPath);
+            Directory.CreateDirectory(Path.GetDirectoryName(shutdownPath));
+            if (File.Exists(shutdownPath))
+            {
+                File.Delete(shutdownPath);
+            }
+
+            // Test hosted service that we'll use below to verify user registered hosted services
+            // are stopped gracefully on shutdown
+            var testService = new TestHostedService();
+
+            var builder = new HostBuilder().ConfigureDefaultTestHost();
+            builder.ConfigureServices(services =>
+            {
+                services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService>(testService));
+            });
+
+            var host = builder.Build();
+            using (host)
+            {
+                // start running the host without blocking so we can trigger the shutdown below
+                Task runTask = host.RunAsync();
+
+                // wait for the host to start completely
+                var loggerProvider = host.GetTestLoggerProvider();
+                IEnumerable<LogMessage> logs = null;
+                await TestHelpers.Await(() =>
+                {
+                    logs = loggerProvider.GetAllLogMessages();
+                    var startedLog = logs.SingleOrDefault(p => p.FormattedMessage == "Job host started");
+                    return startedLog != null;
+                });
+
+                // trigger the shutdown
+                File.Create(shutdownPath);
+
+                // wait for the run task to complete successfully
+                await TestHelpers.Await(() =>
+                {
+                    return runTask.IsCompleted;
+                });
+
+                Assert.True(testService.StartCalled);
+                Assert.True(testService.StopCalled);
+
+                logs = loggerProvider.GetAllLogMessages();
+                Assert.NotNull(logs.SingleOrDefault(p => p.FormattedMessage == "Stopping JobHost"));
+                Assert.NotNull(logs.SingleOrDefault(p => p.FormattedMessage == "Job host stopped"));
+            }
+        }
+
         [Fact]
         public async Task CallAsync_WithDictionary()
         {
@@ -747,5 +804,23 @@ public override void Log<TState>(Microsoft.Extensions.Logging.LogLevel logLevel,
                 base.Log(logLevel, eventId, state, exception, formatter);
             }
         }
+
+        public class TestHostedService : IHostedService
+        {
+            public bool StartCalled { get; set; }
+            public bool StopCalled { get; set; }
+
+            public Task StartAsync(CancellationToken cancellationToken)
+            {
+                StartCalled = true;
+                return Task.CompletedTask;
+            }
+
+            public Task StopAsync(CancellationToken cancellationToken)
+            {
+                StopCalled = true;
+                return Task.CompletedTask;
+            }
+        }
     }
 }