Fail in-flight invocations when worker channel shuts down (#11159)

satvu · web-flow · commit 8913ece85f92 · 2025-09-17T13:04:52.000-07:00
diff --git a/release_notes.md b/release_notes.md
@@ -12,3 +12,4 @@
 - Setting current activity status for failed invocations (#11313)
 - Adding test coverage for `Utility.IsAzureMonitorLoggingEnabled` (#11322)
 - Reduce allocations in `Utility.IsAzureMonitorLoggingEnabled` (#11323)
+- Bug fix that fails in-flight invocations when a worker channel shuts down (#11159)
diff --git a/src/WebJobs.Script.Grpc/Channel/GrpcWorkerChannel.cs b/src/WebJobs.Script.Grpc/Channel/GrpcWorkerChannel.cs
@@ -16,12 +16,14 @@
 using System.Threading.Tasks.Dataflow;
 using Google.Protobuf.Collections;
 using Google.Protobuf.WellKnownTypes;
+using Microsoft.Azure.WebJobs.Host;
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Config;
 using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
 using Microsoft.Azure.WebJobs.Script.Diagnostics.OpenTelemetry;
 using Microsoft.Azure.WebJobs.Script.Eventing;
+using Microsoft.Azure.WebJobs.Script.Exceptions;
 using Microsoft.Azure.WebJobs.Script.Extensions;
 using Microsoft.Azure.WebJobs.Script.Grpc.Eventing;
 using Microsoft.Azure.WebJobs.Script.Grpc.Extensions;
@@ -1552,21 +1554,22 @@ public bool IsExecutingInvocation(string invocationId)
             return _executingInvocations.ContainsKey(invocationId);
         }
 
-        public bool TryFailExecutions(Exception workerException)
+        public void Shutdown(Exception workerException)
         {
-            if (workerException == null)
+            var shutdownException = workerException;
+
+            if (workerException is null || workerException is FunctionTimeoutException)
             {
-                return false;
+                shutdownException = new FunctionTimeoutAbortException(workerException?.Message ?? "Worker channel is shutting down. Aborting function.", workerException);
             }
 
             foreach (var invocation in _executingInvocations?.Values)
             {
                 string invocationId = invocation.Context?.ExecutionContext?.InvocationId.ToString();
                 _workerChannelLogger.LogDebug("Worker '{workerId}' encountered a fatal error. Failing invocation: '{invocationId}'", _workerId, invocationId);
-                invocation.Context?.SetException(workerException);
+                invocation.Context?.SetException(shutdownException);
                 RemoveExecutingInvocation(invocationId);
             }
-            return true;
         }
 
         /// <summary>
diff --git a/src/WebJobs.Script.WebHost/WebScriptHostExceptionHandler.cs b/src/WebJobs.Script.WebHost/WebScriptHostExceptionHandler.cs
@@ -1,4 +1,4 @@
-﻿// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -48,14 +48,16 @@ public async Task OnTimeoutExceptionAsync(ExceptionDispatchInfo exceptionInfo, T
             // 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);
+            // await Task.Delay(timeoutGracePeriod);
             IFunctionInvocationDispatcher functionInvocationDispatcher = _functionInvocationDispatcherFactory.GetFunctionDispatcher();
             if (!functionInvocationDispatcher.State.Equals(FunctionInvocationDispatcherState.Default))
             {
                 _logger.LogWarning($"A function timeout has occurred. Restarting worker process executing invocationId '{timeoutException.InstanceId}'.", exceptionInfo.SourceException);
+
                 // If invocation id is not found in any of the workers => worker is already disposed. No action needed.
-                await functionInvocationDispatcher.RestartWorkerWithInvocationIdAsync(timeoutException.InstanceId.ToString());
-                _logger.LogWarning("Restart of language worker process(es) completed.", exceptionInfo.SourceException);
+                await functionInvocationDispatcher.RestartWorkerWithInvocationIdAsync(timeoutException.InstanceId.ToString(), timeoutException);
+
+                _logger.LogWarning("Attempt to restart language worker process(es) completed.", exceptionInfo.SourceException);
             }
             else
             {
diff --git a/src/WebJobs.Script/Http/RetryProxyHandler.cs b/src/WebJobs.Script/Http/RetryProxyHandler.cs
@@ -1,4 +1,4 @@
-﻿// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -59,7 +59,7 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
                 {
                     if (resultSource is not null && (resultSource.Task.IsFaulted || resultSource.Task.IsCanceled))
                     {
-                        _logger.LogWarning("HTTP request will not be retried. The associated function invocation has failed.");
+                        // if the ScriptInvocation has already failed, we do not need to retry any more because there's no request to proxy.
                         throw;
                     }
 
diff --git a/src/WebJobs.Script/Workers/Http/HttpFunctionInvocationDispatcher.cs b/src/WebJobs.Script/Workers/Http/HttpFunctionInvocationDispatcher.cs
@@ -206,7 +206,7 @@ public Task ShutdownAsync()
             return Task.CompletedTask;
         }
 
-        public Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId)
+        public Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId, Exception exception = null)
         {
             // Since there's only one channel for httpworker
             DisposeAndRestartWorkerChannel(_httpWorkerChannel.Id);
diff --git a/src/WebJobs.Script/Workers/IFunctionInvocationDispatcher.cs b/src/WebJobs.Script/Workers/IFunctionInvocationDispatcher.cs
@@ -23,7 +23,7 @@ public interface IFunctionInvocationDispatcher : IDisposable
 
         Task ShutdownAsync();
 
-        Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId);
+        Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId, Exception exception);
 
         Task StartWorkerChannel();
 
diff --git a/src/WebJobs.Script/Workers/Rpc/FunctionRegistration/RpcFunctionInvocationDispatcher.cs b/src/WebJobs.Script/Workers/Rpc/FunctionRegistration/RpcFunctionInvocationDispatcher.cs
@@ -1,4 +1,4 @@
-﻿// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -676,7 +676,7 @@ public void Dispose()
             Dispose(true);
         }
 
-        public async Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId)
+        public async Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId, Exception exception = null)
         {
             // Dispose and restart errored channel with the particular invocation id
             var channels = await GetInitializedWorkerChannelsAsync();
@@ -685,10 +685,13 @@ public async Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId)
                 if (channel.IsExecutingInvocation(invocationId))
                 {
                     _logger.LogDebug($"Restarting channel with workerId: '{channel.Id}' that is executing invocation: '{invocationId}' and timed out.");
-                    await DisposeAndRestartWorkerChannel(_workerRuntime, channel.Id);
+                    await DisposeAndRestartWorkerChannel(_workerRuntime, channel.Id, exception);
                     return true;
                 }
             }
