Fail in-flight invocations when worker channel shuts down

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

@@ -22,6 +22,7 @@
 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;
@@ -1555,21 +1556,21 @@ public bool IsExecutingInvocation(string invocationId)
             return _executingInvocations.ContainsKey(invocationId);
         }
 
-        public bool TryFailExecutions(Exception workerException)
+        public void Shutdown(Exception workerException)
         {
-            if (workerException == null)
-            {
-                return false;
-            }
+            WorkerShutdownException shutdownException = workerException is not null
+                ? new WorkerShutdownException("Worker encountered a fatal error and is shutting down.", workerException)
+                : new WorkerShutdownException("Worker encountered a fatal error and is shutting down.");
 
             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?.ResultSource?.TrySetException(workerException);
+
+                invocation.Context?.ResultSource?.TrySetException(shutdownException);
+
                 RemoveExecutingInvocation(invocationId);
             }
-            return true;
         }
 
         /// <summary>

src/WebJobs.Script.WebHost/WebScriptHostExceptionHandler.cs

@@ -29,6 +29,13 @@ public async Task OnTimeoutExceptionAsync(ExceptionDispatchInfo exceptionInfo, T
         {
             FunctionTimeoutException timeoutException = exceptionInfo.SourceException as FunctionTimeoutException;
 
+            // this seems to happen when the worker channel is already shutting down. Ex. One timeout is being handled and another comes in during shutdown.
+            if (timeoutException is null)
+            {
+                _logger.LogDebug("A timeout exception has occurred, but worker channel is already shutting down", exceptionInfo.SourceException);
+                return;
+            }
+
             if (timeoutException?.Task != null)
             {
                 // We may double the timeoutGracePeriod here by first waiting to see if the initial
@@ -48,13 +55,13 @@ 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());
+                await functionInvocationDispatcher.RestartWorkerWithInvocationIdAsync(timeoutException.InstanceId.ToString(), timeoutException);
                 _logger.LogWarning("Restart of language worker process(es) completed.", exceptionInfo.SourceException);
             }
             else

src/WebJobs.Script/Exceptions/WorkerShutdownException.cs

@@ -0,0 +1,21 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+
+namespace Microsoft.Azure.WebJobs.Script.Exceptions
+{
+    internal sealed class WorkerShutdownException : Exception
+    {
+        public WorkerShutdownException() { }
+
+        public WorkerShutdownException(string message) : base(message) { }
+
+        public WorkerShutdownException(string message, Exception innerException) : base(message, innerException)
+        {
+            Reason = innerException?.Message ?? string.Empty;
+        }
+
+        public string Reason { get; set; }
+    }
+}

src/WebJobs.Script/Http/DefaultHttpProxyService.cs

@@ -21,6 +21,7 @@ internal class DefaultHttpProxyService : IHttpProxyService, IDisposable
         private readonly HttpMessageInvoker _messageInvoker;
         private readonly ForwarderRequestConfig _forwarderRequestConfig;
         private readonly ILogger<DefaultHttpProxyService> _logger;
+        private readonly HttpTransformer _httpTransformer;
 
         public DefaultHttpProxyService(IHttpForwarder httpForwarder, ILogger<DefaultHttpProxyService> logger)
         {
@@ -39,6 +40,8 @@ public DefaultHttpProxyService(IHttpForwarder httpForwarder, ILogger<DefaultHttp
             {
                 ActivityTimeout = TimeSpan.FromSeconds(240)
             };
+
+            _httpTransformer = new ScriptInvocationRequestTransformer();
         }
 
         public void Dispose()
@@ -98,7 +101,11 @@ public void StartForwarding(ScriptInvocationContext context, Uri httpUri)
             // add invocation id as correlation id, override existing header if present
             httpRequest.Headers[ScriptConstants.HttpProxyCorrelationHeader] = context.ExecutionContext.InvocationId.ToString();
 
-            var forwardingTask = _httpForwarder.SendAsync(httpContext, httpUri.ToString(), _messageInvoker, _forwarderRequestConfig).AsTask();
+            // Add the script invocation context for later observation of the ScriptInvocationResult task.
+            // This helps track failures/cancellations that should halt retrying the http request.
+            httpContext.Items[ScriptConstants.HttpProxyScriptInvocationContext] = context;
+
+            var forwardingTask = _httpForwarder.SendAsync(httpContext, httpUri.ToString(), _messageInvoker, _forwarderRequestConfig, _httpTransformer).AsTask();
             context.Properties[ScriptConstants.HttpProxyTask] = forwardingTask;
         }
     }

