Worker restart on function timeout (#5794)

Restart worker process(es) on function timeout for out-of-proc languages instead of shutting down host.

Author: yojagad

sample/Node/HttpTrigger-Timeout/function.json

@@ -0,0 +1,15 @@
+{
+    "bindings": [
+        {
+            "type": "httpTrigger",
+            "name": "req",
+            "direction": "in",
+            "methods": [ "get" ]
+        },
+        {
+            "type": "http",
+            "name": "$return",
+            "direction": "out"
+        }
+    ]
+}

sample/Node/HttpTrigger-Timeout/index.js

@@ -0,0 +1,35 @@
+var test = require('../Shared/test');
+
+module.exports = function (context, req) {
+    context.log('Node.js HTTP trigger function processed a request. Name=%s', req.query.name);
+
+    var headerValue = req.headers['test-header'];
+    if (headerValue) {
+        context.log('test-header=' + headerValue);
+    }
+
+    var res;
+    if (typeof req.query.name === 'undefined') {
+        res = {
+            status: 400,
+            body: "Please pass a name on the query string",
+            headers: {
+                'Content-Type': 'text/plain'
+            }
+        };
+    }
+    else {
+		var e = new Date().getTime() + (30 * 1000);
+		while (new Date().getTime() <= e) {}
+		res = {
+				status: 200,
+				body: test.greeting(req.query.name),
+				headers: {
+					'Content-Type': 'text/plain',
+					'Shared-Module': test.timestamp
+				}
+		};
+    }
+
+    context.done(null, res);
+};

sample/Node/host.json

@@ -8,7 +8,7 @@
         "healthCheckThreshold": 6,
         "counterThreshold": 0.80
     },
-    "functionTimeout": "00:05:00",
+    "functionTimeout": "00:00:10",
     "logging": {
         "fileLoggingMode": "always"        
     },

src/WebJobs.Script.WebHost/WebScriptHostExceptionHandler.cs

@@ -7,6 +7,7 @@
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.Azure.WebJobs.Host;
 using Microsoft.Azure.WebJobs.Host.Timers;
+using Microsoft.Azure.WebJobs.Script.Workers;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.Azure.WebJobs.Script.WebHost
@@ -15,10 +16,12 @@ public class WebScriptHostExceptionHandler : IWebJobsExceptionHandler
     {
         private readonly IApplicationLifetime _applicationLifetime;
         private readonly ILogger _logger;
+        private readonly IFunctionInvocationDispatcherFactory _functionInvocationDispatcherFactory;
 
-        public WebScriptHostExceptionHandler(IApplicationLifetime applicationLifetime, ILogger<WebScriptHostExceptionHandler> logger)
+        public WebScriptHostExceptionHandler(IApplicationLifetime applicationLifetime, ILogger<WebScriptHostExceptionHandler> logger, IFunctionInvocationDispatcherFactory functionInvocationDispatcherFactory)
         {
             _applicationLifetime = applicationLifetime ?? throw new ArgumentNullException(nameof(applicationLifetime));
+            _functionInvocationDispatcherFactory = functionInvocationDispatcherFactory;
             _logger = logger ?? throw new ArgumentNullException(nameof(logger));
         }
 
@@ -39,17 +42,30 @@ public async Task OnTimeoutExceptionAsync(ExceptionDispatchInfo exceptionInfo, T
                 }
             }
 
-            LogErrorAndFlush("A function timeout has occurred. Host is shutting down.", exceptionInfo.SourceException);
-
             // We can't wait on this as it may cause a deadlock if the timeout was fired
             // by a Listener that cannot stop until it has completed.
             // TODO: DI (FACAVAL) The shutdown call will invoke the host stop... but we may need to do this
             // explicitly in order to pass the timeout.
             // Task ignoreTask = _hostManager.StopAsync();
             // Give the manager and all running tasks some time to shut down gracefully.
             //await Task.Delay(timeoutGracePeriod);
-
-            _applicationLifetime.StopApplication();
+            IFunctionInvocationDispatcher functionInvocationDispatcher = _functionInvocationDispatcherFactory.GetFunctionDispatcher();
+            if (functionInvocationDispatcher.State.Equals(FunctionInvocationDispatcherState.Initialized))
+            {
+                _logger.LogWarning($"A function timeout has occurred. Restarting worker process executing invocationId '{timeoutException.InstanceId}'.", exceptionInfo.SourceException);
+                bool result = await functionInvocationDispatcher.RestartWorkerWithInvocationIdAsync(timeoutException.InstanceId.ToString());
+                if (!result)
+                {
+                    _logger.LogWarning($"Restarting all language worker processes since invocation Id '{timeoutException.InstanceId}' was not found.", exceptionInfo.SourceException);
+                    await functionInvocationDispatcher.RestartAllWorkersAsync();
+                }
+                _logger.LogWarning("Restart of language worker process(es) completed.", exceptionInfo.SourceException);
+            }
+            else
+            {
+                LogErrorAndFlush("A function timeout has occurred. Host is shutting down.", exceptionInfo.SourceException);
+                _applicationLifetime.StopApplication();
+            }
         }
 
         public Task OnUnhandledExceptionAsync(ExceptionDispatchInfo exceptionInfo)

src/WebJobs.Script/ScriptHostBuilderExtensions.cs

