[v3 port] Worker concurrency memory check

Author: alrod

release_notes.md

@@ -2,6 +2,7 @@
 <!-- Please add your release notes in the following format:
 - My change description (#PR)
 -->
+- Worker concurrency memory check (https://github.com/Azure/azure-functions-host/issues/7499)
 
 **Release sprint:** Sprint 128
 [ [bugs](https://github.com/Azure/azure-functions-host/issues?q=is%3Aissue+milestone%3A%22Functions+Sprint+128%22+label%3Abug+is%3Aclosed) | [features](https://github.com/Azure/azure-functions-host/issues?q=is%3Aissue+milestone%3A%22Functions+Sprint+128%22+label%3Afeature+is%3Aclosed) ]

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

@@ -134,7 +134,7 @@ internal GrpcWorkerChannel(
 
         public IDictionary<string, BufferBlock<ScriptInvocationContext>> FunctionInputBuffers => _functionInputBuffers;
 
-        internal IWorkerProcess WorkerProcess => _rpcWorkerProcess;
+        public IWorkerProcess WorkerProcess => _rpcWorkerProcess;
 
         internal RpcWorkerConfig Config => _workerConfig;
 

src/WebJobs.Script/Workers/Http/HttpWorkerChannel.cs

@@ -42,6 +42,8 @@ internal HttpWorkerChannel(
 
         public string Id { get; }
 
+        public IWorkerProcess WorkerProcess => _workerProcess;
+
         public Task InvokeAsync(ScriptInvocationContext context)
         {
             return _httpWorkerService.InvokeAsync(context);

src/WebJobs.Script/Workers/IWorkerChannel.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -10,6 +11,8 @@ public interface IWorkerChannel
     {
         string Id { get; }
 
+        IWorkerProcess WorkerProcess { get; }
+
         Task<WorkerStatus> GetWorkerStatusAsync();
 
         Task StartWorkerProcessAsync(CancellationToken cancellationToken = default);

src/WebJobs.Script/Workers/ProcessManagement/IWorkerProcess.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+using System.Diagnostics;
 using System.Threading.Tasks;
 using Microsoft.Azure.WebJobs.Script.Scale;
 
@@ -10,6 +11,8 @@ public interface IWorkerProcess
     {
         int Id { get; }
 
+        Process Process { get; }
+
         Task StartProcessAsync();
     }
 }

src/WebJobs.Script/Workers/ProcessManagement/WorkerProcess.cs

@@ -53,8 +53,7 @@ internal WorkerProcess(IScriptEventManager eventManager, IProcessRegistry proces
 
         internal Queue<string> ProcessStdErrDataQueue => _processStdErrDataQueue;
 
-        // for testing
-        internal Process Process { get; set; }
+        public Process Process { get; set; }
 
         internal abstract Process CreateWorkerProcess();
 

src/WebJobs.Script/Workers/WorkerConcurrencyManager.cs

@@ -8,6 +8,8 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Host;
+using Microsoft.Azure.WebJobs.Host.Scale;
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Script.Config;
 using Microsoft.Azure.WebJobs.Script.Diagnostics;
@@ -27,6 +29,7 @@ internal class WorkerConcurrencyManager : IHostedService, IDisposable
         private readonly IEnvironment _environment;
         private readonly IFunctionsHostingConfiguration _functionsHostingConfigurations;
         private readonly IApplicationLifetime _applicationLifetime;
+        private readonly long _memoryLimit = AppServicesHostingUtility.GetMemoryLimitBytes();
 
         private IOptions<WorkerConcurrencyOptions> _workerConcurrencyOptions;
         private IFunctionInvocationDispatcher _functionInvocationDispatcher;
@@ -126,8 +129,18 @@ internal async void OnTimer(object sender, System.Timers.ElapsedEventArgs e)
 
                 if (NewWorkerIsRequired(workerStatuses, _addWorkerStopwatch.Elapsed))
                 {
-                    await _functionInvocationDispatcher.StartWorkerChannel();
-                    _addWorkerStopwatch.Restart();
+                    if (_functionInvocationDispatcher is RpcFunctionInvocationDispatcher rpcDispatcher)
+                    {
+                        // Check memory
+                        var currentChannels = (await rpcDispatcher.GetAllWorkerChannelsAsync()).Select(x => x.WorkerProcess.Process.PrivateMemorySize64);
+                        if (IsEnoughMemoryToScale(Process.GetCurrentProcess().PrivateMemorySize64,
+                            currentChannels,
+                            _memoryLimit))
+                        {
+                            await _functionInvocationDispatcher.StartWorkerChannel();
+                            _addWorkerStopwatch.Restart();
+                        }
+                    }
                 }
             }
             catch (Exception ex)
@@ -285,6 +298,26 @@ public void Dispose()
             Dispose(true);
         }
 
+        internal bool IsEnoughMemoryToScale(long hostProcessSize, IEnumerable<long> workerChannelSizes, long memoryLimit)
+        {
+            if (memoryLimit <= 0)
+            {
+                return true;
+            }
+
+            // Checking memory before adding a new worker
+            // By adding `maxWorkerSize` to current memeory consumption we are predicting what will be overall memory consumption after adding a new worker.
+            // We do not want this value to be more then 80%.
+            long maxWorkerSize = workerChannelSizes.Max();
+            long currentMemoryConsumption = workerChannelSizes.Sum() + hostProcessSize;
+            if (currentMemoryConsumption + maxWorkerSize > memoryLimit * 0.8)
+            {
+                _logger.LogDebug($"Starting new language worker canceled: TotalMemory={memoryLimit}, MaxWorkerSize={maxWorkerSize}, CurrentMemoryConsumption={currentMemoryConsumption}");
+                return false;
+            }
+            return true;
+        }
+
         internal class WorkerStatusDetails
         {
             public string WorkerId { get; set; }

test/WebJobs.Script.Tests/Workers/Rpc/TestRpcWorkerChannel.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Dataflow;
@@ -47,6 +48,8 @@ public TestRpcWorkerChannel(string workerId, string runtime = null, IScriptEvent
 
         public bool IsDisposed => _isDisposed;
 
+        public IWorkerProcess WorkerProcess => null;
+
         public IDictionary<string, BufferBlock<ScriptInvocationContext>> FunctionInputBuffers => throw new NotImplementedException();
 
         public List<Task> ExecutionContexts => _executionContexts;

test/WebJobs.Script.Tests/Workers/Rpc/WorkerConcurrencyManagerTest.cs

@@ -3,9 +3,11 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Host.Scale;
 using Microsoft.Azure.WebJobs.Script.Config;
 using Microsoft.Azure.WebJobs.Script.Workers;
 using Microsoft.Azure.WebJobs.Script.Workers.Rpc;
@@ -354,12 +356,35 @@ public async Task ActivateWorkerConcurency_FunctionsHostingConfiguration_WorkAsE
             conf.Setup(x => x.FunctionsWorkerDynamicConcurrencyEnabled).Returns(true);
             WorkerConcurrencyOptions options = new WorkerConcurrencyOptions();
 
-            WorkerConcurrencyManager concurrancyManager = new WorkerConcurrencyManager(functionInvocationDispatcherFactory.Object, testEnvironment, Options.Create(options), conf.Object, _applicationLifetime, _loggerFactory);
+            WorkerConcurrencyManager concurrancyManager = new WorkerConcurrencyManager(functionInvocationDispatcherFactory.Object, testEnvironment, Options.Create(options), conf.Object,
+                _applicationLifetime, _loggerFactory);
             concurrancyManager.ActivationTimerInterval = TimeSpan.FromMilliseconds(100);
             await concurrancyManager.StartAsync(CancellationToken.None);
             await TestHelpers.Await(() => _loggerProvider.GetAllLogMessages().SingleOrDefault(x => x.FormattedMessage.StartsWith("Dynamic worker concurrency monitoring was started by activation timer.")) != null, timeout: 1000, pollingInterval: 100);
             conf.Setup(x => x.FunctionsWorkerDynamicConcurrencyEnabled).Returns(false);
             await TestHelpers.Await(() => _loggerProvider.GetAllLogMessages().SingleOrDefault(x => x.FormattedMessage.StartsWith("Dynamic worker concurrency monitoring is disabled after activation. Shutting down Functions Host.")) != null, timeout: 1000, pollingInterval: 100);
         }
+
+        [Theory]
+        [InlineData(100, 20, new long[] { 20, 20 }, true)]
+        [InlineData(100, 20, new long[] { 20, 10, 10 }, true)]
+        [InlineData(100, 21, new long[] { 20, 10, 10 }, false)]
+        public void IsEnoughMemory_WorkAsExpected(long availableMemory, long hostProcessSize, IEnumerable<long> languageWorkerSizes, bool result)
+        {
+            TestEnvironment testEnvironment = new TestEnvironment();
+            Mock<IFunctionInvocationDispatcher> functionInvocationDispatcher = new Mock<IFunctionInvocationDispatcher>(MockBehavior.Strict);
+            Mock<IFunctionInvocationDispatcherFactory> functionInvocationDispatcherFactory = new Mock<IFunctionInvocationDispatcherFactory>(MockBehavior.Strict);
+            Mock<IFunctionsHostingConfiguration> conf = new Mock<IFunctionsHostingConfiguration>();
+            WorkerConcurrencyOptions options = new WorkerConcurrencyOptions();
+
+            WorkerConcurrencyManager concurrancyManager = new WorkerConcurrencyManager(functionInvocationDispatcherFactory.Object, testEnvironment,
+                Options.Create(options), conf.Object, _applicationLifetime, _loggerFactory);
+
+            Assert.True(concurrancyManager.IsEnoughMemoryToScale(hostProcessSize, languageWorkerSizes, availableMemory) == result);
+            if (!result)
+            {
+                Assert.Contains(_loggerProvider.GetAllLogMessages().Select(x => x.FormattedMessage), x => x.StartsWith("Starting new language worker canceled:"));
+            }
+        }
     }
 }