+
+            _logger.LogDebug("The channel that is executing invocation {invocationId} no longer exists.", invocationId);
+
             return false;
         }
 
diff --git a/src/WebJobs.Script/Workers/Rpc/IRpcWorkerChannel.cs b/src/WebJobs.Script/Workers/Rpc/IRpcWorkerChannel.cs
@@ -32,6 +32,6 @@ public interface IRpcWorkerChannel : IWorkerChannel
 
         bool IsExecutingInvocation(string invocationId);
 
-        bool TryFailExecutions(Exception workerException);
+        void Shutdown(Exception workerException);
     }
 }
diff --git a/src/WebJobs.Script/Workers/Rpc/JobHostRpcWorkerChannelManager.cs b/src/WebJobs.Script/Workers/Rpc/JobHostRpcWorkerChannelManager.cs
@@ -49,7 +49,7 @@ public Task<bool> ShutdownChannelIfExistsAsync(string channelId, Exception worke
                     {
                         string id = rpcChannel.Id;
                         _logger.LogDebug("Disposing language worker channel with id:{workerId}", id);
-                        rpcChannel.TryFailExecutions(workerException);
+                        rpcChannel.Shutdown(workerException);
 
                         (rpcChannel as IDisposable)?.Dispose();
                         _logger.LogDebug("Disposed language worker channel with id:{workerId}", id);
diff --git a/src/WebJobs.Script/Workers/Rpc/WebHostRpcWorkerChannelManager.cs b/src/WebJobs.Script/Workers/Rpc/WebHostRpcWorkerChannelManager.cs
@@ -1,4 +1,4 @@
-﻿// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -269,7 +269,7 @@ public Task<bool> ShutdownChannelIfExistsAsync(string language, string workerId,
                                 if (workerChannel != null)
                                 {
                                     _logger.LogDebug("Disposing WebHost channel for workerId: {channelId}, for runtime:{language}", workerId, language);
-                                    workerChannel.TryFailExecutions(workerException);
+                                    workerChannel.Shutdown(workerException);
                                     (channelTask.Result as IDisposable)?.Dispose();
                                 }
                             }
@@ -295,7 +295,7 @@ public Task<bool> ShutdownChannelIfExistsAsync(string language, string workerId,
                             if (workerChannel != null)
                             {
                                 _logger.LogDebug("Disposing WebHost channel for workerId: {channelId}, for runtime:{language}", workerId, language);
-                                workerChannel.TryFailExecutions(workerException);
+                                workerChannel.Shutdown(workerException);
                                 (channelTask.Result as IDisposable)?.Dispose();
                             }
                         }
diff --git a/test/WebJobs.Script.Tests.Integration/Host/WebScriptHostManagerTimeoutTests.cs b/test/WebJobs.Script.Tests.Integration/Host/WebScriptHostManagerTimeoutTests.cs
@@ -1,4 +1,4 @@
-﻿// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -71,7 +71,7 @@ await Assert.ThrowsAsync<ApplicationException>(() => TestHelpers.Await(() => !(j
                 var messages = host.GetScriptHostLogMessages().Where(t => t?.FormattedMessage != null);
                 Assert.Contains(messages, t => t.FormattedMessage.StartsWith("A function timeout has occurred. Restarting worker process executing invocationId "));
                 Assert.Contains(messages, t => t.FormattedMessage.StartsWith("Restarting channel"));
-                Assert.Contains(messages, t => t.FormattedMessage == "Restart of language worker process(es) completed.");
+                Assert.Contains(messages, t => t.FormattedMessage == "Attempt to restart language worker process(es) completed.");
             }
         }
 
