fix(directory): support rolling upgrades to distributed directory

src/Orleans.Runtime/GrainDirectory/DistributedGrainDirectory.cs

@@ -63,6 +63,7 @@ internal sealed partial class DistributedGrainDirectory : SystemTarget, IGrainDi
     private readonly IServiceProvider _serviceProvider;
     private readonly ImmutableArray<GrainDirectoryPartition> _partitions;
     private readonly CancellationTokenSource _stoppedCts = new();
+    private readonly ClusterMemberCancellationTokens _clusterMemberCancellationTokens;
     private readonly DirectoryInstruments _directoryInstruments;
 
     internal CancellationToken OnStoppedToken => _stoppedCts.Token;
@@ -96,6 +97,7 @@ public DistributedGrainDirectory(
         _membershipService = membershipService;
         _logger = logger;
         _directoryInstruments = directoryInstruments;
+        _clusterMemberCancellationTokens = new(_stoppedCts.Token);
         var partitionsPerSilo = membershipService.PartitionsPerSilo;
         var partitions = ImmutableArray.CreateBuilder<GrainDirectoryPartition>(partitionsPerSilo);
         for (var i = 0; i < partitionsPerSilo; i++)
@@ -224,6 +226,9 @@ private async Task<TResult> InvokeAsync<TState, TResult>(
     internal static bool CanInvokeClusterMember(ClusterMembershipSnapshot snapshot, SiloAddress siloAddress)
         => snapshot.GetSiloStatus(siloAddress) is SiloStatus.Active or SiloStatus.Joining or SiloStatus.ShuttingDown;
 
+    internal CancellationToken GetClusterMemberCancellationToken(SiloAddress siloAddress) =>
+        _clusterMemberCancellationTokens.GetToken(siloAddress);
+
     public async ValueTask<Immutable<List<GrainAddress>>> RecoverRegisteredActivations(MembershipVersion membershipVersion, RingRange range, SiloAddress siloAddress, int partitionIndex)
     {
         foreach (var partition in _partitions)
@@ -345,7 +350,6 @@ async Task OnShuttingDown(CancellationToken token)
             {
                 tasks.Add(partition.OnShuttingDown(token));
             }
-
             await Task.WhenAll(tasks).SuppressThrowing();
         }
     }
@@ -364,6 +368,7 @@ private async Task ProcessMembershipUpdates()
                 await foreach (var update in _membershipService.ViewUpdates.WithCancellation(_stoppedCts.Token))
                 {
                     tasks.RemoveAll(t => t.IsCompleted);
+                    _clusterMemberCancellationTokens.Update(update.ClusterMembershipSnapshot);
                     var changes = update.ClusterMembershipSnapshot.CreateUpdate(previousUpdate);
 
                     foreach (var change in changes.Changes)
@@ -405,6 +410,7 @@ private async Task ProcessMembershipUpdates()
         }
 
         await Task.WhenAll(tasks).SuppressThrowing();
+        _clusterMemberCancellationTokens.Dispose();
     }
 
     private async Task ObserveMembershipUpdateTask(Task task)
@@ -488,3 +494,75 @@ internal interface ITestHooks
     )]
     private static partial void LogTraceInvokingOperation(ILogger logger, string operation, SiloAddress owner, GrainId grainId);
 }
