Fix race condition during language worker specialization (#4794)

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}");
-            _functionDispatcher.Invoke(invocationContext);
+            await _functionDispatcher.InvokeAsync(invocationContext);
             result = await invocationContext.ResultSource.Task;
 
             await BindOutputsAsync(triggerValue, context.Binder, result);

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

@@ -132,10 +132,10 @@ internal async void InitializeWebhostLanguageWorkerChannel()
             workerChannel.SendFunctionLoadRequests();
         }
 
-        internal void ShutdownWebhostLanguageWorkerChannels()
+        internal async void ShutdownWebhostLanguageWorkerChannels()
         {
             _logger.LogDebug("{workerRuntimeConstant}={value}. Will shutdown all the worker channels that started in placeholder mode", LanguageWorkerConstants.FunctionWorkerRuntimeSettingName, _workerRuntime);
-            _webHostLanguageWorkerChannelManager.ShutdownChannels();
+            await _webHostLanguageWorkerChannelManager?.ShutdownChannelsAsync();
         }
 
         private void StartWorkerProcesses(int startIndex, Action startAction)
@@ -187,16 +187,20 @@ public async Task InitializeAsync(IEnumerable<FunctionMetadata> functions)
             if (Utility.IsSupportedRuntime(_workerRuntime, _workerConfigs))
             {
                 State = FunctionDispatcherState.Initializing;
-                IEnumerable<ILanguageWorkerChannel> initializedChannels = _webHostLanguageWorkerChannelManager.GetChannels(_workerRuntime);
-                if (initializedChannels != null)
+                Dictionary<string, TaskCompletionSource<ILanguageWorkerChannel>> webhostLanguageWorkerChannels = _webHostLanguageWorkerChannelManager.GetChannels(_workerRuntime);
+                if (webhostLanguageWorkerChannels != null)
                 {
-                    foreach (var initializedChannel in initializedChannels)
+                    foreach (string workerId in webhostLanguageWorkerChannels.Keys)
                     {
-                        _logger.LogDebug("Found initialized language worker channel for runtime: {workerRuntime} workerId:{workerId}", _workerRuntime, initializedChannel.Id);
-                        initializedChannel.SetupFunctionInvocationBuffers(_functions);
-                        initializedChannel.SendFunctionLoadRequests();
+                        if (webhostLanguageWorkerChannels.TryGetValue(workerId, out TaskCompletionSource<ILanguageWorkerChannel> initializedLanguageWorkerChannelTask))
+                        {
+                            _logger.LogDebug("Found initialized language worker channel for runtime: {workerRuntime} workerId:{workerId}", _workerRuntime, workerId);
+                            ILanguageWorkerChannel initializedLanguageWorkerChannel = await initializedLanguageWorkerChannelTask.Task;
+                            initializedLanguageWorkerChannel.SetupFunctionInvocationBuffers(_functions);
+                            initializedLanguageWorkerChannel.SendFunctionLoadRequests();
+                        }
                     }
-                    StartWorkerProcesses(initializedChannels.Count(), InitializeWebhostLanguageWorkerChannel);
+                    StartWorkerProcesses(webhostLanguageWorkerChannels.Count(), InitializeWebhostLanguageWorkerChannel);
                     State = FunctionDispatcherState.Initialized;
                 }
                 else
@@ -207,11 +211,11 @@ public async Task InitializeAsync(IEnumerable<FunctionMetadata> functions)
             }
         }
 
-        public void Invoke(ScriptInvocationContext invocationContext)
+        public async Task InvokeAsync(ScriptInvocationContext invocationContext)
         {
             try
             {
-                IEnumerable<ILanguageWorkerChannel> workerChannels = GetInitializedWorkerChannels();
+                IEnumerable<ILanguageWorkerChannel> workerChannels = await GetInitializedWorkerChannelsAsync();
                 var languageWorkerChannel = _functionDispatcherLoadBalancer.GetLanguageWorkerChannel(workerChannels, _maxProcessCount);
                 if (languageWorkerChannel.FunctionInputBuffers.TryGetValue(invocationContext.FunctionMetadata.FunctionId, out BufferBlock<ScriptInvocationContext> bufferBlock))
                 {
@@ -229,9 +233,21 @@ public void Invoke(ScriptInvocationContext invocationContext)
             }
         }
 
-        internal IEnumerable<ILanguageWorkerChannel> GetInitializedWorkerChannels()
+        internal async Task<IEnumerable<ILanguageWorkerChannel>> GetInitializedWorkerChannelsAsync()
         {
-            IEnumerable<ILanguageWorkerChannel> webhostChannels = _webHostLanguageWorkerChannelManager.GetChannels(_workerRuntime);
+            Dictionary<string, TaskCompletionSource<ILanguageWorkerChannel>> webhostChannelDictionary = _webHostLanguageWorkerChannelManager.GetChannels(_workerRuntime);
+            List<ILanguageWorkerChannel> webhostChannels = null;
+            if (webhostChannelDictionary != null)
+            {
+                webhostChannels = new List<ILanguageWorkerChannel>();
+                foreach (string workerId in webhostChannelDictionary.Keys)
+                {
+                    if (webhostChannelDictionary.TryGetValue(workerId, out TaskCompletionSource<ILanguageWorkerChannel> initializedLanguageWorkerChannelTask))
+                    {
+                        webhostChannels.Add(await initializedLanguageWorkerChannelTask.Task);
+                    }
+                }
+            }
             IEnumerable<ILanguageWorkerChannel> workerChannels = webhostChannels == null ? _jobHostLanguageWorkerChannelManager.GetChannels() : webhostChannels.Union(_jobHostLanguageWorkerChannelManager.GetChannels());
             IEnumerable<ILanguageWorkerChannel> initializedWorkers = workerChannels.Where(ch => ch.State == LanguageWorkerChannelState.Initialized);
             if (initializedWorkers.Count() > _maxProcessCount)
@@ -247,7 +263,7 @@ public async void WorkerError(WorkerErrorEvent workerError)
             {
                 _logger.LogDebug("Handling WorkerErrorEvent for runtime:{runtime}, workerId:{workerId}", workerError.Language, workerError.WorkerId);
                 _languageWorkerErrors.Add(workerError.Exception);
-                bool isPreInitializedChannel = _webHostLanguageWorkerChannelManager.ShutdownChannelIfExists(workerError.Language, workerError.WorkerId);
+                bool isPreInitializedChannel = await _webHostLanguageWorkerChannelManager.ShutdownChannelIfExistsAsync(workerError.Language, workerError.WorkerId);
                 if (!isPreInitializedChannel)
                 {
                     _logger.LogDebug("Disposing errored channel for workerId: {channelId}, for runtime:{language}", workerError.WorkerId, workerError.Language);
@@ -257,8 +273,11 @@ public async void WorkerError(WorkerErrorEvent workerError)
                         _jobHostLanguageWorkerChannelManager.DisposeAndRemoveChannel(erroredChannel);
                     }
                 }
-                _logger.LogDebug("Restarting worker channel for runtime:{runtime}", workerError.Language);
-                await RestartWorkerChannel(workerError.Language, workerError.WorkerId);
+                if (_workerRuntime.Equals(workerError.Language, StringComparison.InvariantCultureIgnoreCase))
+                {
+                    _logger.LogDebug("Restarting worker channel for runtime:{runtime}", workerError.Language);
+                    await RestartWorkerChannel(workerError.Language, workerError.WorkerId);
+                }
             }
         }
 