diff --git a/test/WebJobs.Script.Tests/Handlers/WebScriptHostExceptionHandlerTests.cs b/test/WebJobs.Script.Tests/Handlers/WebScriptHostExceptionHandlerTests.cs
@@ -0,0 +1,87 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Runtime.ExceptionServices;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Hosting;
+using Microsoft.Azure.WebJobs.Host;
+using Microsoft.Azure.WebJobs.Script.WebHost;
+using Microsoft.Azure.WebJobs.Script.Workers;
+using Microsoft.Extensions.Logging;
+using Moq;
+using Xunit;
+
+namespace Microsoft.Azure.WebJobs.Script.Tests.Handlers
+{
+    public class WebScriptHostExceptionHandlerTests
+    {
+        private readonly Mock<IApplicationLifetime> _mockApplicationLifetime;
+        private readonly Mock<ILogger<WebScriptHostExceptionHandler>> _mockLogger;
+        private readonly Mock<IFunctionInvocationDispatcherFactory> _mockDispatcherFactory;
+        private readonly Mock<IFunctionInvocationDispatcher> _mockDispatcher;
+        private readonly WebScriptHostExceptionHandler _exceptionHandler;
+
+        public WebScriptHostExceptionHandlerTests()
+        {
+            _mockApplicationLifetime = new Mock<IApplicationLifetime>();
+            _mockLogger = new Mock<ILogger<WebScriptHostExceptionHandler>>();
+            _mockDispatcherFactory = new Mock<IFunctionInvocationDispatcherFactory>();
+            _mockDispatcher = new Mock<IFunctionInvocationDispatcher>();
+
+            _mockDispatcherFactory.Setup(f => f.GetFunctionDispatcher())
+                .Returns(_mockDispatcher.Object);
+
+            _exceptionHandler = new WebScriptHostExceptionHandler(
+                _mockApplicationLifetime.Object,
+                _mockLogger.Object,
+                _mockDispatcherFactory.Object);
+        }
+
+        [Fact]
+        public async Task OnTimeoutExceptionAsync_CallsRestartWorkerWithInvocationIdAsync_WithTimeoutException()
+        {
+            var task = Task.CompletedTask;
+            var timeoutException = new FunctionTimeoutException("Test timeout");
+            var exceptionInfo = ExceptionDispatchInfo.Capture(timeoutException);
+            var timeoutGracePeriod = TimeSpan.FromSeconds(5);
+
+            _mockDispatcher.Setup(d => d.State)
+                .Returns(FunctionInvocationDispatcherState.Initialized);
+            _mockDispatcher.Setup(d => d.RestartWorkerWithInvocationIdAsync(It.IsAny<string>(), It.IsAny<Exception>()))
+                .Returns(Task.FromResult(true));
+
+            await _exceptionHandler.OnTimeoutExceptionAsync(exceptionInfo, timeoutGracePeriod);
+
+            _mockDispatcher.Verify(d => d.RestartWorkerWithInvocationIdAsync(
+                It.IsAny<string>(),
+                timeoutException), Times.Once);
+        }
+
+        [Fact]
+        public async Task OnTimeoutExceptionAsync_WhenTaskDoesNotCompleteWithinGracePeriod_RestartsWorker()
+        {
+            // Arrange
+            var invocationId = Guid.NewGuid();
+            var taskCompletionSource = new TaskCompletionSource<bool>();
+            var timeoutException = new FunctionTimeoutException("Test timeout");
+            var exceptionInfo = ExceptionDispatchInfo.Capture(timeoutException);
+            var timeoutGracePeriod = TimeSpan.FromMilliseconds(100); // Short grace period
+
+            _mockDispatcher.Setup(d => d.State)
+                .Returns(FunctionInvocationDispatcherState.Initialized);
+            _mockDispatcher.Setup(d => d.RestartWorkerWithInvocationIdAsync(It.IsAny<string>(), It.IsAny<Exception>()))
+                .Returns(Task.FromResult(true));
+
+            // Don't complete the task to simulate it not finishing within the grace period
+
+            // Act
+            await _exceptionHandler.OnTimeoutExceptionAsync(exceptionInfo, timeoutGracePeriod);
+
+            // Assert
+            _mockDispatcher.Verify(d => d.RestartWorkerWithInvocationIdAsync(
+                It.IsAny<string>(),
+                timeoutException), Times.Once);
+        }
+    }
+}
diff --git a/test/WebJobs.Script.Tests/Workers/Rpc/GrpcWorkerChannelTests.cs b/test/WebJobs.Script.Tests/Workers/Rpc/GrpcWorkerChannelTests.cs
diff --git a/test/WebJobs.Script.Tests/Workers/Rpc/TestRpcWorkerChannel.cs b/test/WebJobs.Script.Tests/Workers/Rpc/TestRpcWorkerChannel.cs
diff --git a/test/WebJobs.Script.Tests/Workers/Rpc/TestRpcWorkerChannelManager.cs b/test/WebJobs.Script.Tests/Workers/Rpc/TestRpcWorkerChannelManager.cs