+
+internal sealed class ClusterMemberCancellationTokens(CancellationToken shutdownToken) : IDisposable
+{
+    private static readonly CancellationToken CanceledToken = new(canceled: true);
+    private Dictionary<SiloAddress, Entry> _entries = [];
+
+    internal int Count => Volatile.Read(ref _entries).Count;
+
+    public CancellationToken GetToken(SiloAddress siloAddress)
+    {
+        var entries = Volatile.Read(ref _entries);
+        return entries.TryGetValue(siloAddress, out var entry) ? entry.Token : CanceledToken;
+    }
+
+    public void Update(ClusterMembershipSnapshot snapshot)
+    {
+        var previous = Volatile.Read(ref _entries);
+        var updated = new Dictionary<SiloAddress, Entry>(snapshot.Members.Count);
+        foreach (var (siloAddress, _) in snapshot.Members)
+        {
+            if (!DistributedGrainDirectory.CanInvokeClusterMember(snapshot, siloAddress))
+            {
+                continue;
+            }
+
+            updated.Add(
+                siloAddress,
+                previous.TryGetValue(siloAddress, out var entry) ? entry : new Entry(shutdownToken));
+        }
+
+        Volatile.Write(ref _entries, updated);
+        DisposeRemoved(previous, updated);
+    }
+
+    public void Dispose()
+    {
+        var previous = Interlocked.Exchange(ref _entries, []);
+        DisposeRemoved(previous, []);
+    }
+
+    private static void DisposeRemoved(
+        Dictionary<SiloAddress, Entry> previous,
+        Dictionary<SiloAddress, Entry> updated)
+    {
+        foreach (var (siloAddress, entry) in previous)
+        {
+            if (!updated.ContainsKey(siloAddress))
+            {
+                entry.Dispose();
+            }
+        }
+    }
+
+    private sealed class Entry : IDisposable
+    {
+        private readonly CancellationTokenSource _source;
+
+        public Entry(CancellationToken shutdownToken)
+        {
+            _source = CancellationTokenSource.CreateLinkedTokenSource(shutdownToken);
+            Token = _source.Token;
+        }
+
+        public CancellationToken Token { get; }
+
+        public void Dispose()
+        {
+            _source.Cancel();
+            _source.Dispose();
+        }
+    }
+}

src/Orleans.Runtime/GrainDirectory/DistributedRemoteGrainDirectory.cs

@@ -76,7 +76,7 @@ private async Task EnsureDirectoryInitializedAsync(CancellationToken cancellatio
         TaskScheduler = TaskScheduler.Current,
     };
 
