Ensure invocation buffers are setup before routing invocations to worker channel (#5807)

Author: pragnagopa

src/WebJobs.Script/Workers/Rpc/FunctionRegistration/RpcFunctionInvocationDispatcher.cs

@@ -117,18 +117,18 @@ internal Task InitializeJobhostLanguageWorkerChannelAsync(int attemptCount)
             rpcWorkerChannel.SetupFunctionInvocationBuffers(_functions);
             _jobHostLanguageWorkerChannelManager.AddChannel(rpcWorkerChannel);
             rpcWorkerChannel.StartWorkerProcessAsync().ContinueWith(workerInitTask =>
-                 {
-                     if (workerInitTask.IsCompleted)
-                     {
-                         _logger.LogDebug("Adding jobhost language worker channel for runtime: {language}. workerId:{id}", _workerRuntime, rpcWorkerChannel.Id);
-                         rpcWorkerChannel.SendFunctionLoadRequests(_managedDependencyOptions.Value);
-                         State = FunctionInvocationDispatcherState.Initialized;
-                     }
-                     else
-                     {
-                         _logger.LogWarning("Failed to start language worker process for runtime: {language}. workerId:{id}", _workerRuntime, rpcWorkerChannel.Id);
-                     }
-                 });
+            {
+                if (workerInitTask.IsCompleted)
+                {
+                    _logger.LogDebug("Adding jobhost language worker channel for runtime: {language}. workerId:{id}", _workerRuntime, rpcWorkerChannel.Id);
+                    rpcWorkerChannel.SendFunctionLoadRequests(_managedDependencyOptions.Value);
+                    State = FunctionInvocationDispatcherState.Initialized;
+                }
+                else
+                {
+                    _logger.LogWarning("Failed to start language worker process for runtime: {language}. workerId:{id}", _workerRuntime, rpcWorkerChannel.Id);
+                }
+            });
             return Task.CompletedTask;
         }
 
@@ -264,7 +264,7 @@ internal async Task<IEnumerable<IRpcWorkerChannel>> GetInitializedWorkerChannels
                 }
             }
             IEnumerable<IRpcWorkerChannel> workerChannels = webhostChannels == null ? _jobHostLanguageWorkerChannelManager.GetChannels() : webhostChannels.Union(_jobHostLanguageWorkerChannelManager.GetChannels());
-            IEnumerable<IRpcWorkerChannel> initializedWorkers = workerChannels.Where(ch => ch.State == RpcWorkerChannelState.Initialized);
+            IEnumerable<IRpcWorkerChannel> initializedWorkers = workerChannels.Where(ch => ch.IsChannelReadyForInvocations());
             if (initializedWorkers.Count() > _maxProcessCount)
             {
                 throw new InvalidOperationException($"Number of initialized language workers exceeded:{initializedWorkers.Count()} exceeded maxProcessCount: {_maxProcessCount}");
@@ -400,4 +400,4 @@ public void Dispose()
             Dispose(true);
         }
     }
-}
+}

src/WebJobs.Script/Workers/Rpc/IRpcWorkerChannel.cs

@@ -16,7 +16,7 @@ public interface IRpcWorkerChannel
 
         IDictionary<string, BufferBlock<ScriptInvocationContext>> FunctionInputBuffers { get; }
 
-        RpcWorkerChannelState State { get; }
+        bool IsChannelReadyForInvocations();
 
         void SetupFunctionInvocationBuffers(IEnumerable<FunctionMetadata> functions);
 

src/WebJobs.Script/Workers/Rpc/RpcWorkerChannel.cs

@@ -107,10 +107,13 @@ internal RpcWorkerChannel(
 
         public IDictionary<string, BufferBlock<ScriptInvocationContext>> FunctionInputBuffers => _functionInputBuffers;
 
-        public RpcWorkerChannelState State => _state;
-
         internal IWorkerProcess WorkerProcess => _rpcWorkerProcess;
 
+        public bool IsChannelReadyForInvocations()
+        {
+            return _state.HasFlag(RpcWorkerChannelState.InvocationBuffersInitialized | RpcWorkerChannelState.Initialized);
+        }
+
         public async Task StartWorkerProcessAsync()
         {
             _startSubscription = _inboundWorkerEvents.Where(msg => msg.MessageType == MsgType.StartStream)
@@ -120,7 +123,7 @@ public async Task StartWorkerProcessAsync()
 
             _workerChannelLogger.LogDebug("Initiating Worker Process start up");
             await _rpcWorkerProcess.StartProcessAsync();
-            _state = RpcWorkerChannelState.Initializing;
+            _state = _state | RpcWorkerChannelState.Initializing;
             await _workerInitTask.Task;
         }
 
@@ -167,7 +170,7 @@ internal void WorkerInitResponse(RpcEvent initEvent)
                 _workerInitTask.SetResult(false);
                 return;
             }
-            _state = RpcWorkerChannelState.Initialized;
+            _state = _state | RpcWorkerChannelState.Initialized;
             _workerCapabilities.UpdateCapabilities(_initMessage.Capabilities);
             _workerInitTask.SetResult(true);
         }