Original file line number	Diff line number	Diff line change
`@@ -16,12 +16,14 @@`
`16`	`16`	`using System.Threading.Tasks.Dataflow;`
`17`	`17`	`using Google.Protobuf.Collections;`
`18`	`18`	`using Google.Protobuf.WellKnownTypes;`
	`19`	`+using Microsoft.Azure.WebJobs.Host;`
`19`	`20`	`using Microsoft.Azure.WebJobs.Logging;`
`20`	`21`	`using Microsoft.Azure.WebJobs.Script.Config;`
`21`	`22`	`using Microsoft.Azure.WebJobs.Script.Description;`
`22`	`23`	`using Microsoft.Azure.WebJobs.Script.Diagnostics;`
`23`	`24`	`using Microsoft.Azure.WebJobs.Script.Diagnostics.OpenTelemetry;`
`24`	`25`	`using Microsoft.Azure.WebJobs.Script.Eventing;`
	`26`	`+using Microsoft.Azure.WebJobs.Script.Exceptions;`
`25`	`27`	`using Microsoft.Azure.WebJobs.Script.Extensions;`
`26`	`28`	`using Microsoft.Azure.WebJobs.Script.Grpc.Eventing;`
`27`	`29`	`using Microsoft.Azure.WebJobs.Script.Grpc.Extensions;`
`@@ -1552,21 +1554,22 @@ public bool IsExecutingInvocation(string invocationId)`
`1552`	`1554`	`return _executingInvocations.ContainsKey(invocationId);`
`1553`	`1555`	`}`
`1554`	`1556`
`1555`		`- public bool TryFailExecutions(Exception workerException)`
	`1557`	`+ public void Shutdown(Exception workerException)`
