Initialize invoke BufferBlock during function load (#4492)

Author: pragnagopa

src/WebJobs.Script/Description/Rpc/WorkerLanguageInvoker.cs

@@ -75,7 +75,7 @@ protected override async Task<object> InvokeCore(object[] parameters, FunctionIn
 
             ScriptInvocationResult result;
             _logger.LogDebug($"Sending invocation id:{invocationId}");
-            await _functionDispatcher.InvokeAsync(invocationContext);
+            _functionDispatcher.Invoke(invocationContext);
             result = await invocationContext.ResultSource.Task;
 
             await BindOutputsAsync(triggerValue, context.Binder, result);
@@ -87,7 +87,7 @@ private async Task DelayUntilFunctionDispatcherInitialized()
             if (_functionDispatcher != null && _functionDispatcher.State == FunctionDispatcherState.Initializing)
             {
                 _logger.LogDebug($"functionDispatcher state: {_functionDispatcher.State}");
-                await Utility.DelayAsync(ScriptConstants.HostTimeoutSeconds, ScriptConstants.HostPollingIntervalMilliseconds, () =>
+                await Utility.DelayAsync(LanguageWorkerConstants.ProcessStartTimeoutSeconds, 25, () =>
                 {
                     return _functionDispatcher.State != FunctionDispatcherState.Initialized;
                 });

src/WebJobs.Script/Rpc/FunctionRegistration/FunctionDispatcher.cs

@@ -179,35 +179,28 @@ public async Task InitializeAsync(IEnumerable<FunctionMetadata> functions)
             }
         }
 
-        public async Task InvokeAsync(ScriptInvocationContext invocationContext)
+        public void Invoke(ScriptInvocationContext invocationContext)
         {
             try
             {
                 IEnumerable<ILanguageWorkerChannel> workerChannels = GetInitializedWorkerChannels();
                 var languageWorkerChannel = _functionDispatcherLoadBalancer.GetLanguageWorkerChannel(workerChannels, _maxProcessCount);
-                BufferBlock<ScriptInvocationContext> bufferBlock = await GetFunctionInvocationBufferBlock(languageWorkerChannel, invocationContext.FunctionMetadata.FunctionId);
-                _logger.LogDebug("Posting invocation id:{InvocationId} on workerId:{workerChannelId}", invocationContext.ExecutionContext.InvocationId, languageWorkerChannel.Id);
-                languageWorkerChannel.FunctionInputBuffers[invocationContext.FunctionMetadata.FunctionId].Post(invocationContext);
+                if (languageWorkerChannel.FunctionInputBuffers.TryGetValue(invocationContext.FunctionMetadata.FunctionId, out BufferBlock<ScriptInvocationContext> bufferBlock))
+                {
+                    _logger.LogDebug("Posting invocation id:{InvocationId} on workerId:{workerChannelId}", invocationContext.ExecutionContext.InvocationId, languageWorkerChannel.Id);
+                    languageWorkerChannel.FunctionInputBuffers[invocationContext.FunctionMetadata.FunctionId].Post(invocationContext);
+                }
+                else
+                {
+                    throw new InvalidOperationException($"Function:{invocationContext.FunctionMetadata.Name} is not loaded by the language worker: {languageWorkerChannel.Id}");
+                }
             }
             catch (Exception invokeEx)
             {
                 invocationContext.ResultSource.TrySetException(invokeEx);
             }
         }
 
-        private async Task<BufferBlock<ScriptInvocationContext>> GetFunctionInvocationBufferBlock(ILanguageWorkerChannel languageWorkerChannel, string functionId)
-        {
-            BufferBlock<ScriptInvocationContext> bufferBlock = null;
-            if (languageWorkerChannel != null)
-            {
-                await Utility.DelayAsync(ScriptConstants.HostTimeoutSeconds, ScriptConstants.HostPollingIntervalMilliseconds, () =>
-                {
-                    return !languageWorkerChannel.FunctionInputBuffers.TryGetValue(functionId, out bufferBlock);
-                });
-            }
-            return bufferBlock;
-        }
-
         internal IEnumerable<ILanguageWorkerChannel> GetInitializedWorkerChannels()
         {
             IEnumerable<ILanguageWorkerChannel> webhostChannels = _languageWorkerChannelManager.GetChannels(_workerRuntime);

src/WebJobs.Script/Rpc/FunctionRegistration/IFunctionDispatcher.cs

@@ -17,7 +17,7 @@ public interface IFunctionDispatcher : IDisposable
         // Tests if the function metadata is supported by a known language worker
         bool IsSupported(FunctionMetadata metadata, string language);
 
-        Task InvokeAsync(ScriptInvocationContext invocationContext);
+        void Invoke(ScriptInvocationContext invocationContext);
 
         Task InitializeAsync(IEnumerable<FunctionMetadata> functions);
     }

src/WebJobs.Script/Rpc/LanguageWorkerChannel.cs

@@ -28,7 +28,6 @@ namespace Microsoft.Azure.WebJobs.Script.Rpc
 {
     internal class LanguageWorkerChannel : ILanguageWorkerChannel
     {
-        private readonly TimeSpan processStartTimeout = TimeSpan.FromSeconds(40);
         private readonly TimeSpan workerInitTimeout = TimeSpan.FromSeconds(30);
         private readonly string _rootScriptPath;
         private readonly IScriptEventManager _eventManager;
@@ -242,7 +241,7 @@ private void OnErrorDataReceived(object sender, DataReceivedEventArgs e)
         public Task StartWorkerProcessAsync()
         {
             _startSubscription = _inboundWorkerEvents.Where(msg => msg.MessageType == MsgType.StartStream)
-                .Timeout(processStartTimeout)
+                .Timeout(TimeSpan.FromSeconds(LanguageWorkerConstants.ProcessStartTimeoutSeconds))
                 .Take(1)
                 .Subscribe(SendWorkerInitRequest, HandleWorkerError);
 
@@ -353,6 +352,9 @@ public void SendFunctionEnvironmentReloadRequest()
 
         internal void SendFunctionLoadRequest(FunctionMetadata metadata)
         {
+            _workerChannelLogger.LogDebug("Sending FunctionLoadRequest for function:{functionName} with functionId:{id}", metadata.Name, metadata.FunctionId);
+            _functionInputBuffers[metadata.FunctionId] = new BufferBlock<ScriptInvocationContext>();
+
             // send a load request for the registered function
             FunctionLoadRequest request = new FunctionLoadRequest()
             {
@@ -387,9 +389,7 @@ internal void SendFunctionLoadRequest(FunctionMetadata metadata)
 
         internal void LoadResponse(FunctionLoadResponse loadResponse)
         {
-            // associate the invocation input buffer with the function
-            _functionInputBuffers[loadResponse.FunctionId] = new BufferBlock<ScriptInvocationContext>();
-
+            _workerChannelLogger.LogDebug("Received FunctionLoadResponse for functionId:{functionId}", loadResponse.FunctionId);
             if (loadResponse.Result.IsFailure(out Exception ex))
             {
                 //Cache function load errors to replay error messages on invoking failed functions
@@ -403,6 +403,7 @@ internal void LoadResponse(FunctionLoadResponse loadResponse)
 
             // link the invocation inputs to the invoke call
             var invokeBlock = new ActionBlock<ScriptInvocationContext>(ctx => SendInvocationRequest(ctx));
+            // associate the invocation input buffer with the function
             var disposableLink = _functionInputBuffers[loadResponse.FunctionId].LinkTo(invokeBlock);
             _inputLinks.Add(disposableLink);
         }

src/WebJobs.Script/Rpc/LanguageWorkerConstants.cs

@@ -1,10 +1,13 @@
 // 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.Rpc
 {
     public static class LanguageWorkerConstants
     {
+        public const int ProcessStartTimeoutSeconds = 60;
         public const string FunctionWorkerRuntimeSettingName = "FUNCTIONS_WORKER_RUNTIME";
         public const string FunctionsWorkerProcessCountSettingName = "FUNCTIONS_WORKER_PROCESS_COUNT";
         public const string DotNetLanguageWorkerName = "dotnet";