Fix a potential NullReferenceException in PlacementWorker (#9386)

Author: ledjon-behluli

src/Orleans.Runtime/Placement/PlacementService.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
@@ -246,7 +247,7 @@ private class PlacementWorker
 #pragma warning restore IDE0052 // Remove unread private members
             private readonly object _lockObj = new();
             private readonly PlacementService _placementService;
-            private List<(Message Message, TaskCompletionSource<bool> Completion)> _messages = new();
+            private List<(Message Message, TaskCompletionSource Completion)> _messages = new();
 
             public PlacementWorker(PlacementService placementService)
             {
@@ -259,7 +260,7 @@ public PlacementWorker(PlacementService placementService)
 
             public Task AddressMessage(Message message)
             {
-                var completion = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+                var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
 
                 lock (_lockObj)
                 {
@@ -271,7 +272,7 @@ public Task AddressMessage(Message message)
                 return completion.Task;
             }
 
-            private List<(Message Message, TaskCompletionSource<bool> Completion)> GetMessages()
+            private List<(Message Message, TaskCompletionSource Completion)> GetMessages()
             {
                 lock (_lockObj)
                 {
@@ -346,36 +347,31 @@ private void AddressWaitingMessages(GrainPlacementWorkItem completedWorkItem)
             {
                 var resultTask = completedWorkItem.Result;
                 var messages = completedWorkItem.Messages;
-                if (resultTask.IsCompletedSuccessfully)
+
+                try
                 {
+                    var siloAddress = resultTask.Result;
                     foreach (var message in messages)
                     {
-                        var siloAddress = resultTask.Result;
                         _placementService.SetMessageTargetPlacement(message.Message, siloAddress);
-                        message.Completion.TrySetResult(true);
+                        message.Completion.TrySetResult();
                     }
-
-                    messages.Clear();
                 }
-                else
+                catch (Exception exception)
                 {
-                    foreach (var message in messages)
+                    var originalException = exception switch
                     {
-                        message.Completion.TrySetException(OriginalException(resultTask.Exception));
-                    }
-
-                    messages.Clear();
-                }
+                        AggregateException ae when ae.InnerExceptions.Count == 1 => ae.InnerException,
+                        _ => exception,
+                    };
 
-                static Exception OriginalException(AggregateException exception)
-                {
-                    if (exception.InnerExceptions.Count == 1)
+                    foreach (var message in messages)
                     {
-                        return exception.InnerException;
+                        message.Completion.TrySetException(originalException);
                     }
-
-                    return exception;
                 }
+
+                messages.Clear();
             }
 
             private async Task<SiloAddress> GetOrPlaceActivationAsync(Message firstMessage)
@@ -411,7 +407,7 @@ private async Task<SiloAddress> GetOrPlaceActivationAsync(Message firstMessage)
 
             private class GrainPlacementWorkItem
             {
-                public List<(Message Message, TaskCompletionSource<bool> Completion)> Messages { get; } = new();
+                public List<(Message Message, TaskCompletionSource Completion)> Messages { get; } = new();
 
                 public Task<SiloAddress> Result { get; set; }
             }