Best effort to avoid GC during cold start specialization (#8521)

Author: safihamid

src/WebJobs.Script.WebHost/Middleware/ClrOptimizationMiddleware.cs

@@ -0,0 +1,93 @@
+// 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;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Middleware
+{
+    /// <summary>
+    /// A middleware responsible for Optimizing CLR settings like GC to help with cold start
+    /// </summary>
+    internal class ClrOptimizationMiddleware
+    {
+        // This is double the amount of memory allocated during cold start specialization.
+        // This value is calculated based on prod profiles across all languages observed for an extended period of time.
+        // This value is just a best effort and if for any reason CLR needs to allocate more memory then it will ignore this value.
+        private const long AllocationBudgetForGCDuringSpecialization = 16 * 1024 * 1024;
+        private readonly ILogger _logger;
+        private readonly RequestDelegate _next;
+        private readonly IScriptWebHostEnvironment _webHostEnvironment;
+        private RequestDelegate _invoke;
+        private double _specialized = 0;
+
+        public ClrOptimizationMiddleware(RequestDelegate next, IScriptWebHostEnvironment webHostEnvironment, ILogger<SystemTraceMiddleware> logger)
+        {
+            _webHostEnvironment = webHostEnvironment;
+            _logger = logger;
+            _next = next;
+            _invoke = InvokeClrOptimizationCheck;
+        }
+
+        public Task Invoke(HttpContext context)
+        {
+            return _invoke(context);
+        }
+
+        private Task InvokeClrOptimizationCheck(HttpContext context)
+        {
+            var task = _next.Invoke(context).ContinueWith(task =>
+            {
+                // We are tweaking GC behavior in ClrOptimizationMiddleware as this is one of the last call stacks that get executed during standby mode as well as function exection.
+                // We force a GC and enter no GC region in standby mode and exit no GC region after first function execution during specialization.
+                StartStopGCAsBestEffort();
+            }, TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.OnlyOnRanToCompletion);
+
+            return task;
+        }
+
+        private void StartStopGCAsBestEffort()
+        {
+            try
+            {
+                if (_webHostEnvironment.InStandbyMode)
+                {
+                    // This is just to make sure we do not enter NoGCRegion multiple times during standby mode.
+                    if (GCSettings.LatencyMode != GCLatencyMode.NoGCRegion)
+                    {
+                        // In standby mode, we enforce a GC then enter NoGCRegion mode as best effort.
+                        // This is to try to avoid GC during cold start specialization.
+                        GC.Collect();
+                        if (!GC.TryStartNoGCRegion(AllocationBudgetForGCDuringSpecialization, disallowFullBlockingGC: false))
+                        {
+                            _logger.LogError($"CLR runtime failed to commit the requested amount of memory: {AllocationBudgetForGCDuringSpecialization}");
+                        }
+                        _logger.LogInformation($"Collection count for gen 0: {GC.CollectionCount(0)}, gen 1: {GC.CollectionCount(1)}, gen 2: {GC.CollectionCount(2)}");
+                    }
+                }
+                else
+                {
+                    // if not in standby mode and we are in NoGCRegion then we end NoGCRegion.
+                    if (GCSettings.LatencyMode == GCLatencyMode.NoGCRegion)
+                    {
+                        GC.EndNoGCRegion();
+                        _logger.LogInformation($"Collection count for gen 0: {GC.CollectionCount(0)}, gen 1: {GC.CollectionCount(1)}, gen 2: {GC.CollectionCount(2)}");
+                    }
+
+                    // This logic needs to run only once during specialization, so replacing the RequestDelegate after specialization
+                    if (Interlocked.CompareExchange(ref _specialized, 1, 0) == 0)
+                    {
+                        Interlocked.Exchange(ref _invoke, _next);
+                    }
+                }
+            }
+            catch (Exception ex)
+            {
+                _logger.LogError(ex.Message);
+            }
+        }
+    }
+}

src/WebJobs.Script.WebHost/WebJobsApplicationBuilderExtension.cs

@@ -33,6 +33,8 @@ public static IApplicationBuilder UseWebJobsScriptHost(this IApplicationBuilder
             HttpBodyControlOptions httpBodyControlOptions = httpBodyControlOptionsMonitor.CurrentValue;
             httpBodyControlOptionsMonitor.OnChange(newOptions => httpBodyControlOptions = newOptions);
 
+            // Ensure the ClrOptimizationMiddleware is registered before all middleware
+            builder.UseMiddleware<ClrOptimizationMiddleware>();
             builder.UseMiddleware<HttpRequestBodySizeMiddleware>();
             builder.UseMiddleware<SystemTraceMiddleware>();
             builder.UseMiddleware<HostnameFixupMiddleware>();

src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs

@@ -159,7 +159,7 @@ internal bool ShouldMonitorHostHealth
         {
             get
             {
-                return _healthMonitorOptions.Value.Enabled && _environment.IsAppService();
+                return _healthMonitorOptions.Value.Enabled && _environment.IsAppService() && !_scriptWebHostEnvironment.InStandbyMode;
             }
         }
 

test/WebJobs.Script.Tests.Integration/WebHostEndToEnd/SpecializationE2ETests.cs

@@ -7,6 +7,7 @@
 using System.Linq;
 using System.Net;
 using System.Net.Http;
+using System.Runtime;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Storage.Blobs;
@@ -243,6 +244,31 @@ public async Task Specialization_ResetsSharedLoadContext()
             }
         }
 
+        [Fact]
+        public async Task Specialization_GCMode()
+        {
+            var builder = CreateStandbyHostBuilder("FunctionExecutionContext");
+
+            using (var testServer = new TestServer(builder))
+            {
+                var client = testServer.CreateClient();
+
+                // GC's LatencyMode should be Interactive as default, switch to NoGCRegion in placeholder mode and back to Interactive when specialization is complete.
+                Assert.True(GCSettings.LatencyMode != GCLatencyMode.NoGCRegion, "GCLatencyMode should *not* be NoGCRegion at the beginning");
+
+                var response = await client.GetAsync("api/warmup");
+                response.EnsureSuccessStatusCode();
+
+                _environment.SetEnvironmentVariable(EnvironmentSettingNames.AzureWebsiteContainerReady, "1");
+                _environment.SetEnvironmentVariable(EnvironmentSettingNames.AzureWebsitePlaceholderMode, "0");
+
+                response = await client.GetAsync("api/functionexecutioncontext");
+                response.EnsureSuccessStatusCode();
+
+                Assert.True(GCSettings.LatencyMode != GCLatencyMode.NoGCRegion, "GCLatencyMode should *not* be NoGCRegion at the end of specialization");
+            }
+        }
+
         [Fact]
         public async Task Specialization_ResetsSecretManagerRepository()
         {