fix(runtime): quiesce client route publication before shutdown

src/Orleans.Runtime/GrainDirectory/ClientDirectory.cs

@@ -38,7 +38,7 @@ internal sealed partial class ClientDirectory : SystemTarget, ILocalClientDirect
     private readonly SiloAddress _localSilo;
     private readonly IClusterMembershipService _clusterMembershipService;
     private readonly SiloMessagingOptions _messagingOptions;
-    private readonly CancellationTokenSource _shutdownCts = new();
+    private readonly CancellationTokenSource _stoppingCts = new();
 #if NET9_0_OR_GREATER
     private readonly Lock _lockObj = new();
 #else
@@ -349,12 +349,12 @@ private void UpdateRoutingTable(ImmutableDictionary<SiloAddress, (ImmutableHashS
 
     private async Task Run()
     {
-        var membershipUpdates = _clusterMembershipService.MembershipUpdates.GetAsyncEnumerator(_shutdownCts.Token);
+        var membershipUpdates = _clusterMembershipService.MembershipUpdates.GetAsyncEnumerator(_stoppingCts.Token);
 
         Task<bool>? membershipTask = null;
         Task<bool>? timerTask = _refreshTimer.NextTick(RandomTimeSpan.Next(_messagingOptions.ClientRegistrationRefresh));
 
-        while (!_shutdownCts.IsCancellationRequested)
+        while (!_stoppingCts.IsCancellationRequested)
         {
             try
             {
@@ -389,7 +389,7 @@ private async Task Run()
                     await PublishUpdates();
                 }
             }
-            catch (OperationCanceledException) when (_shutdownCts.IsCancellationRequested)
+            catch (OperationCanceledException) when (_stoppingCts.IsCancellationRequested)
             {
                 // Ignore during shutdown.
                 break;
@@ -403,9 +403,19 @@ private async Task Run()
 
     private bool ShouldPublish()
     {
+        if (_stoppingCts.IsCancellationRequested)
+        {
+            return false;
+        }
+
         EnsureRefreshed();
         lock (_lockObj)
         {
+            if (_stoppingCts.IsCancellationRequested)
+            {
+                return false;
+            }
+
             if (_nextPublishTask is Task task && !task.IsCompleted)
             {
                 return false;
@@ -432,7 +442,7 @@ private void SchedulePublishUpdates()
     {
         lock (_lockObj)
         {
-            if (_nextPublishTask is Task task && !task.IsCompleted)
+            if (_stoppingCts.IsCancellationRequested || _nextPublishTask is Task task && !task.IsCompleted)
             {
                 return;
             }
@@ -443,20 +453,30 @@ private void SchedulePublishUpdates()
 
     private async Task PublishUpdates()
     {
-        // Publish clients to the next two silos in the ring
-        var successor = _consistentRing.Successor;
-        if (successor is null)
+        SiloAddress? successor;
+        ImmutableDictionary<SiloAddress, (ImmutableHashSet<GrainId> ConnectedClients, long Version)> newRoutes;
+        ImmutableDictionary<SiloAddress, (ImmutableHashSet<GrainId> ConnectedClients, long Version)>? previousRoutes;
+        lock (_lockObj)
         {
-            return;
-        }
+            if (_stoppingCts.IsCancellationRequested)
+            {
+                return;
+            }
 
-        if (successor.Equals(_previousSuccessor))
-        {
-            _publishedTable = null;
-        }
+            successor = _consistentRing.Successor;
+            if (successor is null)
+            {
+                return;
+            }
 
-        var newRoutes = _table;
-        var previousRoutes = _publishedTable;
+            if (!successor.Equals(_previousSuccessor))
+            {
+                _publishedTable = null;
+            }
+
+            newRoutes = _table;
+            previousRoutes = _publishedTable;
+        }
 
         if (ReferenceEquals(previousRoutes, newRoutes))
         {
@@ -466,18 +486,13 @@ private async Task PublishUpdates()
 
         // Try to find the minimum amount of information required to update the successor.
         var builder = newRoutes.ToBuilder();
+        builder.Remove(successor);
         if (previousRoutes is not null)
         {
             foreach (var pair in previousRoutes)
             {
                 var silo = pair.Key;
                 var (_, version) = pair.Value;
-                if (silo.Equals(successor))
-                {
-                    // No need to publish updates to the silo which originated them.
-                    continue;
-                }
-
                 if (!builder.TryGetValue(silo, out var published))
                 {
                     continue;
@@ -491,29 +506,44 @@ private async Task PublishUpdates()
             }
         }
 
+        var update = builder.ToImmutable();
         try
         {
             LogDebugPublishingRoutes(successor);
 
             var remote = _grainFactory.GetSystemTarget<IRemoteClientDirectory>(Constants.ClientDirectoryType, successor);
-            await remote.OnUpdateClientRoutes(_table).WaitAsync(_shutdownCts.Token);
+            if (_stoppingCts.IsCancellationRequested)
+            {
+                return;
+            }
+
+            var publishTask = remote.OnUpdateClientRoutes(update);
+            await publishTask.WaitAsync(_stoppingCts.Token);
 
             // Record the current lower bound of what the successor knows, so that it can be used to minimize
             // data transfer next time an update is performed.
-            if (ReferenceEquals(_publishedTable, previousRoutes))
+            LogDebugSuccessfullyPublishedRoutes(successor);
+
+            lock (_lockObj)
             {
-                _publishedTable = newRoutes;
-                _previousSuccessor = successor;
-            }
+                if (ReferenceEquals(_publishedTable, previousRoutes))
+                {
+                    _publishedTable = newRoutes;
+                    _previousSuccessor = successor;
+                }
 
-            LogDebugSuccessfullyPublishedRoutes(successor);
+                _nextPublishTask = null;
+            }
 
-            _nextPublishTask = null;
             if (ShouldPublish())
             {
                 _schedulePublishUpdate();
             }
         }
+        catch (OperationCanceledException) when (_stoppingCts.IsCancellationRequested)
+        {
+            // Publication is intentionally canceled while the silo is quiescing.
+        }
         catch (Exception exception)
         {
             LogErrorPublishingClientRoutingTableToSilo(exception, successor);
@@ -522,6 +552,12 @@ private async Task PublishUpdates()
 
     void ILifecycleParticipant<ISiloLifecycle>.Participate(ISiloLifecycle lifecycle)
     {
+        lifecycle.Subscribe(
+            $"{nameof(ClientDirectory)}.Quiesce",
+            ServiceLifecycleStage.Active,
+            static _ => Task.CompletedTask,
+            QuiescePublishingRoutingTable);
+
         lifecycle.Subscribe(
             nameof(ClientDirectory),
             ServiceLifecycleStage.RuntimeGrainServices,
@@ -530,31 +566,62 @@ void ILifecycleParticipant<ISiloLifecycle>.Participate(ISiloLifecycle lifecycle)
 
         Task StartPublishingRoutingTable(CancellationToken ct)
         {
-            this.RunOrQueueTask(() => _runTask = this.Run()).Ignore();
+            var runTask = this.RunOrQueueTask(Run);
+            lock (_lockObj)
+            {
+                _runTask = runTask;
+            }
+
+            runTask.Ignore();
             return Task.CompletedTask;
         }
 
-        async Task StopPublishingRoutingTable(CancellationToken ct)
+        async Task QuiescePublishingRoutingTable(CancellationToken ct)
         {
-            _shutdownCts.Cancel();
-            _refreshTimer?.Dispose();
+            Task? runTask;
+            Task? publishTask;
+            if (!_stoppingCts.IsCancellationRequested)
+            {
+                _stoppingCts.Cancel();
+                _refreshTimer.Dispose();
+            }
+
+            lock (_lockObj)
+            {
+                runTask = _runTask;
+                publishTask = _nextPublishTask;
+            }
 
-            if (_runTask is Task task)
+            if (runTask is not null)
             {
-                await task.WaitAsync(ct).SuppressThrowing();
+                await runTask.WaitAsync(ct).SuppressThrowing();
             }
 
-            if (_nextPublishTask is Task publishTask)
+            if (publishTask is not null)
             {
                 await publishTask.WaitAsync(ct).SuppressThrowing();
             }
         }
+
+        Task StopPublishingRoutingTable(CancellationToken ct) => QuiescePublishingRoutingTable(ct);
     }
 
     internal class TestAccessor(ClientDirectory instance)
     {
         public Action SchedulePublishUpdate { get => instance._schedulePublishUpdate; set => instance._schedulePublishUpdate = value; }
         public long ObservedConnectedClientsVersion { get => instance._observedConnectedClientsVersion; set => instance._observedConnectedClientsVersion = value; }
+        public bool PublishTasksCompleted
+        {
+            get
+            {
+                lock (instance._lockObj)
+                {
+                    return instance._runTask is not { IsCompleted: false }
+                        && instance._nextPublishTask is not { IsCompleted: false };
+                }
+            }
+        }
+        public void SchedulePublishUpdates() => instance.SchedulePublishUpdates();
         public Task PublishUpdates() => instance.PublishUpdates();
     }
 

test/Orleans.Core.Tests/Directory/ClientDirectoryTests.cs

@@ -359,6 +359,7 @@ IRemoteClientDirectory CreateRemoteDirectory(SiloAddress silo)
 
                     if (callNumber == 1)
                     {
+                        Assert.DoesNotContain(silo, update);
                         Assert.True(update.TryGetValue(otherRemoteSilo, out var siloUpdate));
                         Assert.Contains(remoteClientId2, siloUpdate.ConnectedClients);
                     }
@@ -392,17 +393,78 @@ IRemoteClientDirectory CreateRemoteDirectory(SiloAddress silo)
             await _directory.OnUpdateClientRoutes(builder.ToImmutable());
             Assert.Equal(1, totalUpdateCalls[0]);
 
-            var oldSuccessor = calledSilos.Last();
-            _clusterMembershipService.UpdateSiloStatus(oldSuccessor, SiloStatus.Dead, "blah");
-            var newSuccessor = GetOtherRemoteSilo(oldSuccessor);
-            totalUpdateCalls[0] = 0;
-
-            // Add clients locally and see that they are propagated to the new successor.
+            var successor = Assert.Single(calledSilos);
             SetLocalClients(new List<GrainId> { remoteClientId, remoteClientId2 });
-            builder = ImmutableDictionary.CreateBuilder<SiloAddress, (ImmutableHashSet<GrainId>, long)>();
-            builder[oldSuccessor] = (ImmutableHashSet.CreateRange(new[] { remoteClientId2 }), 4);
-            await _directory.OnUpdateClientRoutes(builder.ToImmutable());
-            Assert.Equal(1, totalUpdateCalls[0]);
+            await _directory.OnUpdateClientRoutes(
+                ImmutableDictionary<SiloAddress, (ImmutableHashSet<GrainId>, long)>.Empty);
+
+            Assert.Equal(2, totalUpdateCalls[0]);
+            Assert.All(calledSilos, calledSilo => Assert.Equal(successor, calledSilo));
+        }
+
+        [Fact]
+        public async Task PublishingStopsBeforeMembershipShutdownBegins()
+        {
+            var cancellationToken = TestContext.Current.CancellationToken;
+            var remoteSilo = Silo("127.0.0.1:222@100");
+            var remoteDirectory = _remoteDirectories.GetOrAdd(remoteSilo, Substitute.For<IRemoteClientDirectory>());
+            var publicationStarted = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            var publicationRelease = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            remoteDirectory.OnUpdateClientRoutes(default!).ReturnsForAnyArgs(_ =>
+            {
+                publicationStarted.TrySetResult(true);
+                return publicationRelease.Task;
+            });
+
+            _clusterMembershipService.UpdateSiloStatus(remoteSilo, SiloStatus.Active, "remoteSilo");
+            ((ILifecycleParticipant<ISiloLifecycle>)_directory).Participate(_lifecycle);
+
+            var publicationStoppedBeforeMembershipShutdown = false;
+            _lifecycle.Subscribe(
+                "MembershipShutdownObserver",
+                ServiceLifecycleStage.BecomeActive,
+                static _ => Task.CompletedTask,
+                _ =>
+                {
+                    publicationStoppedBeforeMembershipShutdown = _testAccessor.PublishTasksCompleted;
+                    return Task.CompletedTask;
+                });
+
+            var lifecycleStopped = false;
+            try
+            {
+                var localClient = Client("local");
+                SetLocalClients([localClient]);
+                Assert.True(_directory.TryLocalLookup(localClient, out _));
+
+                await _lifecycle.OnStart(cancellationToken);
+                await publicationStarted.Task.WaitAsync(cancellationToken);
+
+                await _lifecycle.OnStop(cancellationToken);
+                lifecycleStopped = true;
+
+                Assert.True(
+                    publicationStoppedBeforeMembershipShutdown,
+                    "Client route publication remained active when the membership shutdown stage began.");
+                Assert.True(_testAccessor.PublishTasksCompleted);
+
+                var update = ImmutableDictionary<SiloAddress, (ImmutableHashSet<GrainId>, long)>.Empty.Add(
+                    remoteSilo,
+                    (ImmutableHashSet.Create(Client("remote")), 2));
+                await _directory.OnUpdateClientRoutes(update);
+                _testAccessor.SchedulePublishUpdates();
+
+                _ = remoteDirectory.Received(1).OnUpdateClientRoutes(
+                    Arg.Any<ImmutableDictionary<SiloAddress, (ImmutableHashSet<GrainId>, long)>>());
+            }
+            finally
+            {
+                publicationRelease.TrySetResult(true);
+                if (!lifecycleStopped)
+                {
+                    await _lifecycle.OnStop(CancellationToken.None);
+                }
+            }
         }
 
         private static SiloAddress Silo(string value) => SiloAddress.FromParsableString(value);