Add 5 second grace period for langauge worker channels before JobHost shutsdown (#4986)

yojagad · pragnagopa · commit 4a9aa46e9442 · 2019-10-22T13:59:47.000-07:00
diff --git a/src/WebJobs.Script.WebHost/DependencyInjection/DependencyValidator/DependencyValidator.cs b/src/WebJobs.Script.WebHost/DependencyInjection/DependencyValidator/DependencyValidator.cs
@@ -44,6 +44,7 @@ private static ExpectedDependencyBuilder CreateExpectedDependencies()
                 .Expect<PrimaryHostCoordinator>()
                 .Expect<FileMonitoringService>()
                 .Expect<LanguageWorkerConsoleLogService>()
+                .Expect<FunctionDispatcherShutdownManager>()
                 .Optional<FunctionsScaleMonitorService>()
                 .Optional<FuncAppFileProvisioningService>() // Used by powershell.
                 .Optional<JobHostService>() // Missing when host is offline.
diff --git a/src/WebJobs.Script/OutOfProc/FunctionDispatcherShutdownManager.cs b/src/WebJobs.Script/OutOfProc/FunctionDispatcherShutdownManager.cs
@@ -0,0 +1,29 @@
+﻿// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Hosting;
+
+namespace Microsoft.Azure.WebJobs.Script.Rpc
+{
+    internal sealed class FunctionDispatcherShutdownManager : IHostedService
+    {
+        private readonly IFunctionDispatcher _functionDispatcher;
+
+        public FunctionDispatcherShutdownManager(IFunctionDispatcher functionDispatcher)
+        {
+            _functionDispatcher = functionDispatcher;
+        }
+
+        public Task StartAsync(CancellationToken cancellationToken)
+        {
+            return Task.CompletedTask;
+        }
+
+        public async Task StopAsync(CancellationToken cancellationToken)
+        {
+            await _functionDispatcher.ShutdownAsync();
+        }
+    }
+}
diff --git a/src/WebJobs.Script/OutOfProc/Http/HttpFunctionInvocationDispatcher.cs b/src/WebJobs.Script/OutOfProc/Http/HttpFunctionInvocationDispatcher.cs
@@ -171,5 +171,10 @@ public void Dispose()
             _disposing = true;
             Dispose(true);
         }
+
+        public Task ShutdownAsync()
+        {
+            return Task.CompletedTask;
+        }
     }
 }
diff --git a/src/WebJobs.Script/OutOfProc/IFunctionDispatcher.cs b/src/WebJobs.Script/OutOfProc/IFunctionDispatcher.cs
@@ -15,5 +15,7 @@ public interface IFunctionDispatcher : IDisposable
         Task InvokeAsync(ScriptInvocationContext invocationContext);
 
         Task InitializeAsync(IEnumerable<FunctionMetadata> functions);
+
+        Task ShutdownAsync();
     }
 }
diff --git a/src/WebJobs.Script/OutOfProc/Rpc/FunctionRegistration/RpcFunctionInvocationDispatcher.cs b/src/WebJobs.Script/OutOfProc/Rpc/FunctionRegistration/RpcFunctionInvocationDispatcher.cs
@@ -30,6 +30,7 @@ internal class RpcFunctionInvocationDispatcher : IFunctionDispatcher
         private readonly IScriptJobHostEnvironment _scriptJobHostEnvironment;
         private readonly int _debounceSeconds = 10;
         private readonly int _maxAllowedProcessCount = 10;
+        private readonly TimeSpan _shutdownTimeout = TimeSpan.FromSeconds(10);
         private readonly TimeSpan thresholdBetweenRestarts = TimeSpan.FromMinutes(OutOfProcConstants.WorkerRestartErrorIntervalThresholdInMinutes);
 
         private IScriptEventManager _eventManager;
@@ -211,6 +212,23 @@ public async Task InitializeAsync(IEnumerable<FunctionMetadata> functions)
             }
         }
 
+        public async Task ShutdownAsync()
+        {
+            _logger.LogDebug($"Waiting for {nameof(RpcFunctionInvocationDispatcher)} to shutdown");
+            Task timeoutTask = Task.Delay(_shutdownTimeout);
+            IList<Task> workerChannelTasks = (await GetInitializedWorkerChannelsAsync()).Select(a => a.DrainInvocationsAsync()).ToList();
+            Task completedTask = await Task.WhenAny(timeoutTask, Task.WhenAll(workerChannelTasks));
+
+            if (completedTask.Equals(timeoutTask))
+            {
+                _logger.LogDebug($"Draining invocations from language worker channel timed out. Shutting down '{nameof(RpcFunctionInvocationDispatcher)}'");
+            }
+            else
+            {
+                _logger.LogDebug($"Draining invocations from language worker channel completed. Shutting down '{nameof(RpcFunctionInvocationDispatcher)}'");
+            }
+        }
+
         public async Task InvokeAsync(ScriptInvocationContext invocationContext)
         {
             try
diff --git a/src/WebJobs.Script/OutOfProc/Rpc/ILanguageWorkerChannel.cs b/src/WebJobs.Script/OutOfProc/Rpc/ILanguageWorkerChannel.cs
@@ -24,5 +24,7 @@ public interface ILanguageWorkerChannel
         Task SendFunctionEnvironmentReloadRequest();
 
         Task StartWorkerProcessAsync();
+
+        Task DrainInvocationsAsync();
     }
 }
