fix(directory): make rolling upgrade handoffs durable

src/Orleans.Runtime/GrainDirectory/DistributedRemoteGrainDirectory.cs

@@ -150,19 +150,6 @@ private async Task ExecutePendingOperations()
         }
     }
 
-    private void DestroyDuplicateActivations(Dictionary<SiloAddress, List<GrainAddress>>? duplicates)
-    {
-        if (duplicates is null || duplicates.Count == 0)
-        {
-            return;
-        }
-
-        EnqueueOperation(
-            nameof(DestroyDuplicateActivations),
-            duplicates,
-            static (self, state) => self.DestroyDuplicateActivationsAsync((Dictionary<SiloAddress, List<GrainAddress>>)state));
-    }
-
     private async Task DestroyDuplicateActivationsAsync(Dictionary<SiloAddress, List<GrainAddress>> duplicates)
     {
         while (duplicates.Count > 0)
@@ -205,7 +192,6 @@ await Parallel.ForEachAsync(Enumerable.Range(0, pendingRegistrations.Count), opt
             }
         });
 
-        Dictionary<SiloAddress, List<GrainAddress>>? duplicates = null;
         Exception? failure = null;
         for (var i = pendingRegistrations.Count - 1; i >= 0; i--)
         {
@@ -221,10 +207,10 @@ await Parallel.ForEachAsync(Enumerable.Range(0, pendingRegistrations.Count), opt
             {
                 if (registration.SiloAddress is { } siloAddress)
                 {
-                    if (duplicates is null || !duplicates.TryGetValue(siloAddress, out var activations))
+                    if (!batch.DuplicateActivations.TryGetValue(siloAddress, out var activations))
                     {
                         activations = [];
-                        (duplicates ??= []).Add(siloAddress, activations);
+                        batch.DuplicateActivations.Add(siloAddress, activations);
                     }
 
                     activations.Add(registration);
@@ -234,7 +220,7 @@ await Parallel.ForEachAsync(Enumerable.Range(0, pendingRegistrations.Count), opt
             pendingRegistrations.RemoveAt(i);
         }
 
-        DestroyDuplicateActivations(duplicates);
+        await DestroyDuplicateActivationsAsync(batch.DuplicateActivations);
 
         if (failure is not null)
         {
@@ -243,6 +229,7 @@ await Parallel.ForEachAsync(Enumerable.Range(0, pendingRegistrations.Count), opt
         }
 
         LogInformationAcceptSplitPartitionCompleted(_logger, Silo, batch.InitialCount);
+        batch.Completion.TrySetResult();
     }
 
     public async Task<AddressAndTag> RegisterAsync(GrainAddress address, int hopCount)
@@ -337,15 +324,29 @@ await Parallel.ForEachAsync(Enumerable.Range(0, grainAndETagList.Count), options
     public Task AcceptSplitPartition(List<GrainAddress> singleActivations)
     {
         LogInformationAcceptSplitPartitionStarted(_logger, Silo, singleActivations.Count);
-        if (singleActivations.Count > 0)
+        if (singleActivations.Count == 0)
+        {
+            return Task.CompletedTask;
+        }
+
+        lock (_pendingOperations)
         {
+            foreach (var operation in _pendingOperations)
+            {
+                if (operation.State is SplitPartitionRegistrationBatch existingBatch
+                    && existingBatch.Matches(singleActivations))
+                {
+                    return existingBatch.Completion.Task.WaitAsync(_directory.OnStoppedToken);
+                }
+            }
+
+            var batch = new SplitPartitionRegistrationBatch(singleActivations);
             EnqueueOperation(
                 nameof(AcceptSplitPartition),
-                new SplitPartitionRegistrationBatch([.. singleActivations]),
+                batch,
                 static (self, state) => self.ProcessSplitPartitionRegistrationsAsync((SplitPartitionRegistrationBatch)state));
+            return batch.Completion.Task.WaitAsync(_directory.OnStoppedToken);
         }
-
-        return Task.CompletedTask;
     }
 
     [LoggerMessage(
@@ -378,9 +379,15 @@ public Task AcceptSplitPartition(List<GrainAddress> singleActivations)
     )]
     private static partial void LogWarningOperationFailedRetry(ILogger logger, Exception exception, string operation);
 
-    private sealed class SplitPartitionRegistrationBatch(List<GrainAddress> pendingRegistrations)
+    private sealed class SplitPartitionRegistrationBatch(List<GrainAddress> registrations)
     {
-        public int InitialCount { get; } = pendingRegistrations.Count;
-        public List<GrainAddress> PendingRegistrations { get; } = pendingRegistrations;
+        private readonly List<GrainAddress> _registrations = [.. registrations];
+
+        public int InitialCount => _registrations.Count;
+        public List<GrainAddress> PendingRegistrations { get; } = [.. registrations];
+        public Dictionary<SiloAddress, List<GrainAddress>> DuplicateActivations { get; } = [];
+        public TaskCompletionSource Completion { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+        public bool Matches(List<GrainAddress> registrations) => _registrations.SequenceEqual(registrations);
     }
 }

test/Orleans.GrainDirectory.Tests/GrainDirectory/GrainDirectoryRollingUpgradeTests.cs

@@ -883,6 +883,7 @@ await ValidateTrackedDirectoryCheckpointAsync(
                     staleCacheEvidence);
             }
 
+            AssertSplitPartitionHandoffsAreDurable(logs);
             var finalWorkerProgress = traffic.GetWorkerProgress();
             phase.Set("all-distributed-final");
             Assert.Equal(SiloCount, cluster.Silos.Count);
@@ -1305,6 +1306,46 @@ private void AssertNoImpactfulErrors(PhaseAwareLogCapture logs)
         Assert.Empty(errors);
     }
 
+    private static void AssertSplitPartitionHandoffsAreDurable(PhaseAwareLogCapture logs)
+    {
+        var completedHandoffs = new Dictionary<(string Silo, int Count), int>();
+        var removedHandoffs = 0;
+        foreach (var entry in logs.ToArray())
+        {
+            var isCompleted = string.Equals(
+                entry.EventId.Name,
+                "LogInformationAcceptSplitPartitionCompleted",
+                StringComparison.Ordinal);
+            var isRemoved = string.Equals(
+                entry.EventId.Name,
+                "LogInformationRemovedTransferredEntries",
+                StringComparison.Ordinal);
+            if (!isCompleted && !isRemoved)
+            {
+                continue;
+            }
+
+            Assert.False(string.IsNullOrEmpty(entry.HandoffSilo));
+            Assert.True(entry.HandoffCount.HasValue);
+            var handoff = (entry.HandoffSilo!, entry.HandoffCount.Value);
+            completedHandoffs.TryGetValue(handoff, out var count);
+            if (isCompleted)
+            {
+                completedHandoffs[handoff] = count + 1;
+            }
+            else
+            {
+                removedHandoffs++;
+                Assert.True(
+                    count > 0,
+                    $"Sender-side registrations were removed before the recipient completed the handoff: {entry}");
+                completedHandoffs[handoff] = count - 1;
+            }
+        }
+
+        Assert.True(removedHandoffs > 0, "The rolling upgrade did not exercise a non-empty split-partition handoff.");
+    }
+
     private static bool IsExpectedIntentionalRestartLog(PhaseAwareLogEntry entry)
     {
         if (!entry.Phase.StartsWith("restart-", StringComparison.Ordinal))
@@ -1842,9 +1883,42 @@ public void Log<TState>(
             {
                 if (IsEnabled(logLevel))
                 {
-                    capture.Add(siloName, category, logLevel, eventId, formatter(state, exception), exception);
+                    var (handoffSilo, handoffCount) = GetHandoffIdentity(state);
+                    capture.Add(
+                        siloName,
+                        category,
+                        logLevel,
+                        eventId,
+                        formatter(state, exception),
+                        exception,
+                        handoffSilo,
+                        handoffCount);
                 }
             }
+
+            private static (string? Silo, int? Count) GetHandoffIdentity<TState>(TState state)
+            {
+                if (state is not IEnumerable<KeyValuePair<string, object?>> properties)
+                {
+                    return default;
+                }
+
+                string? silo = null;
+                int? count = null;
+                foreach (var property in properties)
+                {
+                    if (property.Key is "Silo" or "AddedSilo")
+                    {
+                        silo = property.Value?.ToString();
+                    }
+                    else if (property.Key == "Count" && property.Value is int value)
+                    {
+                        count = value;
+                    }
+                }
+
+                return (silo, count);
+            }
         }
     }
 
@@ -1858,7 +1932,9 @@ public void Add(
             LogLevel level,
             EventId eventId,
             string message,
-            Exception? exception)
+            Exception? exception,
+            string? handoffSilo,
+            int? handoffCount)
         {
             var baseException = exception?.GetBaseException();
             _entries.Enqueue(
@@ -1871,7 +1947,9 @@ public void Add(
                     eventId,
                     message,
                     baseException?.GetType().FullName,
-                    baseException?.Message));
+                    baseException?.Message,
+                    handoffSilo,
+                    handoffCount));
         }
 
         public PhaseAwareLogEntry[] ToArray() => _entries.ToArray();
@@ -1886,7 +1964,9 @@ private sealed record PhaseAwareLogEntry(
         EventId EventId,
         string Message,
         string? ExceptionType,
-        string? ExceptionMessage)
+        string? ExceptionMessage,
+        string? HandoffSilo,
+        int? HandoffCount)
     {
         public override string ToString() =>
             $"{Timestamp:O} phase='{Phase}' silo='{SiloName}' [{Level}] [{Category}] ({EventId.Id}:{EventId.Name}) "