Block for environmentreload request to complete

Author: pragnagopa

src/WebJobs.Script/Eventing/Rpc/WorkerProcessReadyEvent.cs

@@ -21,7 +21,7 @@ public interface ILanguageWorkerChannel
 
         void SendFunctionLoadRequests();
 
-        void SendFunctionEnvironmentReloadRequest();
+        Task SendFunctionEnvironmentReloadRequest();
 
         Task StartWorkerProcessAsync();
     }

src/WebJobs.Script/Rpc/ILanguageWorkerChannel.cs

@@ -52,6 +52,7 @@ internal class LanguageWorkerChannel : ILanguageWorkerChannel, IDisposable
         private Capabilities _workerCapabilities;
         private ILogger _workerChannelLogger;
         private ILanguageWorkerProcess _languageWorkerProcess;
+        private TaskCompletionSource<bool> _reloadTask = new TaskCompletionSource<bool>();
 
         internal LanguageWorkerChannel(
            string workerId,
@@ -139,15 +140,19 @@ internal void SendWorkerInitRequest(RpcEvent startEvent)
             });
         }
 
-        internal void PublishWorkerProcessReadyEvent(FunctionEnvironmentReloadResponse res)
+        internal void FunctionEnvironmentReloadResponse(FunctionEnvironmentReloadResponse res)
         {
-            if (_disposing)
+            _workerChannelLogger.LogDebug("Received FunctionEnvironmentReloadResponse");
+            if (_reloadTask.Task.IsCompleted)
             {
-                // do not publish ready events when disposing
-                return;
+                throw new InvalidOperationException("FunctionEnvironmentReloadResponse received more than once");
             }
-            WorkerProcessReadyEvent wpEvent = new WorkerProcessReadyEvent(_workerId, _runtime);
-            _eventManager.Publish(wpEvent);
+            if (res.Result.IsFailure(out Exception relaodEnvironmentVariablesException))
+            {
+                _workerChannelLogger.LogError(relaodEnvironmentVariablesException, "Failed to reload environment variables");
+                _reloadTask.SetResult(false);
+            }
+            _reloadTask.SetResult(true);
         }
 
         internal void PublishRpcChannelReadyEvent(RpcEvent initEvent)
@@ -206,10 +211,14 @@ public void SendFunctionLoadRequests()
             }
         }
 
-        public void SendFunctionEnvironmentReloadRequest()
+        public Task SendFunctionEnvironmentReloadRequest()
         {
-            _eventSubscriptions.Add(_inboundWorkerEvents.Where(msg => msg.MessageType == MsgType.FunctionEnvironmentReloadResponse)
-      .Subscribe((msg) => PublishWorkerProcessReadyEvent(msg.Message.FunctionEnvironmentReloadResponse)));
+            _workerChannelLogger.LogDebug("Sending FunctionEnvironmentReloadRequest");
+            _eventSubscriptions
+                .Add(_inboundWorkerEvents.Where(msg => msg.MessageType == MsgType.FunctionEnvironmentReloadResponse)
+                .Timeout(workerInitTimeout)
+                .Take(1)
+                .Subscribe((msg) => FunctionEnvironmentReloadResponse(msg.Message.FunctionEnvironmentReloadResponse)));
 
             IDictionary processEnv = Environment.GetEnvironmentVariables();
 
@@ -223,6 +232,7 @@ public void SendFunctionEnvironmentReloadRequest()
             {
                 FunctionEnvironmentReloadRequest = request
             });
+            return _reloadTask.Task;
         }
 
         internal void SendFunctionLoadRequest(FunctionMetadata metadata)

src/WebJobs.Script/Rpc/LanguageWorkerChannel.cs

@@ -94,13 +94,7 @@ public async Task SpecializeAsync()
             if (_workerRuntime != null && languageWorkerChannel != null)
             {
                 _logger.LogInformation("Loading environment variables for runtime: {runtime}", _workerRuntime);
-                IObservable<WorkerProcessReadyEvent> processReadyEvents = _eventManager.OfType<WorkerProcessReadyEvent>()
-                .Where(msg => string.Equals(msg.Language, _workerRuntime, StringComparison.OrdinalIgnoreCase))
-                .Timeout(workerInitTimeout);
-                languageWorkerChannel.SendFunctionEnvironmentReloadRequest();
-
-                // Wait for response from language worker process
-                await processReadyEvents.FirstAsync();
+                await languageWorkerChannel.SendFunctionEnvironmentReloadRequest();
             }
             _shutdownStandbyWorkerChannels();
         }

