Added support to gracefully shutdown language worker

src/WebJobs.Script.Grpc/Channel/GrpcWorkerChannel.cs

@@ -14,6 +14,7 @@
 using System.Threading.Tasks;
 using System.Threading.Tasks.Dataflow;
 using Google.Protobuf.Collections;
+using Google.Protobuf.WellKnownTypes;
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
@@ -901,10 +902,36 @@ protected virtual void Dispose(bool disposing)
 
         public void Dispose()
         {
+            StopWorkerProcess();
             _disposing = true;
             Dispose(true);
         }
 
+        private void StopWorkerProcess()
+        {
+            bool capabilityEnabled = !string.IsNullOrEmpty(_workerCapabilities.GetCapabilityState(RpcWorkerConstants.HandlesWorkerTerminateMessage));
+            if (!capabilityEnabled)
+            {
+                return;
+            }
+
+            int gracePeriod = WorkerConstants.WorkerTerminateGracePeriodInSeconds;
+
+            var workerTerminate = new WorkerTerminate()
+            {
+                GracePeriod = Duration.FromTimeSpan(TimeSpan.FromSeconds(gracePeriod))
+            };
+
+            _workerChannelLogger.LogDebug($"Sending WorkerTerminate message with grace period {gracePeriod} seconds.");
+
+            SendStreamingMessage(new StreamingMessage
+            {
+                WorkerTerminate = workerTerminate
+            });
+
+            WorkerProcess.WaitForProcessExitInMilliSeconds(gracePeriod * 1000);
+        }
+
         public async Task DrainInvocationsAsync()
         {
             _workerChannelLogger.LogDebug($"Count of in-buffer invocations waiting to be drained out: {_executingInvocations.Count}");

src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs

@@ -434,7 +434,7 @@ public async Task StopAsync(CancellationToken cancellationToken)
             var currentHost = ActiveHost;
             ActiveHost = null;
             Task stopTask = Orphan(currentHost, cancellationToken);
-            Task result = await Task.WhenAny(stopTask, Task.Delay(TimeSpan.FromSeconds(10), cancellationToken));
+            Task result = await Task.WhenAny(stopTask, Task.Delay(TimeSpan.FromSeconds(10)));
 
             if (result != stopTask)
             {

src/WebJobs.Script/Workers/ProcessManagement/EmptyProcessRegistry.cs

@@ -7,6 +7,10 @@ namespace Microsoft.Azure.WebJobs.Script.Workers
 {
     internal class EmptyProcessRegistry : IProcessRegistry
     {
-        public bool Register(Process process) => true;
+        public bool Register(WorkerProcess process) => true;
+
+        public void Close()
+        {
+        }
     }
 }

src/WebJobs.Script/Workers/ProcessManagement/IProcessRegistry.cs

@@ -9,6 +9,8 @@ namespace Microsoft.Azure.WebJobs.Script.Workers
     internal interface IProcessRegistry
     {
         // Registers processes to ensure that they are cleaned up on host exit.
-        bool Register(Process process);
+        bool Register(WorkerProcess process);
+
+        void Close();
     }
 }

src/WebJobs.Script/Workers/ProcessManagement/IWorkerProcess.cs

@@ -11,5 +11,7 @@ public interface IWorkerProcess
         int Id { get; }
 
         Task StartProcessAsync();
+
+        void WaitForProcessExitInMilliSeconds(int waitTime);
     }
 }

src/WebJobs.Script/Workers/ProcessManagement/JobObjectRegistry.cs

@@ -12,6 +12,7 @@ internal class JobObjectRegistry : IProcessRegistry, IDisposable
     {
         private IntPtr _handle;
         private bool _disposed = false;
+        private WorkerProcess _workerProcess = null;
 
         public JobObjectRegistry()
         {
@@ -37,9 +38,10 @@ public JobObjectRegistry()
             }
         }
 