@@ -180,6 +183,7 @@ public void SetupFunctionInvocationBuffers(IEnumerable<FunctionMetadata> functio
                 _workerChannelLogger.LogDebug("Setting up FunctionInvocationBuffer for function:{functionName} with functionId:{id}", metadata.Name, metadata.FunctionId);
                 _functionInputBuffers[metadata.FunctionId] = new BufferBlock<ScriptInvocationContext>();
             }
+            _state = _state | RpcWorkerChannelState.InvocationBuffersInitialized;
         }
 
         public void SendFunctionLoadRequests(ManagedDependencyOptions managedDependencyOptions)

src/WebJobs.Script/Workers/Rpc/RpcWorkerChannelState.cs

@@ -1,23 +1,31 @@
 // 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.Workers.Rpc
 {
+    [Flags]
     public enum RpcWorkerChannelState
     {
         /// <summary>
         /// The Default state of LanguageWorkerChannel.
         /// </summary>
-        Default,
+        Default = 0,
+
+        /// <summary>
+        /// LanguageWorkerChannel is created. InvocationBuffers per function are setup
+        /// </summary>
+        InvocationBuffersInitialized = 1,
 
         /// <summary>
         /// The LanguageWorkerChannel is created.Worker process is starting
         /// </summary>
-        Initializing,
+        Initializing = 2,
 
         /// <summary>
         /// LanguageWorkerChannel is created. Worker process is Initialized. Rpc Channel is established.
         /// </summary>
-        Initialized
+        Initialized = 4,
     }
 }

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

@@ -465,7 +465,7 @@ await TestHelpers.Await(() =>
                 currentChannelCount = functionDispatcher.JobHostLanguageWorkerChannelManager.GetChannels().Count();
                 if (currentChannelCount == expectedCount)
                 {
-                    return functionDispatcher.JobHostLanguageWorkerChannelManager.GetChannels().All(ch => ch.State == RpcWorkerChannelState.Initialized);
+                    return functionDispatcher.JobHostLanguageWorkerChannelManager.GetChannels().All(ch => ch.IsChannelReadyForInvocations());
                 }
                 return false;
             }, pollingInterval: expectedCount * 5 * 1000, timeout: 60 * 1000);
@@ -514,4 +514,4 @@ private IEnumerable<FunctionMetadata> GetTestFunctionsList(string runtime)
             };
         }
     }
-}
+}

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

@@ -71,13 +71,24 @@ public RpcWorkerChannelTests()
         }
 
         [Fact]
-        public async Task StartWorkerProcessAsync_Invoked()
+        public async Task StartWorkerProcessAsync_Invoked_SetupFunctionBuffers_Verify_ReadyForInvocation()
         {
             var initTask = _workerChannel.StartWorkerProcessAsync();
             _testFunctionRpcService.PublishStartStreamEvent(_workerId);
             _testFunctionRpcService.PublishWorkerInitResponseEvent();
             await initTask;
             _mockrpcWorkerProcess.Verify(m => m.StartProcessAsync(), Times.Once);
+            Assert.False(_workerChannel.IsChannelReadyForInvocations());
+            _workerChannel.SetupFunctionInvocationBuffers(GetTestFunctionsList("node"));
+            Assert.True(_workerChannel.IsChannelReadyForInvocations());
+        }
+
+        [Fact]
+        public void SetupFunctionBuffers_Verify_ReadyForInvocation_Returns_False()
+        {
+            Assert.False(_workerChannel.IsChannelReadyForInvocations());
+            _workerChannel.SetupFunctionInvocationBuffers(GetTestFunctionsList("node"));
+            Assert.False(_workerChannel.IsChannelReadyForInvocations());
         }
 
         [Fact]

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

@@ -41,14 +41,13 @@ public TestRpcWorkerChannel(string workerId, string runtime = null, IScriptEvent
 
         public List<Task> ExecutionContexts => _executionContexts;
 
-        public RpcWorkerChannelState State => _state;
-
         public void Dispose()
         {
         }
 
         public void SetupFunctionInvocationBuffers(IEnumerable<FunctionMetadata> functions)
         {
+            _state = _state | RpcWorkerChannelState.InvocationBuffersInitialized;
             _testLogger.LogInformation("SetupFunctionInvocationBuffers called");
         }
 
@@ -82,7 +81,7 @@ public async Task StartWorkerProcessAsync()
             {
                 { "test", "testSupported" }
             };
-            _state = RpcWorkerChannelState.Initialized;
+            _state = _state | RpcWorkerChannelState.Initialized;
         }
 
         public void RaiseWorkerError()
@@ -107,5 +106,10 @@ public async Task DrainInvocationsAsync()
             await Task.WhenAll(ExecutionContexts);
             ExecutionContexts.Clear();
         }
+
+        public bool IsChannelReadyForInvocations()
+        {
+            return _state.HasFlag(RpcWorkerChannelState.InvocationBuffersInitialized | RpcWorkerChannelState.Initialized);
+        }
     }
 }