@@ -142,7 +142,6 @@ public static IHostBuilder AddScriptHostCore(this IHostBuilder builder, ScriptAp
 
                 //Worker Function Invocation dispatcher
                 services.AddSingleton<IFunctionInvocationDispatcherFactory, FunctionInvocationDispatcherFactory>();
-
                 services.AddSingleton<IScriptJobHost>(p => p.GetRequiredService<ScriptHost>());
                 services.AddSingleton<IJobHost>(p => p.GetRequiredService<ScriptHost>());
                 services.AddSingleton<IFunctionMetadataManager, FunctionMetadataManager>();

src/WebJobs.Script/Workers/Http/HttpFunctionInvocationDispatcher.cs

@@ -185,5 +185,16 @@ public Task ShutdownAsync()
         {
             return Task.CompletedTask;
         }
+
+        public async Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId)
+        {
+            await DisposeAndRestartWorkerChannel(_httpWorkerChannel.Id);    // Since there's only one channel for httpworker
+            return true;
+        }
+
+        public async Task RestartAllWorkersAsync()
+        {
+            await DisposeAndRestartWorkerChannel(_httpWorkerChannel.Id);    // Since there's only one channel for httpworker
+        }
     }
 }

src/WebJobs.Script/Workers/IFunctionInvocationDispatcher.cs

@@ -20,5 +20,9 @@ public interface IFunctionInvocationDispatcher : IDisposable
         Task InitializeAsync(IEnumerable<FunctionMetadata> functions, CancellationToken cancellationToken = default);
 
         Task ShutdownAsync();
+
+        Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId);
+
+        Task RestartAllWorkersAsync();
     }
 }

src/WebJobs.Script/Workers/Rpc/FunctionRegistration/RpcFunctionInvocationDispatcher.cs

@@ -248,7 +248,7 @@ public async Task InvokeAsync(ScriptInvocationContext invocationContext)
             }
         }
 
-        internal async Task<IEnumerable<IRpcWorkerChannel>> GetInitializedWorkerChannelsAsync()
+        internal async Task<IEnumerable<IRpcWorkerChannel>> GetAllWorkerChannelsAsync()
         {
             Dictionary<string, TaskCompletionSource<IRpcWorkerChannel>> webhostChannelDictionary = _webHostLanguageWorkerChannelManager.GetChannels(_workerRuntime);
             List<IRpcWorkerChannel> webhostChannels = null;
@@ -264,11 +264,18 @@ internal async Task<IEnumerable<IRpcWorkerChannel>> GetInitializedWorkerChannels
                 }
             }
             IEnumerable<IRpcWorkerChannel> workerChannels = webhostChannels == null ? _jobHostLanguageWorkerChannelManager.GetChannels() : webhostChannels.Union(_jobHostLanguageWorkerChannelManager.GetChannels());
+            return workerChannels;
+        }
+
+        internal async Task<IEnumerable<IRpcWorkerChannel>> GetInitializedWorkerChannelsAsync()
+        {
+            IEnumerable<IRpcWorkerChannel> workerChannels = await GetAllWorkerChannelsAsync();
             IEnumerable<IRpcWorkerChannel> initializedWorkers = workerChannels.Where(ch => ch.IsChannelReadyForInvocations());
             if (initializedWorkers.Count() > _maxProcessCount)
             {
                 throw new InvalidOperationException($"Number of initialized language workers exceeded:{initializedWorkers.Count()} exceeded maxProcessCount: {_maxProcessCount}");
             }
+
             return initializedWorkers;
         }
 
@@ -399,5 +406,33 @@ public void Dispose()
             State = FunctionInvocationDispatcherState.Disposing;
             Dispose(true);
         }
+
+        public async Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId)
+        {
+            // Dispose and restart errored channel with the particular invocation id
+            State = FunctionInvocationDispatcherState.WorkerProcessRestarting;
+            var channels = await GetInitializedWorkerChannelsAsync();
+            foreach (var channel in channels)
+            {
+                if (channel.IsExecutingInvocation(invocationId))
+                {
+                    _logger.LogInformation($"Restarting channel '{channel.Id}' that is executing invocation '{invocationId}' and timed out.");
+                    await DisposeAndRestartWorkerChannel(_workerRuntime, channel.Id);
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        public async Task RestartAllWorkersAsync()
+        {
+            State = FunctionInvocationDispatcherState.WorkerProcessRestarting;
+            var channels = await GetAllWorkerChannelsAsync();
+            foreach (var channel in channels)
+            {
+                _logger.LogInformation($"Restarting channel '{channel.Id}' that is as part of restarting all channels.");
+                await DisposeAndRestartWorkerChannel(_workerRuntime, channel.Id);
+            }
+        }
     }
 }

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

@@ -27,5 +27,7 @@ public interface IRpcWorkerChannel
         Task StartWorkerProcessAsync();
 
         Task DrainInvocationsAsync();
+
+        bool IsExecutingInvocation(string invocationId);
     }
 }

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

@@ -114,7 +114,7 @@ internal RpcWorkerChannel(
 
         public bool IsChannelReadyForInvocations()
         {
-            return _state.HasFlag(RpcWorkerChannelState.InvocationBuffersInitialized | RpcWorkerChannelState.Initialized);
+            return !_disposing && !_disposed && _state.HasFlag(RpcWorkerChannelState.InvocationBuffersInitialized | RpcWorkerChannelState.Initialized);
         }
 
         public async Task StartWorkerProcessAsync()
@@ -513,5 +513,10 @@ public async Task DrainInvocationsAsync()
                 await currContext.ResultSource.Task;
             }
         }
+
+        public bool IsExecutingInvocation(string invocationId)
+        {
+            return _executingInvocations.ContainsKey(invocationId);
+        }
     }
 }