src/WebJobs.Script/Rpc/WebHostLanguageWorkerChannelManager.cs

@@ -82,7 +82,7 @@ await TestHelpers.Await(() =>
             Assert.NotSame(GetCachedTimeZoneInfo(), _originalTimeZoneInfoCache);
         }
 
-        [Fact(Skip = "https://github.com/Azure/azure-functions-host/issues/4230")]
+        [Fact]
         public async Task StandbyModeE2E_Java()
         {
             _settings.Add(EnvironmentSettingNames.AzureWebsiteInstanceId, Guid.NewGuid().ToString());

test/WebJobs.Script.Tests.Integration/Host/StandbyManager/StandbyManagerE2ETests_Windows.cs

@@ -150,6 +150,20 @@ public void FunctionLoadRequest_IsExpected()
             Assert.True(proxyFunctionLoadRequest.Metadata.IsProxy);
         }
 
+        [Fact]
+        public void Multiple_FunctionEnvironmentReloadResponse_Throws()
+        {
+            _workerChannel.SendFunctionEnvironmentReloadRequest();
+            _testFunctionRpcService.PublishFunctionEnvironmentReloadResponseEvent();
+
+            var traces = _logger.GetLogMessages();
+
+            Assert.True(traces.Any(m => string.Equals(m.FormattedMessage, "Received FunctionEnvironmentReloadResponse")));
+            Assert.True(traces.Any(m => string.Equals(m.FormattedMessage, "Sending FunctionEnvironmentReloadRequest")));
+            var ex = Assert.Throws<InvalidOperationException>(() => _workerChannel.FunctionEnvironmentReloadResponse(TestFunctionRpcService.GetTestFunctionEnvReloadResponse()));
+            Assert.Contains("FunctionEnvironmentReloadResponse received more than once", ex.Message);
+        }
+
         private IEnumerable<FunctionMetadata> GetTestFunctionsList(string runtime)
         {
             return new List<FunctionMetadata>()

test/WebJobs.Script.Tests/Rpc/LanguageWorkerChannelTests.cs

@@ -49,6 +49,29 @@ public void PublishFunctionLoadResponseEvent(string functionId)
             _eventManager.Publish(new InboundEvent(_workerId, responseMessage));
         }
 
+        public void PublishFunctionEnvironmentReloadResponseEvent()
+        {
+            FunctionEnvironmentReloadResponse relaodEnvResponse = GetTestFunctionEnvReloadResponse();
+            StreamingMessage responseMessage = new StreamingMessage()
+            {
+                FunctionEnvironmentReloadResponse = relaodEnvResponse
+            };
+            _eventManager.Publish(new InboundEvent(_workerId, responseMessage));
+        }
+
+        public static FunctionEnvironmentReloadResponse GetTestFunctionEnvReloadResponse()
+        {
+            StatusResult statusResult = new StatusResult()
+            {
+                Status = StatusResult.Types.Status.Success
+            };
+            FunctionEnvironmentReloadResponse relaodEnvResponse = new FunctionEnvironmentReloadResponse()
+            {
+                Result = statusResult
+            };
+            return relaodEnvResponse;
+        }
+
         public void PublishInvocationResponseEvent()
         {
             StatusResult statusResult = new StatusResult()

test/WebJobs.Script.Tests/Rpc/TestFunctionRpcService.cs

@@ -50,9 +50,10 @@ public void SendFunctionLoadRequests()
             _testLogger.LogInformation("RegisterFunctions called");
         }
 
-        public void SendFunctionEnvironmentReloadRequest()
+        public Task SendFunctionEnvironmentReloadRequest()
         {
             _testLogger.LogInformation("SendFunctionEnvironmentReloadRequest called");
+            return Task.CompletedTask;
         }
 
         public void SendInvocationRequest(ScriptInvocationContext context)