src/WebJobs.Script/Http/RetryProxyHandler.cs

@@ -5,6 +5,8 @@
 using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Script.Description;
+using Microsoft.Azure.WebJobs.Script.Exceptions;
 using Microsoft.Extensions.Logging;
 
 namespace Microsoft.Azure.WebJobs.Script.Http
@@ -30,11 +32,22 @@ public RetryProxyHandler(HttpMessageHandler innerHandler, ILogger logger)
 
         protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
         {
+            TaskCompletionSource<ScriptInvocationResult> resultSource = null;
+            if (request.Options.TryGetValue(ScriptConstants.HttpProxyScriptInvocationContext, out ScriptInvocationContext scriptInvocationContext))
+            {
+                resultSource = scriptInvocationContext.ResultSource;
+            }
+
             var currentDelay = InitialDelay;
             for (int attemptCount = 1; attemptCount <= MaxRetries; attemptCount++)
             {
                 try
                 {
+                    if (resultSource is not null && resultSource.Task.IsFaulted)
+                    {
+                        throw resultSource.Task.Exception?.InnerException ?? new InvalidOperationException("The function invocation tied to this HTTP request failed.");
+                    }
+
                     return await base.SendAsync(request, cancellationToken);
                 }
                 catch (TaskCanceledException) when (cancellationToken.IsCancellationRequested)
@@ -51,6 +64,11 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
 
                     currentDelay = Math.Min(currentDelay * 2, MaximumDelay);
                 }
+                catch (WorkerShutdownException)
+                {
+                    _logger.LogDebug("Language worker channel is shutting down. Request will not be retried.");
+                    throw;
+                }
                 catch (Exception ex)
                 {
                     var message = attemptCount == MaxRetries

src/WebJobs.Script/Http/ScriptInvocationRequestTransformer.cs

@@ -0,0 +1,28 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.Description;
+using Yarp.ReverseProxy.Forwarder;
+
+namespace Microsoft.Azure.WebJobs.Script.Http
+{
+    internal class ScriptInvocationRequestTransformer : HttpTransformer
+    {
+        public override async ValueTask TransformRequestAsync(HttpContext httpContext, HttpRequestMessage proxyRequest, string destinationPrefix, CancellationToken cancellationToken)
+        {
+            // this preserves previous behavior (which called the default transformer) - base method is also called inside of here
+            await HttpTransformer.Default.TransformRequestAsync(httpContext, proxyRequest, destinationPrefix, cancellationToken);
+
+            if (httpContext.Items.TryGetValue(ScriptConstants.HttpProxyScriptInvocationContext, out object result)
+                && result is ScriptInvocationContext scriptContext)
+            {
+                proxyRequest.Options.TryAdd(ScriptConstants.HttpProxyScriptInvocationContext, scriptContext);
+            }
+        }
+    }
+}

src/WebJobs.Script/ScriptConstants.cs

@@ -261,6 +261,7 @@ public static class ScriptConstants
         public static readonly string HttpProxyingEnabled = "HttpProxyingEnabled";
         public static readonly string HttpProxyCorrelationHeader = "x-ms-invocation-id";
         public static readonly string HttpProxyTask = "HttpProxyTask";
+        public static readonly string HttpProxyScriptInvocationContext = "HttpProxyScriptInvocationContext";
 
         public static readonly string OperationNameKey = "OperationName";
 

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);

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();
 

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

@@ -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,7 +685,7 @@ 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;
                 }
             }

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);
     }
 }

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);

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

@@ -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();
                             }
                         }

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);
+        }
+    }
+}

test/WebJobs.Script.Tests/HttpProxyService/DefaultHttpProxyServiceTests.cs

@@ -12,7 +12,7 @@
 using Xunit;
 using Yarp.ReverseProxy.Forwarder;
 
-namespace Microsoft.Azure.WebJobs.Script.Tests
+namespace Microsoft.Azure.WebJobs.Script.Tests.Http
 {
     public class DefaultHttpProxyServiceTests
     {