Worker concurrency memory check (#8585)

Author: alrod

release_notes.md

@@ -9,6 +9,7 @@
 - Updated Java Worker Version to [2.4.0](https://github.com/Azure/azure-functions-java-worker/releases/tag/2.4.0)
 - Update PowerShell Worker 7.2 to 4.0.2258 [Release Note](https://github.com/Azure/azure-functions-powershell-worker/releases/tag/v4.0.2258)
 - Adding support to conditionally include empty entries from trigger payload when sending to OOP workers. ([#8499](https://github.com/Azure/azure-functions-host/issues/8499))
+- 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

@@ -137,7 +137,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

@@ -41,6 +41,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();
 
         void WaitForProcessExitInMilliSeconds(int waitTime);

src/WebJobs.Script/Workers/WorkerConcurrencyManager.cs

@@ -3,10 +3,13 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 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;
@@ -26,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;
@@ -125,8 +129,18 @@ internal async void OnTimer(object sender, System.Timers.ElapsedEventArgs e)
 
                 if (NewWorkerIsRequired(workerStatuses, _addWorkerStopwatch.GetElapsedTime()))
                 {
-                    await _functionInvocationDispatcher.StartWorkerChannel();
-                    _addWorkerStopwatch = ValueStopwatch.StartNew();
+                    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 = ValueStopwatch.StartNew();
+                        }
+                    }
                 }
             }
             catch (Exception ex)
@@ -284,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:"));
+            }
+        }
     }
 }