`1556`	`1558`	`{`
`1557`		`- if (workerException == null)`
	`1559`	`+ var shutdownException = workerException;`
	`1560`	`+`
	`1561`	`+ if (workerException is null \|\| workerException is FunctionTimeoutException)`
`1558`	`1562`	`{`
`1559`		`- return false;`
	`1563`	`+ shutdownException = new FunctionTimeoutAbortException(workerException?.Message ?? "Worker channel is shutting down. Aborting function.", workerException);`
`1560`	`1564`	`}`
`1561`	`1565`
`1562`	`1566`	`foreach (var invocation in _executingInvocations?.Values)`
`1563`	`1567`	`{`
`1564`	`1568`	`string invocationId = invocation.Context?.ExecutionContext?.InvocationId.ToString();`
`1565`	`1569`	`_workerChannelLogger.LogDebug("Worker '{workerId}' encountered a fatal error. Failing invocation: '{invocationId}'", _workerId, invocationId);`
`1566`		`- invocation.Context?.SetException(workerException);`
	`1570`	`+ invocation.Context?.SetException(shutdownException);`
`1567`	`1571`	`RemoveExecutingInvocation(invocationId);`
`1568`	`1572`	`}`
`1569`		`- return true;`
`1570`	`1573`	`}`
`1571`	`1574`
`1572`	`1575`	`/// <summary>`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-// Copyright (c) .NET Foundation. All rights reserved.`
	`1`	`+// Copyright (c) .NET Foundation. All rights reserved.`
`2`	`2`	`// Licensed under the MIT License. See License.txt in the project root for license information.`
`3`	`3`
`4`	`4`	`using System;`
`@@ -59,7 +59,7 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage`
`59`	`59`	`{`
`60`	`60`	`if (resultSource is not null && (resultSource.Task.IsFaulted \|\| resultSource.Task.IsCanceled))`
`61`	`61`	`{`
`62`		`- _logger.LogWarning("HTTP request will not be retried. The associated function invocation has failed.");`
	`62`	`+ // if the ScriptInvocation has already failed, we do not need to retry any more because there's no request to proxy.`
`63`	`63`	`throw;`
`64`	`64`	`}`
`65`	`65`
Original file line number	Diff line number	Diff line change
`@@ -206,7 +206,7 @@ public Task ShutdownAsync()`
`206`	`206`	`return Task.CompletedTask;`
`207`	`207`	`}`
`208`	`208`
`209`		`- public Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId)`
	`209`	`+ public Task<bool> RestartWorkerWithInvocationIdAsync(string invocationId, Exception exception = null)`
`210`	`210`	`{`
`211`	`211`	`// Since there's only one channel for httpworker`
`212`	`212`	`DisposeAndRestartWorkerChannel(_httpWorkerChannel.Id);`
Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,6 @@ public interface IRpcWorkerChannel : IWorkerChannel`
`32`	`32`
`33`	`33`	`bool IsExecutingInvocation(string invocationId);`
`34`	`34`
`35`		`- bool TryFailExecutions(Exception workerException);`
	`35`	`+ void Shutdown(Exception workerException);`
`36`	`36`	`}`
`37`	`37`	`}`
Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,7 @@ public Task<bool> ShutdownChannelIfExistsAsync(string channelId, Exception worke`
`49`	`49`	`{`
`50`	`50`	`string id = rpcChannel.Id;`
`51`	`51`	`_logger.LogDebug("Disposing language worker channel with id:{workerId}", id);`
`52`		`- rpcChannel.TryFailExecutions(workerException);`
	`52`	`+ rpcChannel.Shutdown(workerException);`
`53`	`53`
`54`	`54`	`(rpcChannel as IDisposable)?.Dispose();`
`55`	`55`	`_logger.LogDebug("Disposed language worker channel with id:{workerId}", id);`