diff --git a/src/WebJobs.Script/OutOfProc/Rpc/LanguageWorkerChannel.cs b/src/WebJobs.Script/OutOfProc/Rpc/LanguageWorkerChannel.cs
@@ -480,5 +480,14 @@ private bool IsTriggerMetadataPopulatedByWorker()
         {
             return !string.IsNullOrEmpty(_workerCapabilities.GetCapabilityState(LanguageWorkerConstants.RpcHttpTriggerMetadataRemoved));
         }
+
+        public async Task DrainInvocationsAsync()
+        {
+            _workerChannelLogger.LogDebug($"Count of in-buffer invocations waiting to be drained out: {_executingInvocations.Count}");
+            foreach (ScriptInvocationContext currContext in _executingInvocations.Values)
+            {
+                await currContext.ResultSource.Task;
+            }
+        }
     }
 }
diff --git a/src/WebJobs.Script/ScriptHostBuilderExtensions.cs b/src/WebJobs.Script/ScriptHostBuilderExtensions.cs
@@ -178,6 +178,7 @@ public static IHostBuilder AddScriptHostCore(this IHostBuilder builder, ScriptAp
 
                 services.AddSingleton<IHostedService, LanguageWorkerConsoleLogService>();
                 services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService, PrimaryHostCoordinator>());
+                services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService, FunctionDispatcherShutdownManager>());
 
                 if (SystemEnvironment.Instance.IsRuntimeScaleMonitoringEnabled())
                 {
diff --git a/test/WebJobs.Script.Tests/FunctionDispatcherShutdownManagerTests.cs b/test/WebJobs.Script.Tests/FunctionDispatcherShutdownManagerTests.cs
@@ -0,0 +1,25 @@
+﻿// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Script.Rpc;
+using Microsoft.Extensions.Logging;
+using Moq;
+using Xunit;
+
+namespace Microsoft.Azure.WebJobs.Script.Tests
+{
+    public class FunctionDispatcherShutdownManagerTests
+    {
+        [Fact]
+        public async Task Test_StopAsync()
+        {
+            Mock<IFunctionDispatcher> functionDispatcher = new Mock<IFunctionDispatcher>();
+            functionDispatcher.Setup(a => a.ShutdownAsync()).Returns(Task.CompletedTask);
+            var functionDispatcherShutdownManager = new FunctionDispatcherShutdownManager(functionDispatcher.Object);
+            await functionDispatcherShutdownManager.StopAsync(CancellationToken.None);
+            functionDispatcher.Verify(a => a.ShutdownAsync(), Times.Once);
+        }
+    }
+}
diff --git a/test/WebJobs.Script.Tests/Rpc/FunctionDispatcherTests.cs b/test/WebJobs.Script.Tests/Rpc/FunctionDispatcherTests.cs
@@ -46,6 +46,48 @@ public async void Starting_MultipleWebhostChannels_Succeeds()
             Assert.Equal(0, finalJobhostChannelCount);
         }
 
