Update http proxy to handle client disconnect scenario (#10688)

Azure · Jan 9, 2025 · e9887fc · e9887fc
1 parent c99ef0a
commit e9887fc
Show file tree

Hide file tree

Showing 8 changed files with 105 additions and 1 deletion.
diff --git a/release_notes.md b/release_notes.md
@@ -3,4 +3,7 @@
 <!-- Please add your release notes in the following format:
 - My change description (#PR)
 -->
+
 - Updated Microsoft.Extensions.DependencyModel to 6.0.2 and 8.0.2 for .NET 6 and .NET 8, respectively. (#10661)
+- Update the `DefaultHttpProxyService` to better handle client disconnect scenarios (#10688)
+  - Replaced `InvalidOperationException` with `HttpForwardingException` when there is a ForwarderError
diff --git a/src/WebJobs.Script.Grpc/Exceptions/HttpForwardingException.cs b/src/WebJobs.Script.Grpc/Exceptions/HttpForwardingException.cs
@@ -0,0 +1,22 @@
+// 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.Grpc.Exceptions
+{
+    internal class HttpForwardingException : Exception
+    {
+        public HttpForwardingException()
+        {
+        }
+
+        public HttpForwardingException(string message) : base(message)
+        {
+        }
+
+        public HttpForwardingException(string message, Exception innerException) : base(message, innerException)
+        {
+        }
+    }
+}
diff --git a/src/WebJobs.Script.Grpc/Server/DefaultHttpProxyService.cs b/src/WebJobs.Script.Grpc/Server/DefaultHttpProxyService.cs
@@ -8,6 +8,7 @@
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Azure.WebJobs.Script.Description;
+using Microsoft.Azure.WebJobs.Script.Grpc.Exceptions;
 using Microsoft.Azure.WebJobs.Script.Workers;
 using Microsoft.Extensions.Logging;
 using Yarp.ReverseProxy.Forwarder;
@@ -20,10 +21,12 @@ internal class DefaultHttpProxyService : IHttpProxyService, IDisposable
         private readonly IHttpForwarder _httpForwarder;
         private readonly HttpMessageInvoker _messageInvoker;
         private readonly ForwarderRequestConfig _forwarderRequestConfig;
+        private readonly ILogger<DefaultHttpProxyService> _logger;
 
         public DefaultHttpProxyService(IHttpForwarder httpForwarder, ILogger<DefaultHttpProxyService> logger)
         {
             _httpForwarder = httpForwarder ?? throw new ArgumentNullException(nameof(httpForwarder));
+            _logger = logger ?? throw new ArgumentNullException(nameof(logger));
 
             _handler = new SocketsHttpHandler
             {
@@ -56,13 +59,23 @@ public async Task EnsureSuccessfulForwardingAsync(ScriptInvocationContext contex
 
             if (httpProxyTaskResult is not ForwarderError.None)
             {
+                _logger.LogDebug($"The function invocation failed to proxy the request with ForwarderError: {httpProxyTaskResult}");
+
                 Exception forwarderException = null;
                 if (context.TryGetHttpRequest(out HttpRequest request))
                 {
                     forwarderException = request.HttpContext.GetForwarderErrorFeature()?.Exception;
                 }
 
-                throw new InvalidOperationException($"Failed to proxy request with ForwarderError: {httpProxyTaskResult}", forwarderException);
+                if (request.HttpContext.RequestAborted.IsCancellationRequested)
+                {
+                    // Client disconnected, nothing we can do here but inform
+                    _logger.LogInformation("The client disconnected while the function was processing the request.");
+                }
+                else
+                {
+                    throw new HttpForwardingException($"Failed to proxy request with ForwarderError: {httpProxyTaskResult}", forwarderException);
+                }
             }
         }
 

diff --git a/src/WebJobs.Script.Grpc/Server/RetryProxyHandler.cs b/src/WebJobs.Script.Grpc/Server/RetryProxyHandler.cs
@@ -32,6 +32,11 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
                 {
                     return await base.SendAsync(request, cancellationToken);
                 }
+                catch (TaskCanceledException) when (cancellationToken.IsCancellationRequested)
+                {
+                    _logger.LogDebug("Request was canceled. Stopping retries.");
+                    throw new OperationCanceledException(cancellationToken);
+                }
                 catch (HttpRequestException) when (attemptCount < _maxRetries)
                 {
                     currentDelay *= attemptCount;

diff --git a/src/WebJobs.Script.WebHost/Diagnostics/Extensions/ScriptHostServiceLoggerExtension.cs b/src/WebJobs.Script.WebHost/Diagnostics/Extensions/ScriptHostServiceLoggerExtension.cs
@@ -188,6 +188,12 @@ public static class ScriptHostServiceLoggerExtension
                 new EventId(530, nameof(LogHostInitializationSettings)),
                 "{hostInitializationSettings}");
 
+        private static readonly Action<ILogger, string, Exception> _requestAborted =
+            LoggerMessage.Define<string>(
+                LogLevel.Debug,
+                new EventId(531, nameof(RequestAborted)),
+                "The request was aborted by the client (requestId: '{mS_ActivityId}').");
+
         public static void HostStateChanged(this ILogger logger, ScriptHostState previousHostState, ScriptHostState newHostState)
         {
             var newState = newHostState.ToString();
@@ -205,6 +211,11 @@ public static void ExecutedHttpRequest(this ILogger logger, string mS_ActivityId
             _executedHttpRequest(logger, mS_ActivityId, identities, statusCode, duration, null);
         }
 
+        public static void RequestAborted(this ILogger logger, string mS_ActivityId)
+        {
+            _requestAborted(logger, mS_ActivityId, null);
+        }
+
         public static void ScriptHostServiceInitCanceledByRuntime(this ILogger logger)
         {
             _scriptHostServiceInitCanceledByRuntime(logger, null);

diff --git a/src/WebJobs.Script.WebHost/Middleware/HandleCancellationMiddleware.cs b/src/WebJobs.Script.WebHost/Middleware/HandleCancellationMiddleware.cs
@@ -0,0 +1,48 @@
+// 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.IO;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.WebHost.Diagnostics.Extensions;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Middleware
+{
+    internal class HandleCancellationMiddleware
+    {
+        private readonly ILogger _logger;
+        private readonly RequestDelegate _next;
+
+        public HandleCancellationMiddleware(RequestDelegate next, ILogger<HandleCancellationMiddleware> logger)
+        {
+            _logger = logger;
+            _next = next;
+        }
+
+        public async Task Invoke(HttpContext context)
+        {
+            var requestId = context.Request.HttpContext.Items[ScriptConstants.AzureFunctionsRequestIdKey].ToString();
+            try
+            {
+                await _next.Invoke(context);
+
+                if (context.RequestAborted.IsCancellationRequested && !context.Response.HasStarted)
+                {
+                    _logger.RequestAborted(requestId);
+                    context.Response.StatusCode = StatusCodes.Status499ClientClosedRequest;
+                }
+            }
+            catch (Exception ex) when ((ex is OperationCanceledException || ex is IOException) && context.RequestAborted.IsCancellationRequested)
+            {
+                _logger.RequestAborted(requestId);
+
+                if (!context.Response.HasStarted)
+                {
+                    context.Response.StatusCode = StatusCodes.Status499ClientClosedRequest;
+                }
+            }
+        }
+    }
+}
diff --git a/src/WebJobs.Script.WebHost/Middleware/SystemTraceMiddleware.cs b/src/WebJobs.Script.WebHost/Middleware/SystemTraceMiddleware.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Diagnostics;
+using System.IO;
 using System.Linq;
 using System.Text;
 using System.Threading.Tasks;

diff --git a/src/WebJobs.Script.WebHost/WebJobsApplicationBuilderExtension.cs b/src/WebJobs.Script.WebHost/WebJobsApplicationBuilderExtension.cs
@@ -37,6 +37,7 @@ public static IApplicationBuilder UseWebJobsScriptHost(this IApplicationBuilder
             builder.UseMiddleware<ClrOptimizationMiddleware>();
             builder.UseMiddleware<HttpRequestBodySizeMiddleware>();
             builder.UseMiddleware<SystemTraceMiddleware>();
+            builder.UseMiddleware<HandleCancellationMiddleware>();
             builder.UseMiddleware<HostnameFixupMiddleware>();
             if (environment.IsAnyLinuxConsumption() || environment.IsAnyKubernetesEnvironment())
             {