-    private CancellationTokenSource CreateTimeoutCts() => new(OperationTimeout);
+    private CancellationTokenSource CreateTimeoutCts() => CreateTimeoutCts(CancellationToken.None);
 
     private CancellationTokenSource CreateTimeoutCts(CancellationToken cancellationToken)
     {

src/Orleans.Runtime/GrainDirectory/GrainDirectoryPartition.cs

@@ -762,14 +762,34 @@ private async Task<T> InvokeOnClusterMember<T>(SiloAddress siloAddress, Func<Tas
         var attempt = 0;
         while (!ShutdownToken.IsCancellationRequested)
         {
-            if (!DistributedGrainDirectory.CanInvokeClusterMember(_owner.LatestClusterMembershipSnapshot, siloAddress))
+            var memberCancellationToken = _owner.GetClusterMemberCancellationToken(siloAddress);
+            if (memberCancellationToken.IsCancellationRequested)
             {
                 break;
             }
 
             try
             {
-                return await func();
+                var operationTask = func();
+                try
+                {
+                    if (operationTask.IsCompleted)
+                    {
+                        return await operationTask;
+                    }
+
+                    return await operationTask.WaitAsync(memberCancellationToken);
+                }
+                catch (OperationCanceledException) when (memberCancellationToken.IsCancellationRequested)
+                {
+                    if (operationTask.IsCompleted)
+                    {
+                        return await operationTask;
+                    }
+
+                    operationTask.Ignore();
+                    break;
+                }
             }
             catch (Exception ex)
             {
@@ -779,7 +799,8 @@ private async Task<T> InvokeOnClusterMember<T>(SiloAddress siloAddress, Func<Tas
                 }
 
                 await _owner.RefreshViewAsync(default, ShutdownToken);
-                if (!DistributedGrainDirectory.CanInvokeClusterMember(_owner.LatestClusterMembershipSnapshot, siloAddress))
+                if (memberCancellationToken.IsCancellationRequested
+                    || !DistributedGrainDirectory.CanInvokeClusterMember(_owner.LatestClusterMembershipSnapshot, siloAddress))
                 {
                     break;
                 }

src/Orleans.Runtime/Messaging/MessageCenter.cs

@@ -98,8 +98,9 @@ public async Task StopAsync()
         /// Indicates that application messages should be blocked from being sent or received.
         /// This method is used by the "fast stop" process.
         /// <para>
-        /// Specifically, all outbound application messages are dropped, except for rejections and messages to the membership table grain.
-        /// Inbound application requests are rejected, and other inbound application messages are dropped.
+        /// Specifically, all outbound application requests and one-way messages are rejected or dropped,
+        /// while responses and messages to the membership table grain are allowed to complete.
+        /// Inbound application requests are rejected with cache invalidation, and other inbound application messages are dropped.
         /// </para>
         /// </summary>
         public void BlockApplicationMessages()
@@ -145,10 +146,27 @@ public void SendMessage(Message msg)
             Debug.Assert(!msg.IsLocalOnly);
 
             // Note that if we identify or add other grains that are required for proper stopping, we will need to treat them as we do the membership table grain here.
-            if (IsBlockingApplicationMessages && !msg.IsSystemMessage && msg.Result is not Message.ResponseTypes.Rejection && !Constants.SystemMembershipTableType.Equals(msg.TargetGrain))
+            var isBlockedApplicationMessage = IsBlockingApplicationMessages
+                && !msg.IsSystemMessage
+                && msg.Direction is not Message.Directions.Response
+                && !Constants.SystemMembershipTableType.Equals(msg.TargetGrain);
+            if (isBlockedApplicationMessage)
             {
-                // Drop the message on the floor if it's an application message that isn't a rejection
-                this.messagingTrace.OnDropBlockedApplicationMessage(msg);
+                if (msg.Direction == Message.Directions.Request)
+                {
+                    ProcessRequestToInvalidActivation(
+                        msg,
+                        new GrainAddress { GrainId = msg.TargetGrain, SiloAddress = msg.TargetSilo },
+                        forwardingAddress: null,
+                        failedOperation: "Silo stopping",
+                        rejectMessages: true);
+                }
+                else
+                {
+                    this.messagingTrace.OnDropBlockedApplicationMessage(msg);
+                }
+
+                return;
             }
             else
             {
@@ -315,7 +333,7 @@ internal void ProcessRequestsToInvalidActivation(
             }
         }
 
-        private void ProcessRequestToInvalidActivation(
+        internal void ProcessRequestToInvalidActivation(
             Message message,
             GrainAddress? oldAddress,
             SiloAddress? forwardingAddress,
@@ -334,6 +352,11 @@ private void ProcessRequestToInvalidActivation(
             // IMPORTANT: do not do anything on activation context anymore, since this activation is invalid already.
             if (rejectMessages)
             {
+                if (oldAddress != null)
+                {
+                    message.AddToCacheInvalidationHeader(oldAddress, validAddress: null);
+                }
+
                 this.RejectMessage(message, Message.RejectionTypes.Transient, exc, failedOperation);
             }
             else

src/Orleans.Runtime/Networking/SiloConnection.cs

@@ -84,20 +84,23 @@ protected override void OnReceivedMessage(Message msg)
                 return;
             }
 
-            // If we've stopped application message processing, then filter those out now
+            // If we've stopped application message processing, then reject requests and filter out other messages.
             // Note that if we identify or add other grains that are required for proper stopping, we will need to treat them as we do the membership table grain here.
             if (messageCenter.IsBlockingApplicationMessages && !msg.IsSystemMessage)
             {
-                // We reject new requests, and drop all other messages
+                // We reject new requests with targeted cache invalidation and drop all other messages.
                 if (msg.Direction != Message.Directions.Request)
                 {
                     this.MessagingTrace.OnDropBlockedApplicationMessage(msg);
                     return;
                 }
 
-                MessagingInstrumentation.OnRejectedMessage(msg);
-                var rejection = this.MessageFactory.CreateRejectionResponse(msg, Message.RejectionTypes.Unrecoverable, "Silo stopping", new SiloUnavailableException());
-                this.Send(rejection);
+                messageCenter.ProcessRequestToInvalidActivation(
+                    msg,
+                    new GrainAddress { GrainId = msg.TargetGrain, SiloAddress = msg.TargetSilo },
+                    forwardingAddress: null,
+                    failedOperation: "Silo stopping",
+                    rejectMessages: true);
                 return;
             }