+        [Fact]
+        public async Task ShutdownTests()
+        {
+            int expectedProcessCount = 2;
+            RpcFunctionInvocationDispatcher functionDispatcher = GetTestFunctionDispatcher(expectedProcessCount.ToString());
+            await functionDispatcher.InitializeAsync(GetTestFunctionsList(LanguageWorkerConstants.NodeLanguageWorkerName));
+            await WaitForJobhostWorkerChannelsToStartup(functionDispatcher, expectedProcessCount);
+
+            foreach (var currChannel in functionDispatcher.JobHostLanguageWorkerChannelManager.GetChannels())
+            {
+                var initializedChannel = (TestLanguageWorkerChannel)currChannel;
+                initializedChannel.ExecutionContexts.Add(Task.Factory.StartNew(() => { }));
+            }
+
+            await functionDispatcher.ShutdownAsync();
+            foreach (var currChannel in functionDispatcher.JobHostLanguageWorkerChannelManager.GetChannels())
+            {
+                Assert.True(((TestLanguageWorkerChannel)currChannel).ExecutionContexts.Count == 0);
+            }
+        }
+
+        [Fact]
+        public async Task ShutdownTests_WithInfinitelyRunningTasks_Timesout()
+        {
+            int expectedProcessCount = 2;
+            RpcFunctionInvocationDispatcher functionDispatcher = GetTestFunctionDispatcher(expectedProcessCount.ToString());
+            await functionDispatcher.InitializeAsync(GetTestFunctionsList(LanguageWorkerConstants.NodeLanguageWorkerName));
+            await WaitForJobhostWorkerChannelsToStartup(functionDispatcher, expectedProcessCount);
+
+            foreach (var currChannel in functionDispatcher.JobHostLanguageWorkerChannelManager.GetChannels())
+            {
+                var initializedChannel = (TestLanguageWorkerChannel)currChannel;
+                initializedChannel.ExecutionContexts.Add(new Task<bool>(() => true));   // A task that never starts and therefore never runs to completion
+            }
+
+            await functionDispatcher.ShutdownAsync();
+            foreach (var currChannel in functionDispatcher.JobHostLanguageWorkerChannelManager.GetChannels())
+            {
+                Assert.True(((TestLanguageWorkerChannel)currChannel).ExecutionContexts.Count > 0);
+            }
+        }
+
         [Fact]
         public void MaxProcessCount_Returns_Default()
         {
diff --git a/test/WebJobs.Script.Tests/Rpc/LanguageWorkerChannelTests.cs b/test/WebJobs.Script.Tests/Rpc/LanguageWorkerChannelTests.cs
@@ -140,24 +140,43 @@ public void SendSystemLogMessage_PublishesSystemLogMessage(RpcLog.Types.Level le
         [Fact]
         public void SendInvocationRequest_PublishesOutboundEvent()
         {
-            ScriptInvocationContext scriptInvocationContext = new ScriptInvocationContext()
-            {
-                FunctionMetadata = GetTestFunctionsList("node").FirstOrDefault(),
-                ExecutionContext = new ExecutionContext()
-                {
-                    InvocationId = Guid.NewGuid(),
-                    FunctionName = "js1",
-                    FunctionAppDirectory = _scriptRootPath,
-                    FunctionDirectory = _scriptRootPath
-                },
-                BindingData = new Dictionary<string, object>(),
-                Inputs = new List<(string name, DataType type, object val)>()
-            };
+            ScriptInvocationContext scriptInvocationContext = GetTestScriptInvocationContext(Guid.NewGuid(), null);
             _workerChannel.SendInvocationRequest(scriptInvocationContext);
             var traces = _logger.GetLogMessages();
             Assert.True(traces.Any(m => string.Equals(m.FormattedMessage, _expectedLogMsg)));
         }
 
+        [Fact]
+        public async Task Drain_Verify()
+        {
+            var resultSource = new TaskCompletionSource<ScriptInvocationResult>();
+            Guid invocationId = Guid.NewGuid();
+            LanguageWorkerChannel channel = new LanguageWorkerChannel(
+               _workerId,
+               _scriptRootPath,
+               _eventManager,
+               _testWorkerConfig,
+               _mockLanguageWorkerProcess.Object,
+               _logger,
+               _metricsLogger,
+               0);
+            channel.SetupFunctionInvocationBuffers(GetTestFunctionsList("node"));
+            ScriptInvocationContext scriptInvocationContext = GetTestScriptInvocationContext(invocationId, resultSource);
+            channel.SendInvocationRequest(scriptInvocationContext);
+            Task result = channel.DrainInvocationsAsync();
+            Assert.NotEqual(result.Status, TaskStatus.RanToCompletion);
+            channel.InvokeResponse(new InvocationResponse
+            {
+                InvocationId = invocationId.ToString(),
+                Result = new StatusResult
+                {
+                    Status = StatusResult.Types.Status.Success
+                },
+            });
+            await result;
+            Assert.Equal(result.Status, TaskStatus.RanToCompletion);
+        }
+
         [Fact]
         public void SendLoadRequests_PublishesOutboundEvents()
         {
@@ -289,6 +308,24 @@ private IEnumerable<FunctionMetadata> GetTestFunctionsList(string runtime)
             };
         }
 
+        private ScriptInvocationContext GetTestScriptInvocationContext(Guid invocationId, TaskCompletionSource<ScriptInvocationResult> resultSource)
+        {
+            return new ScriptInvocationContext()
+            {
+                FunctionMetadata = GetTestFunctionsList("node").FirstOrDefault(),
+                ExecutionContext = new ExecutionContext()
+                {
+                    InvocationId = invocationId,
+                    FunctionName = "js1",
+                    FunctionAppDirectory = _scriptRootPath,
+                    FunctionDirectory = _scriptRootPath
+                },
+                BindingData = new Dictionary<string, object>(),
+                Inputs = new List<(string name, DataType type, object val)>(),
+                ResultSource = resultSource
+            };
+        }
+
         private IEnumerable<FunctionMetadata> GetTestFunctionsList_WithDisabled(string runtime, string funcName)
         {
             var disabledList = new List<FunctionMetadata>()
diff --git a/test/WebJobs.Script.Tests/Rpc/TestLanguageWorkerChannel.cs b/test/WebJobs.Script.Tests/Rpc/TestLanguageWorkerChannel.cs
@@ -21,6 +21,7 @@ public class TestLanguageWorkerChannel : ILanguageWorkerChannel
         private string _runtime;
         private ILogger _testLogger;
         private LanguageWorkerChannelState _state;
+        private List<Task> _executionContexts;
 
         public TestLanguageWorkerChannel(string workerId, string runtime = null, IScriptEventManager eventManager = null, ILogger testLogger = null, bool isWebhostChannel = false, bool throwOnProcessStartUp = false)
         {
@@ -30,12 +31,15 @@ public TestLanguageWorkerChannel(string workerId, string runtime = null, IScript
             _runtime = runtime;
             _testLogger = testLogger;
             _throwOnProcessStartUp = throwOnProcessStartUp;
+            _executionContexts = new List<Task>();
         }
 
         public string Id => _workerId;
 
         public IDictionary<string, BufferBlock<ScriptInvocationContext>> FunctionInputBuffers => throw new NotImplementedException();
 
+        public List<Task> ExecutionContexts => _executionContexts;
+
         public LanguageWorkerChannelState State => _state;
 
         public void Dispose()
@@ -96,5 +100,11 @@ public void RaiseWorkerErrorWithCustomTimestamp(DateTime timestamp)
             Exception testEx = new Exception("Test Worker Error");
             _eventManager.Publish(new WorkerErrorEvent(_runtime, Id, testEx, timestamp));
         }
+
+        public async Task DrainInvocationsAsync()
+        {
+            await Task.WhenAll(ExecutionContexts);
+            ExecutionContexts.Clear();
+        }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -171,5 +171,10 @@ public void Dispose()`
`171`	`171`	`_disposing = true;`
`172`	`172`	`Dispose(true);`
`173`	`173`	`}`
	`174`	`+`
	`175`	`+ public Task ShutdownAsync()`
	`176`	`+ {`
	`177`	`+ return Task.CompletedTask;`
	`178`	`+ }`
`174`	`179`	`}`
`175`	`180`	`}`
Original file line number	Diff line number	Diff line change
`@@ -15,5 +15,7 @@ public interface IFunctionDispatcher : IDisposable`
`15`	`15`	`Task InvokeAsync(ScriptInvocationContext invocationContext);`
`16`	`16`
`17`	`17`	`Task InitializeAsync(IEnumerable<FunctionMetadata> functions);`
	`18`	`+`
	`19`	`+ Task ShutdownAsync();`
`18`	`20`	`}`
`19`	`21`	`}`
Original file line number	Diff line number	Diff line change
`@@ -24,5 +24,7 @@ public interface ILanguageWorkerChannel`
`24`	`24`	`Task SendFunctionEnvironmentReloadRequest();`
`25`	`25`
`26`	`26`	`Task StartWorkerProcessAsync();`
	`27`	`+`
	`28`	`+ Task DrainInvocationsAsync();`
`27`	`29`	`}`
`28`	`30`	`}`
Original file line number	Diff line number	Diff line change
`@@ -480,5 +480,14 @@ private bool IsTriggerMetadataPopulatedByWorker()`
`480`	`480`	`{`
`481`	`481`	`return !string.IsNullOrEmpty(_workerCapabilities.GetCapabilityState(LanguageWorkerConstants.RpcHttpTriggerMetadataRemoved));`
`482`	`482`	`}`
	`483`	`+`
	`484`	`+ public async Task DrainInvocationsAsync()`
	`485`	`+ {`
	`486`	`+ _workerChannelLogger.LogDebug($"Count of in-buffer invocations waiting to be drained out: {_executingInvocations.Count}");`
	`487`	`+ foreach (ScriptInvocationContext currContext in _executingInvocations.Values)`
	`488`	`+ {`
	`489`	`+ await currContext.ResultSource.Task;`
	`490`	`+ }`
	`491`	`+ }`
`483`	`492`	`}`
`484`	`493`	`}`
Original file line number	Diff line number	Diff line change
`@@ -178,6 +178,7 @@ public static IHostBuilder AddScriptHostCore(this IHostBuilder builder, ScriptAp`
`178`	`178`
`179`	`179`	`services.AddSingleton<IHostedService, LanguageWorkerConsoleLogService>();`
`180`	`180`	`services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService, PrimaryHostCoordinator>());`
	`181`	`+ services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService, FunctionDispatcherShutdownManager>());`
`181`	`182`
`182`	`183`	`if (SystemEnvironment.Instance.IsRuntimeScaleMonitoringEnabled())`
`183`	`184`	`{`