Fixing http throttle enablement (#6548) (#6867)

mathewc · web-flow · commit 5652e9287273 · 2020-11-05T15:59:31.000-08:00
diff --git a/src/WebJobs.Script.WebHost/Controllers/HostController.cs b/src/WebJobs.Script.WebHost/Controllers/HostController.cs
@@ -154,17 +154,18 @@ public async Task<IActionResult> GetScaleStatus([FromBody] ScaleStatusContext co
 
             // TEMP: Once https://github.com/Azure/azure-functions-host/issues/5161 is fixed, we should take
             // FunctionsScaleManager as a parameter.
-            var scaleManager = (scriptHostManager as IServiceProvider)?.GetService<FunctionsScaleManager>();
-            if (scaleManager == null)
+            if (Utility.TryGetHostService(scriptHostManager, out FunctionsScaleManager scaleManager))
+            {
+                var scaleStatus = await scaleManager.GetScaleStatusAsync(context);
+                return new ObjectResult(scaleStatus);
+            }
+            else
             {
                 // This case should never happen. Because this action is marked RequiresRunningHost,
                 // it's only invoked when the host is running, and if it's running, we'll have access
                 // to the FunctionsScaleManager.
                 return StatusCode(StatusCodes.Status503ServiceUnavailable);
             }
-            var scaleStatus = await scaleManager.GetScaleStatusAsync(context);
-
-            return new ObjectResult(scaleStatus);
         }
 
         [HttpPost]
diff --git a/src/WebJobs.Script.WebHost/Middleware/HttpThrottleMiddleware.cs b/src/WebJobs.Script.WebHost/Middleware/HttpThrottleMiddleware.cs
@@ -8,6 +8,7 @@
 using Microsoft.Azure.WebJobs.Extensions.Http;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
 using Microsoft.Azure.WebJobs.Script.Scale;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
 
@@ -28,9 +29,15 @@ public HttpThrottleMiddleware(RequestDelegate next, ILoggerFactory loggerFactory
             _logger = loggerFactory?.CreateLogger("Host." + ScriptConstants.TraceSourceHttpThrottleMiddleware);
         }
 
-        public static bool ShouldEnable(HttpOptions options)
+        public static bool ShouldEnable(IServiceProvider serviceProvider)
         {
-            return HttpRequestQueue.IsEnabled(options) || options.DynamicThrottlesEnabled;
+            var scriptHostManager = serviceProvider?.GetService<IScriptHostManager>();
+            if (scriptHostManager != null && Utility.TryGetHostService(scriptHostManager, out IOptions<HttpOptions> options))
+            {
+                return HttpRequestQueue.IsEnabled(options.Value) || options.Value.DynamicThrottlesEnabled;
+            }
+
+            return false;
         }
 
         public async Task Invoke(HttpContext httpContext, IOptions<HttpOptions> httpOptions, HttpRequestQueue requestQueue, HostPerformanceManager performanceManager, IMetricsLogger metricsLogger)
diff --git a/src/WebJobs.Script.WebHost/WebJobsApplicationBuilderExtension.cs b/src/WebJobs.Script.WebHost/WebJobsApplicationBuilderExtension.cs
@@ -25,7 +25,7 @@ public static IApplicationBuilder UseWebJobsScriptHost(this IApplicationBuilder
             IEnvironment environment = builder.ApplicationServices.GetService<IEnvironment>() ?? SystemEnvironment.Instance;
             IOptionsMonitor<StandbyOptions> standbyOptions = builder.ApplicationServices.GetService<IOptionsMonitor<StandbyOptions>>();
             IOptionsMonitor<HttpBodyControlOptions> httpBodyControlOptions = builder.ApplicationServices.GetService<IOptionsMonitor<HttpBodyControlOptions>>();
-            IOptionsMonitor<HttpOptions> httpOptions = builder.ApplicationServices.GetService<IOptionsMonitor<HttpOptions>>();
+            IServiceProvider serviceProvider = builder.ApplicationServices;
 
             builder.UseMiddleware<SystemTraceMiddleware>();
             builder.UseMiddleware<HostnameFixupMiddleware>();
@@ -71,10 +71,11 @@ public static IApplicationBuilder UseWebJobsScriptHost(this IApplicationBuilder
             {
                 config.UseMiddleware<HomepageMiddleware>();
             });
-            builder.UseWhen(context => HttpThrottleMiddleware.ShouldEnable(httpOptions.CurrentValue) && !context.Request.IsAdminRequest(), config =>
+            builder.UseWhen(context => !context.Request.IsAdminRequest() && HttpThrottleMiddleware.ShouldEnable(serviceProvider), config =>
             {
                 config.UseMiddleware<HttpThrottleMiddleware>();
             });
+
             builder.UseMiddleware<JobHostPipelineMiddleware>();
             builder.UseMiddleware<FunctionInvocationMiddleware>();
 
diff --git a/src/WebJobs.Script/Scale/HostPerformanceManager.cs b/src/WebJobs.Script/Scale/HostPerformanceManager.cs
@@ -234,8 +234,7 @@ internal ApplicationPerformanceCounters GetPerformanceCounters(ILogger logger =
         private IFunctionInvocationDispatcher GetDispatcherAsync()
         {
             var hostManager = _serviceProvider.GetService<IScriptHostManager>();
-            var dispatcherFactory = (hostManager as IServiceProvider)?.GetService<IFunctionInvocationDispatcherFactory>();
-            if (dispatcherFactory != null)
+            if (Utility.TryGetHostService(hostManager, out IFunctionInvocationDispatcherFactory dispatcherFactory))
             {
                 var dispatcher = dispatcherFactory.GetFunctionDispatcher();
                 if (dispatcher.State == FunctionInvocationDispatcherState.Initialized)
diff --git a/src/WebJobs.Script/Utility.cs b/src/WebJobs.Script/Utility.cs
@@ -16,11 +16,11 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
-using Microsoft.Azure.WebJobs.Host;
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Diagnostics.Extensions;
 using Microsoft.Azure.WebJobs.Script.Workers.Rpc;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Converters;
@@ -44,6 +44,22 @@ public static class Utility
 
         public static int ColdStartDelayMS { get; set; } = 5000;
 
+        internal static bool TryGetHostService<TService>(IScriptHostManager scriptHostManager, out TService service) where TService : class
+        {
+            service = null;
+
+            try
+            {
+                service = (scriptHostManager as IServiceProvider)?.GetService<TService>();
+            }
+            catch
+            {
+                // can get exceptions if the host is being disposed
+            }
+
+            return service != null;
+        }
+
         /// <summary>
         /// Walk from the method up to the containing type, looking for an instance
         /// of the specified attribute type, returning it if found.
diff --git a/test/WebJobs.Script.Tests.Integration/Host/StandbyManager/StandbyManagerE2ETestBase.cs b/test/WebJobs.Script.Tests.Integration/Host/StandbyManager/StandbyManagerE2ETestBase.cs
@@ -12,6 +12,7 @@
 using System.Web.Http;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.TestHost;
+using Microsoft.Azure.WebJobs.Extensions.Http;
 using Microsoft.Azure.WebJobs.Host.Executors;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
 using Microsoft.Azure.WebJobs.Script.WebHost;
@@ -98,6 +99,15 @@ protected async Task<IWebHostBuilder> CreateWebHostBuilderAsync(string testDirNa
                 .ConfigureScriptHostLogging(b =>
                 {
                     b.AddProvider(_loggerProvider);
+                })
+                .ConfigureScriptHostServices(s =>
+                {
+                    s.PostConfigure<HttpOptions>(o =>
+                    {
+                        // disabling dynamic throttles since they can cause sporadic failures for
+                        // tests based on CPU limits being hit resulting in 429 responses
+                        o.DynamicThrottlesEnabled = false;
+                    });
                 });
 
             return webHostBuilder;
diff --git a/test/WebJobs.Script.Tests/Middleware/HttpThrottleMiddlewareTests.cs b/test/WebJobs.Script.Tests/Middleware/HttpThrottleMiddlewareTests.cs
@@ -63,6 +63,31 @@ public HttpThrottleMiddlewareTests()
             _requestQueue = new HttpRequestQueue(new OptionsWrapper<HttpOptions>(_httpOptions));
         }
 
+        [Fact]
+        public void ShouldEnable_ReturnsExpectedValue()
+        {
+            IOptions<HttpOptions> optionsWrapper = null;
+            var scriptHostManagerMock = new Mock<IScriptHostManager>(MockBehavior.Strict);
+            var hostServiceProviderMock = scriptHostManagerMock.As<IServiceProvider>();
+            hostServiceProviderMock.Setup(p => p.GetService(typeof(IOptions<HttpOptions>))).Returns(() => optionsWrapper);
+            var rootServiceProvider = new Mock<IServiceProvider>(MockBehavior.Strict);
+            rootServiceProvider.Setup(p => p.GetService(typeof(IScriptHostManager))).Returns(scriptHostManagerMock.Object);
+
+            Assert.False(HttpThrottleMiddleware.ShouldEnable(null));
+            Assert.False(HttpThrottleMiddleware.ShouldEnable(rootServiceProvider.Object));
+
+            var httpOptions = new HttpOptions();
+            optionsWrapper = new OptionsWrapper<HttpOptions>(httpOptions);
+            Assert.False(HttpThrottleMiddleware.ShouldEnable(rootServiceProvider.Object));
+
+            httpOptions.MaxConcurrentRequests = 5;
+            Assert.True(HttpThrottleMiddleware.ShouldEnable(rootServiceProvider.Object));
+
+            httpOptions.MaxConcurrentRequests = -1;
+            httpOptions.DynamicThrottlesEnabled = true;
+            Assert.True(HttpThrottleMiddleware.ShouldEnable(rootServiceProvider.Object));
+        }
+
         [Fact]
         public async Task Invoke_PropagatesExceptions()
         {
@@ -108,10 +133,11 @@ public async Task Invoke_NoThrottle_DispatchesDirectly()
             Assert.Equal(HttpStatusCode.Accepted, (HttpStatusCode)httpContext.Response.StatusCode);
         }
 
-        [Fact]
-        public async Task Invoke_MaxParallelism_RequestsAreThrottled()
+        [Theory]
+        [InlineData(3)]
+        [InlineData(1)]
+        public async Task Invoke_MaxParallelism_RequestsAreThrottled(int maxParallelism)
         {
-            int maxParallelism = 3;
             _httpOptions = new HttpOptions
             {
                 MaxConcurrentRequests = maxParallelism
diff --git a/test/WebJobs.Script.Tests/UtilityTests.cs b/test/WebJobs.Script.Tests/UtilityTests.cs
@@ -33,6 +33,10 @@ public struct TestStruct
         public string Location { get; set; }
     }
 
+    public class TestClass : ITestInterface
+    {
+    }
+
     public class TestPoco
     {
         public string Name { get; set; }
@@ -55,6 +59,42 @@ public class UtilityTests
     {
         private TestLogger _testLogger = new TestLogger("test");
 
+        [Fact]
+        public void TryGetHostService_ReturnsExpectedResult()
+        {
+            ITestInterface test = new TestClass();
+            var scriptHostManagerMock = new Mock<IScriptHostManager>(MockBehavior.Strict);
+            var serviceProviderMock = scriptHostManagerMock.As<IServiceProvider>();
+            serviceProviderMock.Setup(p => p.GetService(typeof(ITestInterface))).Returns(() => test);
+
+            Assert.True(Utility.TryGetHostService(scriptHostManagerMock.Object, out ITestInterface result));
+            Assert.Same(test, result);
+
+            test = null;
+            Assert.False(Utility.TryGetHostService(scriptHostManagerMock.Object, out result));
+            Assert.Null(result);
+        }
+
+        [Fact]
+        public void TryGetHostService_ObjectDisposed_ReturnsFalse()
+        {
+            var scriptHostManagerMock = new Mock<IScriptHostManager>(MockBehavior.Strict);
+            var serviceProviderMock = scriptHostManagerMock.As<IServiceProvider>();
+            serviceProviderMock.Setup(p => p.GetService(typeof(ITestInterface))).Throws(new ObjectDisposedException("Test"));
+
+            Assert.False(Utility.TryGetHostService(scriptHostManagerMock.Object, out ITestInterface result));
+            Assert.Null(result);
+        }
+
+        [Fact]
+        public void TryGetHostService_NotServiceProvider_ReturnsFalse()
+        {
+            var scriptHostManagerMock = new Mock<IScriptHostManager>(MockBehavior.Strict);
+
+            Assert.False(Utility.TryGetHostService(scriptHostManagerMock.Object, out ITestInterface result));
+            Assert.Null(result);
+        }
+
         [Fact]
         public async Task InvokeWithRetriesAsync_Throws_WhenRetryCountExceeded()
         {

Original file line number	Diff line number	Diff line change
`@@ -234,8 +234,7 @@ internal ApplicationPerformanceCounters GetPerformanceCounters(ILogger logger =`
`234`	`234`	`private IFunctionInvocationDispatcher GetDispatcherAsync()`
`235`	`235`	`{`
`236`	`236`	`var hostManager = _serviceProvider.GetService<IScriptHostManager>();`
`237`		`- var dispatcherFactory = (hostManager as IServiceProvider)?.GetService<IFunctionInvocationDispatcherFactory>();`
`238`		`- if (dispatcherFactory != null)`
	`237`	`+ if (Utility.TryGetHostService(hostManager, out IFunctionInvocationDispatcherFactory dispatcherFactory))`
`239`	`238`	`{`
`240`	`239`	`var dispatcher = dispatcherFactory.GetFunctionDispatcher();`
`241`	`240`	`if (dispatcher.State == FunctionInvocationDispatcherState.Initialized)`