Fixing thread safety issue in the GrpcWorkerChannel.LoadResponse method (#10363) (#10400)

* Making `GrpcWorkerChannel.LoadResponse` thread-safe

* Replace `Task.Delay(500)` with `TestHelpers.Await` in the test.

Author: kshyju

release_notes.md

@@ -32,3 +32,4 @@
 - Update PowerShell 7.2 worker to [4.0.4020](https://github.com/Azure/azure-functions-powershell-worker/releases/tag/v4.0.4020)
 - Update PowerShell 7.4 worker to [4.0.4021](https://github.com/Azure/azure-functions-powershell-worker/releases/tag/v4.0.4021)
 - Trim FunctionsNetHost artifacts (#10364)
+- Resolved thread safety issue in the `GrpcWorkerChannel.LoadResponse` method. (#10352)

src/WebJobs.Script.Grpc/Channel/GrpcWorkerChannel.cs

@@ -70,7 +70,7 @@ internal partial class GrpcWorkerChannel : IRpcWorkerChannel, IDisposable
         private ConcurrentDictionary<string, ExecutingInvocation> _executingInvocations = new();
         private IDictionary<string, BufferBlock<ScriptInvocationContext>> _functionInputBuffers = new ConcurrentDictionary<string, BufferBlock<ScriptInvocationContext>>();
         private ConcurrentDictionary<string, TaskCompletionSource<bool>> _workerStatusRequests = new ConcurrentDictionary<string, TaskCompletionSource<bool>>();
-        private List<IDisposable> _inputLinks = new List<IDisposable>();
+        private ConcurrentBag<IDisposable> _inputLinks = new ConcurrentBag<IDisposable>();
         private List<IDisposable> _eventSubscriptions = new List<IDisposable>();
         private IDisposable _startLatencyMetric;
         private IEnumerable<FunctionMetadata> _functions;

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

@@ -939,6 +939,48 @@ public async Task ReceivesInboundEvent_FunctionLoadResponse()
             Assert.True(traces.Any(m => string.Equals(m.FormattedMessage, "Received FunctionLoadResponse for function: 'js1' with functionId: 'TestFunctionId1'.")), "FunctionLoadResponse TestFunctionId1");
         }
 
+        [Fact]
+        public async Task Receives_Individual_FunctionLoadResponses_Parallel()
+        {
+            await CreateDefaultWorkerChannel();
+
+            var startStreamMessage = new StreamingMessage()
+            {
+                StartStream = new StartStream()
+                {
+                    WorkerId = _workerId
+                }
+            };
+
+            var rpcEvent = new GrpcEvent(_workerId, startStreamMessage);
+            _workerChannel.SendWorkerInitRequest(rpcEvent);
+
+            var functionMetadataList = GetTestFunctionsList("node", numberOfFunctions: 250);
+            _workerChannel.SetupFunctionInvocationBuffers(functionMetadataList);
+            _workerChannel.SendFunctionLoadRequests(managedDependencyOptions: null, TimeSpan.FromSeconds(1));
+
+            var allFunctionIdsAndNames = functionMetadataList.Select(f => new { Id = f.Properties["FunctionId"].ToString(), f.Name }).ToList();
+
+            // Send function load responses for each function, not necessarily in the order the load requests were sent.
+            var publishFunctionLoadResponseTasks = allFunctionIdsAndNames.Select(function =>
+                Task.Run(() => _testFunctionRpcService.PublishFunctionLoadResponseEvent(function.Id)));
+
+            await Task.WhenAll(publishFunctionLoadResponseTasks);
+
+            await TestHelpers.Await(() =>
+            {
+                return _logger.GetLogMessages().Count(m => m.FormattedMessage.StartsWith("Received FunctionLoadResponse")) == allFunctionIdsAndNames.Count;
+            });
+
+            var traces = _logger.GetLogMessages();
+
+            foreach (var function in allFunctionIdsAndNames)
+            {
+                Assert.True(traces.Any(m => string.Equals(m.FormattedMessage, $"Setting up FunctionInvocationBuffer for function: '{function.Name}' with functionId: '{function.Id}'")), $"setup {function.Id}");
+                Assert.True(traces.Any(m => string.Equals(m.FormattedMessage, $"Received FunctionLoadResponse for function: '{function.Name}' with functionId: '{function.Id}'.")), $"FunctionLoadResponse {function.Id}");
+            }
+        }
+
         [Fact]
         public async Task ReceivesInboundEvent_Failed_FunctionLoadResponses()
         {
@@ -1489,41 +1531,7 @@ await TestHelpers.Await(() => GetInvocationLogs().Length == logLoop,
 
         private static IEnumerable<FunctionMetadata> GetTestFunctionsList(string runtime, bool addWorkerProperties = false)
         {
-            var metadata1 = new FunctionMetadata()
-            {
-                Language = runtime,
-                Name = "js1"
-            };
-
-            metadata1.SetFunctionId("TestFunctionId1");
-            metadata1.Properties.Add(LogConstants.CategoryNameKey, "testcat1");
-            metadata1.Properties.Add(ScriptConstants.LogPropertyHostInstanceIdKey, "testhostId1");
-
-            if (addWorkerProperties)
-            {
-                metadata1.Properties.Add("worker.functionId", "fn1");
-            }
-
-            var metadata2 = new FunctionMetadata()
-            {
-                Language = runtime,
-                Name = "js2",
-            };
-
-            metadata2.SetFunctionId("TestFunctionId2");
-            metadata2.Properties.Add(LogConstants.CategoryNameKey, "testcat2");
-            metadata2.Properties.Add(ScriptConstants.LogPropertyHostInstanceIdKey, "testhostId2");
-
-            if (addWorkerProperties)
-            {
-                metadata2.Properties.Add("WORKER.functionId", "fn2");
-            }
-
-            return new List<FunctionMetadata>()
-            {
-                metadata1,
-                metadata2
-            };
+            return GetTestFunctionsList(runtime, numberOfFunctions: 2, addWorkerProperties);
         }
 
         public static ScriptInvocationContext GetTestScriptInvocationContext(Guid invocationId, TaskCompletionSource<ScriptInvocationResult> resultSource,
@@ -1548,6 +1556,33 @@ public static ScriptInvocationContext GetTestScriptInvocationContext(Guid invoca
             };
         }
 
+        private static List<FunctionMetadata> GetTestFunctionsList(string runtime, int numberOfFunctions, bool addWorkerProperties = false)
+        {
+            var functions = new List<FunctionMetadata>();
+
+            for (int i = 1; i <= numberOfFunctions; i++)
+            {
+                var metadata = new FunctionMetadata()
+                {
+                    Language = runtime,
+                    Name = $"js{i}"
+                };
+
+                metadata.SetFunctionId($"TestFunctionId{i}");
+                metadata.Properties.Add(LogConstants.CategoryNameKey, $"testcat1");
+                metadata.Properties.Add(ScriptConstants.LogPropertyHostInstanceIdKey, $"testhostId1");
+
+                if (addWorkerProperties)
+                {
+                    metadata.Properties.Add("worker.functionId", $"fn{i}");
+                }
+
+                functions.Add(metadata);
+            }
+
+            return functions;
+        }
+
         /// <summary>
         /// The <see cref="ScriptInvocationContext"/> would contain inputs that can be transferred over shared memory.
         /// </summary>