-        public bool Register(Process proc)
+        public bool Register(WorkerProcess workerProcess)
         {
-            return AssignProcessToJobObject(_handle, proc.Handle);
+            _workerProcess = workerProcess;
+            return AssignProcessToJobObject(_handle, _workerProcess.Process.Handle);
         }
 
         [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
@@ -78,6 +80,8 @@ private void Dispose(bool disposing)
 
         public void Close()
         {
+            _workerProcess?.ProcessWaitingForTermination.Task.GetAwaiter().GetResult();
+
             if (_handle != IntPtr.Zero)
             {
                 CloseHandle(_handle);

src/WebJobs.Script/Workers/ProcessManagement/WorkerProcess.cs

@@ -11,6 +11,7 @@
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
 using Microsoft.Azure.WebJobs.Script.Eventing;
+using Microsoft.Azure.WebJobs.Script.Workers.Rpc;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 
@@ -53,8 +54,9 @@ internal WorkerProcess(IScriptEventManager eventManager, IProcessRegistry proces
 
         internal Queue<string> ProcessStdErrDataQueue => _processStdErrDataQueue;
 
-        // for testing
-        internal Process Process { get; set; }
+        public Process Process { get; set; }
+
+        public TaskCompletionSource<bool> ProcessWaitingForTermination { get; set; } = new TaskCompletionSource<bool>();
 
         internal abstract Process CreateWorkerProcess();
 
@@ -78,7 +80,7 @@ public Task StartProcessAsync()
                     Process.BeginOutputReadLine();
 
                     // Register process only after it starts
-                    _processRegistry?.Register(Process);
+                    _processRegistry?.Register(this);
 
                     RegisterWithProcessMonitor();
 
@@ -204,10 +206,18 @@ internal void BuildAndLogConsoleLog(string msg, LogLevel level)
 
         internal abstract void HandleWorkerProcessRestart();
 
+        public void WaitForProcessExitInMilliSeconds(int waitTime)
+        {
+            Process.WaitForExit(waitTime);
+        }
+
         public void Dispose()
         {
             Disposing = true;
             // best effort process disposal
+
+            ProcessWaitingForTermination.SetResult(false);
+
             try
             {
                 _eventSubscription?.Dispose();

src/WebJobs.Script/Workers/Rpc/RpcWorkerConstants.cs

@@ -51,6 +51,7 @@ public static class RpcWorkerConstants
         public const string AcceptsListOfFunctionLoadRequests = "AcceptsListOfFunctionLoadRequests";
         public const string EnableUserCodeException = "EnableUserCodeException";
         public const string SupportsLoadResponseCollection = "SupportsLoadResponseCollection";
+        public const string HandlesWorkerTerminateMessage = "HandlesWorkerTerminateMessage";
 
         // Host Capabilities
         public const string V2Compatable = "V2Compatable";

src/WebJobs.Script/Workers/WorkerConstants.cs

@@ -9,6 +9,7 @@ public static class WorkerConstants
         public const string HttpScheme = "http";
 
         public const int WorkerReadyCheckPollingIntervalMilliseconds = 25;
+        public const int WorkerTerminateGracePeriodInSeconds = 5;
         public const string WorkerConfigFileName = "worker.config.json";
         public const string DefaultWorkersDirectoryName = "workers";
 

test/WebJobs.Script.Tests/Workers/Rpc/GrpcWorkerChannelTests.cs

@@ -145,6 +145,48 @@ await Assert.ThrowsAsync<TaskCanceledException>(async () =>
             });
         }
 
+        [Fact]
+        public void WorkerChannel_Dispose_With_WorkerTerminateCapability()
+        {
+            var initTask = _workerChannel.StartWorkerProcessAsync(CancellationToken.None);
+
+            IDictionary<string, string> capabilities = new Dictionary<string, string>()
+            {
+                { RpcWorkerConstants.HandlesWorkerTerminateMessage, "1" }
+            };
+
+            StartStream startStream = new StartStream()
+            {
+                WorkerId = _workerId
+            };
+
+            StreamingMessage startStreamMessage = new StreamingMessage()
+            {
+                StartStream = startStream
+            };
+
+            // Send worker init request and enable the capabilities
+            GrpcEvent rpcEvent = new GrpcEvent(_workerId, startStreamMessage);
+            _workerChannel.SendWorkerInitRequest(rpcEvent);
+            _testFunctionRpcService.PublishWorkerInitResponseEvent(capabilities);
+
+            _workerChannel.Dispose();
+            var traces = _logger.GetLogMessages();
+            var expectedLogMsg = $"Sending WorkerTerminate message with grace period {WorkerConstants.WorkerTerminateGracePeriodInSeconds} seconds.";
+            Assert.True(traces.Any(m => string.Equals(m.FormattedMessage, expectedLogMsg)));
+        }
+
+        [Fact]
+        public void WorkerChannel_Dispose_Without_WorkerTerminateCapability()
+        {
+            var initTask = _workerChannel.StartWorkerProcessAsync(CancellationToken.None);
+
+            _workerChannel.Dispose();
+            var traces = _logger.GetLogMessages();
+            var expectedLogMsg = $"Sending WorkerTerminate message with grace period {WorkerConstants.WorkerTerminateGracePeriodInSeconds} seconds.";
+            Assert.False(traces.Any(m => string.Equals(m.FormattedMessage, expectedLogMsg)));
+        }
+
         [Fact]
         public async Task StartWorkerProcessAsync_Invoked_SetupFunctionBuffers_Verify_ReadyForInvocation()
         {

test/WebJobs.Script.Tests/Workers/Rpc/RpcWorkerProcessTests.cs

@@ -18,6 +18,7 @@ public class RpcWorkerProcessTests
         private RpcWorkerProcess _rpcWorkerProcess;
         private Mock<IScriptEventManager> _eventManager;
         private Mock<IHostProcessMonitor> _hostProcessMonitorMock;
+        private TestLogger _logger = new TestLogger("test");
 
         public RpcWorkerProcessTests()
         {
@@ -45,7 +46,7 @@ public RpcWorkerProcessTests()
                 _eventManager.Object,
                 workerProcessFactory.Object,
                 processRegistry.Object,
-                new TestLogger("test"),
+                _logger,
                 languageWorkerConsoleLogSource.Object,
                 new TestMetricsLogger(),
                 serviceProviderMock.Object);
@@ -179,5 +180,23 @@ public void HandleWorkerProcessExitError_PublishesWorkerRestartEvent_OnIntention
             _eventManager.Verify(_ => _.Publish(It.IsAny<WorkerRestartEvent>()), Times.Once());
             _eventManager.Verify(_ => _.Publish(It.IsAny<WorkerErrorEvent>()), Times.Never());
         }
+
+        [Fact]
+        public void WorkerProcess_Dispose()
+        {
+            Process process = new Process();
+            ProcessStartInfo startInfo = new ProcessStartInfo();
+            startInfo.WindowStyle = ProcessWindowStyle.Hidden;
+            startInfo.FileName = "cmd.exe";
+            startInfo.Arguments = $"ls";
+            process.StartInfo = startInfo;
+            process.Start();
+
+            _rpcWorkerProcess.Process = process;
+            _rpcWorkerProcess.Dispose();
+            var traces = _logger.GetLogMessages();
+            var disposeLogs = traces.Where(m => string.Equals(m.FormattedMessage, "Worker process has not exited despite waiting for 1000 ms"));
+            Assert.False(disposeLogs.Any());
+        }
     }
 }