@@ -279,6 +298,7 @@ protected virtual void Dispose(bool disposing)
         {
             if (!_disposed && disposing)
             {
+                _logger.LogDebug("Disposing FunctionDispatcher");
                 _workerErrorSubscription.Dispose();
                 _processStartCancellationToken.Cancel();
                 _processStartCancellationToken.Dispose();

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

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

src/WebJobs.Script/Rpc/IWebHostLanguageWorkerChannelManager.cs

@@ -10,12 +10,12 @@ public interface IWebHostLanguageWorkerChannelManager
     {
         Task<ILanguageWorkerChannel> InitializeChannelAsync(string language);
 
-        IEnumerable<ILanguageWorkerChannel> GetChannels(string language);
+        Dictionary<string, TaskCompletionSource<ILanguageWorkerChannel>> GetChannels(string language);
 
         Task SpecializeAsync();
 
-        bool ShutdownChannelIfExists(string language, string workerId);
+        Task<bool> ShutdownChannelIfExistsAsync(string language, string workerId);
 
-        void ShutdownChannels();
+        Task ShutdownChannelsAsync();
     }
 }

src/WebJobs.Script/Rpc/JobHostLanguageWorkerChannelManager.cs

@@ -4,13 +4,20 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using Microsoft.Extensions.Logging;
 
 namespace Microsoft.Azure.WebJobs.Script.Rpc
 {
     internal class JobHostLanguageWorkerChannelManager : IJobHostLanguageWorkerChannelManager
     {
+        private readonly ILogger _logger;
         private ConcurrentDictionary<string, ILanguageWorkerChannel> _channels = new ConcurrentDictionary<string, ILanguageWorkerChannel>();
 
+        public JobHostLanguageWorkerChannelManager(ILoggerFactory loggerFactory)
+        {
+            _logger = loggerFactory.CreateLogger<JobHostLanguageWorkerChannelManager>();
+        }
+
         public void AddChannel(ILanguageWorkerChannel channel)
         {
             _channels.TryAdd(channel.Id, channel);
@@ -20,6 +27,7 @@ public void DisposeAndRemoveChannel(ILanguageWorkerChannel channel)
         {
             if (_channels.TryRemove(channel.Id, out ILanguageWorkerChannel removedChannel))
             {
+                _logger.LogDebug("Disposing language worker channel with id:{workerId}", removedChannel.Id);
                 (removedChannel as IDisposable)?.Dispose();
             }
         }
@@ -30,6 +38,7 @@ public void DisposeAndRemoveChannels()
             {
                 if (_channels.TryRemove(channelId, out ILanguageWorkerChannel removedChannel))
                 {
+                    _logger.LogDebug("Disposing language worker channel with id:{workerId}", removedChannel.Id);
                     (removedChannel as IDisposable)?.Dispose();
                 }
             }

src/WebJobs.Script/Rpc/RpcInitializationService.cs

@@ -65,11 +65,10 @@ public async Task StartAsync(CancellationToken cancellationToken)
             _logger.LogDebug("Rpc Initialization Service started.");
         }
 
-        public Task StopAsync(CancellationToken cancellationToken)
+        public async Task StopAsync(CancellationToken cancellationToken)
         {
             _logger.LogDebug("Shuttingdown Rpc Channels Manager");
-            _webHostlanguageWorkerChannelManager.ShutdownChannels();
-            return Task.CompletedTask;
+            await _webHostlanguageWorkerChannelManager.ShutdownChannelsAsync();
         }
 
         public async Task OuterStopAsync(CancellationToken cancellationToken)