Call EnableDrainModeAsync on ApplicationStopping (#7262)

Co-authored-by: Pragna Gopa <[email protected]>

Author: ahmelsayed

release_notes.md

@@ -4,8 +4,8 @@
 -->
 - Added a feature flag to opt out of the default behavior where the host sets the environment name to `Development` when running in debug mode. To disable the behavior, set the app setting: `AzureWebJobsFeatureFlags` to `DisableDevModeInDebug`
 - Reorder CORS and EasyAuth middleware to prevent EasyAuth from blocking CORS requests (#7315)
-
 - Updated PowerShell Worker (PS7) to [3.0.833](https://github.com/Azure/azure-functions-powershell-worker/releases/tag/v3.0.833)
+- Call EnableDrainModeAsync on ApplicationStopping (#7262)
 
 **Release sprint:** Sprint 100
 [ [bugs](https://github.com/Azure/azure-functions-host/issues?q=is%3Aissue+milestone%3A%22Functions+Sprint+100%22+label%3Abug+is%3Aclosed) | [features](https://github.com/Azure/azure-functions-host/issues?q=is%3Aissue+milestone%3A%22Functions+Sprint+100%22+label%3Afeature+is%3Aclosed) ]

src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs

@@ -11,6 +11,7 @@
 using Microsoft.ApplicationInsights.AspNetCore.Extensions;
 using Microsoft.ApplicationInsights.Extensibility;
 using Microsoft.ApplicationInsights.Extensibility.Implementation.ApplicationId;
+using Microsoft.Azure.WebJobs.Host;
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Configuration;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
@@ -693,6 +694,17 @@ private void RegisterApplicationLifetimeEvents()
             _applicationLifetime.ApplicationStopping.Register(() =>
             {
                 Interlocked.Exchange(ref _applicationStopping, 1);
+                if (_environment.DrainOnApplicationStoppingEnabled())
+                {
+                    var drainModeManager = _host?.Services.GetService<IDrainModeManager>();
+                    if (drainModeManager != null)
+                    {
+                        _logger.LogDebug("Application Stopping: initiate drain mode");
+                        drainModeManager.EnableDrainModeAsync(CancellationToken.None);
+                        // Workaround until https://github.com/Azure/azure-functions-host/issues/7188 is addressed.
+                        Thread.Sleep(TimeSpan.FromMinutes(10));
+                    }
+                }
             });
 
             _applicationLifetime.ApplicationStopped.Register(() =>

src/WebJobs.Script/Environment/EnvironmentExtensions.cs

@@ -447,5 +447,11 @@ public static bool SupportsAzureFileShareMount(this IEnvironment environment)
             return string.Equals(environment.GetFunctionsWorkerRuntime(), RpcWorkerConstants.PowerShellLanguageWorkerName,
                 StringComparison.OrdinalIgnoreCase);
         }
+
+        public static bool DrainOnApplicationStoppingEnabled(this IEnvironment environment)
+        {
+            return !string.IsNullOrEmpty(environment.GetEnvironmentVariable(KubernetesServiceHost)) ||
+                (bool.TryParse(environment.GetEnvironmentVariable(DrainOnApplicationStopping), out bool v) && v);
+        }
     }
 }

src/WebJobs.Script/Environment/EnvironmentSettingNames.cs

@@ -120,5 +120,7 @@ public static class EnvironmentSettingNames
         public const string AntaresComputerName = "COMPUTERNAME";
 
         public const string AppKind = "APP_KIND";
+
+        public const string DrainOnApplicationStopping = "FUNCTIONS_ENABLE_DRAIN_ON_APP_STOPPING";
     }
 }

test/WebJobs.Script.Tests/Extensions/EnvironmentExtensionsTests.cs

@@ -219,5 +219,23 @@ public void Returns_SupportsAzureFileShareMount(string workerRuntime, bool useLo
             environment.SetEnvironmentVariable(FunctionWorkerRuntime, workerRuntime);
             Assert.Equal(supportsAzureFileShareMount, environment.SupportsAzureFileShareMount());
         }
+
+        [Theory]
+        [InlineData(null, null, false)]
+        [InlineData("", null, false)]
+        [InlineData("", "", false)]
+        [InlineData("", "false", false)]
+        [InlineData("", "true", true)]
+        [InlineData("10.0.0.1", "", true)]
+        [InlineData("10.0.0.1", "true", true)]
+        [InlineData("10.0.0.1", "false", true)]
+        [InlineData("10.0.0.1", null, true)]
+        public void IsDrainOnApplicationStopping_ReturnsExpectedResult(string serviceHostValue, string drainOnStoppingValue, bool expected)
+        {
+            var environment = new TestEnvironment();
+            environment.SetEnvironmentVariable(KubernetesServiceHost, serviceHostValue);
+            environment.SetEnvironmentVariable(DrainOnApplicationStopping, drainOnStoppingValue);
+            Assert.Equal(expected, environment.DrainOnApplicationStoppingEnabled());
+        }
     }
 }