diff --git a/src/Orleans.Runtime/GrainDirectory/ClientDirectory.cs b/src/Orleans.Runtime/GrainDirectory/ClientDirectory.cs index 13bf800720f..b44ae357c68 100644 --- a/src/Orleans.Runtime/GrainDirectory/ClientDirectory.cs +++ b/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 @@ -46,6 +46,7 @@ internal sealed partial class ClientDirectory : SystemTarget, ILocalClientDirect #endif private readonly GrainId _localHostedClientId; private readonly IConnectedClientCollection _connectedClients; + private Func _onPublishRegistered = static () => Task.CompletedTask; private Action _schedulePublishUpdate; private Task? _runTask; private MembershipVersion _observedMembershipVersion = MembershipVersion.MinValue; @@ -58,6 +59,10 @@ internal sealed partial class ClientDirectory : SystemTarget, ILocalClientDirect // For synchronization with remote silos. private Task? _nextPublishTask; + private Task? _inflightPublishTask; + private long _publishRequestVersion; + private SiloAddress? _requestedSuccessor; + private ImmutableDictionary ConnectedClients, long Version)>? _requestedTable; private SiloAddress? _previousSuccessor; private ImmutableDictionary ConnectedClients, long Version)>? _publishedTable; @@ -349,12 +354,12 @@ private void UpdateRoutingTable(ImmutableDictionary? membershipTask = null; Task? timerTask = _refreshTimer.NextTick(RandomTimeSpan.Next(_messagingOptions.ClientRegistrationRefresh)); - while (!_shutdownCts.IsCancellationRequested) + while (!_stoppingCts.IsCancellationRequested) { try { @@ -386,10 +391,10 @@ private async Task Run() if (ShouldPublish()) { - await PublishUpdates(); + _schedulePublishUpdate(); } } - catch (OperationCanceledException) when (_shutdownCts.IsCancellationRequested) + catch (OperationCanceledException) when (_stoppingCts.IsCancellationRequested) { // Ignore during shutdown. break; @@ -403,28 +408,31 @@ private async Task Run() private bool ShouldPublish() { + if (_stoppingCts.IsCancellationRequested) + { + return false; + } + EnsureRefreshed(); lock (_lockObj) { - if (_nextPublishTask is Task task && !task.IsCompleted) + if (_stoppingCts.IsCancellationRequested) { return false; } - if (!ReferenceEquals(_table, _publishedTable)) + var successor = _consistentRing.Successor; + if (successor is null) { - return true; + return false; } - // If there is no successor, or the successor is equal to the successor the last time the table was published, - // then there is no need to publish. - var successor = _consistentRing.Successor; - if (successor is null || successor.Equals(_previousSuccessor)) + if (!ReferenceEquals(_table, _publishedTable)) { - return false; + return true; } - return true; + return !successor.Equals(_previousSuccessor); } } @@ -432,52 +440,103 @@ private void SchedulePublishUpdates() { lock (_lockObj) { + if (_stoppingCts.IsCancellationRequested) + { + return; + } + + var successor = _consistentRing.Successor; + if (successor is null) + { + return; + } + + var table = _table; if (_nextPublishTask is Task task && !task.IsCompleted) { + if (ReferenceEquals(table, _requestedTable) && successor.Equals(_requestedSuccessor)) + { + return; + } + + _requestedTable = table; + _requestedSuccessor = successor; + ++_publishRequestVersion; return; } - _nextPublishTask = this.RunOrQueueTask(PublishUpdates); + _requestedTable = table; + _requestedSuccessor = successor; + var requestVersion = ++_publishRequestVersion; + _nextPublishTask = this.QueueTask(() => RunScheduledPublish(requestVersion)); } } - private async Task PublishUpdates() + private async Task RunScheduledPublish(long requestVersion) { - // Publish clients to the next two silos in the ring - var successor = _consistentRing.Successor; - if (successor is null) + bool published; + bool newerRequest; + try + { + published = await PublishUpdates(); + } + finally { - return; + lock (_lockObj) + { + _nextPublishTask = null; + newerRequest = _publishRequestVersion > requestVersion; + } } - if (successor.Equals(_previousSuccessor)) + if (newerRequest || published && ShouldPublish()) { - _publishedTable = null; + _schedulePublishUpdate(); } + } + + private async Task PublishUpdates() + { + SiloAddress? successor; + ImmutableDictionary ConnectedClients, long Version)> newRoutes; + ImmutableDictionary ConnectedClients, long Version)>? previousRoutes; + lock (_lockObj) + { + if (_stoppingCts.IsCancellationRequested) + { + return false; + } - var newRoutes = _table; - var previousRoutes = _publishedTable; + successor = _consistentRing.Successor; + if (successor is null) + { + return false; + } + + if (!successor.Equals(_previousSuccessor)) + { + _publishedTable = null; + } + + newRoutes = _table; + previousRoutes = _publishedTable; + } if (ReferenceEquals(previousRoutes, newRoutes)) { LogDebugSkippingPublishingRoutes(); - return; + return false; } // 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,37 +550,127 @@ private async Task PublishUpdates() } } + var update = builder.ToImmutable(); try { LogDebugPublishingRoutes(successor); var remote = _grainFactory.GetSystemTarget(Constants.ClientDirectoryType, successor); - await remote.OnUpdateClientRoutes(_table).WaitAsync(_shutdownCts.Token); + if (_stoppingCts.IsCancellationRequested) + { + return false; + } - // 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)) + var publicationCompletion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + publicationCompletion.Task.Ignore(); + Volatile.Write(ref _inflightPublishTask, publicationCompletion.Task); + await _onPublishRegistered(); + + if (_stoppingCts.IsCancellationRequested) { - _publishedTable = newRoutes; - _previousSuccessor = successor; + publicationCompletion.TrySetResult(false); + return false; } + Task publishTask; + try + { + publishTask = remote.OnUpdateClientRoutes(update); + } + catch (Exception exception) + { + publicationCompletion.TrySetException(exception); + throw; + } + + ObservePublication(publishTask, publicationCompletion).Ignore(); + 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. LogDebugSuccessfullyPublishedRoutes(successor); - _nextPublishTask = null; - if (ShouldPublish()) + lock (_lockObj) { - _schedulePublishUpdate(); + if (ReferenceEquals(_publishedTable, previousRoutes)) + { + _publishedTable = newRoutes; + _previousSuccessor = successor; + } } + + return true; + } + catch (OperationCanceledException) when (_stoppingCts.IsCancellationRequested) + { + // Publication is intentionally canceled while the silo is quiescing. + return false; } catch (Exception exception) { LogErrorPublishingClientRoutingTableToSilo(exception, successor); + return false; + } + + static async Task ObservePublication(Task publishTask, TaskCompletionSource publicationCompletion) + { + try + { + await publishTask; + publicationCompletion.TrySetResult(true); + } + catch (OperationCanceledException exception) + { + publicationCompletion.TrySetCanceled(exception.CancellationToken); + } + catch (Exception exception) + { + publicationCompletion.TrySetException(exception); + } + } + } + + private async Task QuiescePublishingRoutingTable(CancellationToken cancellationToken) + { + Task? runTask; + Task? publishTask; + if (!_stoppingCts.IsCancellationRequested) + { + _stoppingCts.Cancel(); + _refreshTimer.Dispose(); + } + + lock (_lockObj) + { + runTask = _runTask; + publishTask = _nextPublishTask; + } + + if (runTask is not null) + { + await runTask.WaitAsync(cancellationToken).SuppressThrowing(); + } + + if (publishTask is not null) + { + await publishTask.WaitAsync(cancellationToken).SuppressThrowing(); + } + + var inflightPublishTask = Volatile.Read(ref _inflightPublishTask); + if (inflightPublishTask is not null) + { + await inflightPublishTask.WaitAsync(cancellationToken).SuppressThrowing(); } } void ILifecycleParticipant.Participate(ISiloLifecycle lifecycle) { + lifecycle.Subscribe( + $"{nameof(ClientDirectory)}.Quiesce", + ServiceLifecycleStage.Active, + static _ => Task.CompletedTask, + QuiescePublishingRoutingTable); + lifecycle.Subscribe( nameof(ClientDirectory), ServiceLifecycleStage.RuntimeGrainServices, @@ -530,31 +679,40 @@ void ILifecycleParticipant.Participate(ISiloLifecycle lifecycle) Task StartPublishingRoutingTable(CancellationToken ct) { - this.RunOrQueueTask(() => _runTask = this.Run()).Ignore(); - return Task.CompletedTask; - } - - async Task StopPublishingRoutingTable(CancellationToken ct) - { - _shutdownCts.Cancel(); - _refreshTimer?.Dispose(); - - if (_runTask is Task task) + var runTask = this.RunOrQueueTask(Run); + lock (_lockObj) { - await task.WaitAsync(ct).SuppressThrowing(); + _runTask = runTask; } - if (_nextPublishTask is Task publishTask) - { - await publishTask.WaitAsync(ct).SuppressThrowing(); - } + runTask.Ignore(); + return Task.CompletedTask; } + + Task StopPublishingRoutingTable(CancellationToken ct) => QuiescePublishingRoutingTable(ct); } internal class TestAccessor(ClientDirectory instance) { public Action SchedulePublishUpdate { get => instance._schedulePublishUpdate; set => instance._schedulePublishUpdate = value; } + public Func OnPublishRegistered { set => instance._onPublishRegistered = value; } public long ObservedConnectedClientsVersion { get => instance._observedConnectedClientsVersion; set => instance._observedConnectedClientsVersion = value; } + public CancellationToken StoppingToken => instance._stoppingCts.Token; + public Task DrainScheduler() => instance.RunOrQueueTask(static () => Task.CompletedTask); + public Task Quiesce(CancellationToken cancellationToken) => instance.QuiescePublishingRoutingTable(cancellationToken); + public bool PublishTasksCompleted + { + get + { + lock (instance._lockObj) + { + return instance._runTask is not { IsCompleted: false } + && instance._nextPublishTask is not { IsCompleted: false } + && Volatile.Read(ref instance._inflightPublishTask) is not { IsCompleted: false }; + } + } + } + public void SchedulePublishUpdates() => instance.SchedulePublishUpdates(); public Task PublishUpdates() => instance.PublishUpdates(); } diff --git a/test/Orleans.Core.Tests/Directory/ClientDirectoryTests.cs b/test/Orleans.Core.Tests/Directory/ClientDirectoryTests.cs index 81477accc07..9f9b3d874cb 100644 --- a/test/Orleans.Core.Tests/Directory/ClientDirectoryTests.cs +++ b/test/Orleans.Core.Tests/Directory/ClientDirectoryTests.cs @@ -338,6 +338,7 @@ public async Task PublishChangesSuccessTests() var totalUpdateCalls = new[] { 0 }; var calledSilos = new List(); + var publishedUpdates = new List ConnectedClients, long Version)>>(); SiloAddress GetOtherRemoteSilo(SiloAddress silo) => silo.Equals(remoteSilo) ? remoteSilo2 : remoteSilo; IRemoteClientDirectory CreateRemoteDirectory(SiloAddress silo) { @@ -354,26 +355,9 @@ IRemoteClientDirectory CreateRemoteDirectory(SiloAddress silo) remoteDirectory.OnUpdateClientRoutes(default!).ReturnsForAnyArgs(info => { calledSilos.Add(silo); - var callNumber = ++totalUpdateCalls[0]; + ++totalUpdateCalls[0]; var update = info.ArgAt ConnectedClients, long Version)>>(0); - - if (callNumber == 1) - { - Assert.True(update.TryGetValue(otherRemoteSilo, out var siloUpdate)); - Assert.Contains(remoteClientId2, siloUpdate.ConnectedClients); - } - else if (callNumber == 2) - { - // There should only be one silo in this update since the other remote silo is dead and this silo already has its own latest state. - Assert.Single(update); - Assert.True(update.TryGetValue(_localSilo, out var siloUpdate)); - Assert.Equal(3, siloUpdate.ConnectedClients.Count); - } - else - { - throw new InvalidOperationException("Unexpected call"); - } - + publishedUpdates.Add(update); return Task.CompletedTask; }); @@ -392,17 +376,296 @@ 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; + var successor = Assert.Single(calledSilos); + var initialUpdate = Assert.Single(publishedUpdates); + Assert.DoesNotContain(successor, initialUpdate); + Assert.True(initialUpdate.TryGetValue(GetOtherRemoteSilo(successor), out var remoteUpdate)); + Assert.Contains(remoteClientId2, remoteUpdate.ConnectedClients); - // Add clients locally and see that they are propagated to the new successor. SetLocalClients(new List { remoteClientId, remoteClientId2 }); - builder = ImmutableDictionary.CreateBuilder, long)>(); - builder[oldSuccessor] = (ImmutableHashSet.CreateRange(new[] { remoteClientId2 }), 4); - await _directory.OnUpdateClientRoutes(builder.ToImmutable()); - Assert.Equal(1, totalUpdateCalls[0]); + await _directory.OnUpdateClientRoutes( + ImmutableDictionary, long)>.Empty); + + Assert.Equal(2, totalUpdateCalls[0]); + Assert.All(calledSilos, calledSilo => Assert.Equal(successor, calledSilo)); + var followUpUpdate = publishedUpdates[1]; + Assert.Single(followUpUpdate); + Assert.True(followUpUpdate.TryGetValue(_localSilo, out var localUpdate)); + Assert.Equal(3, localUpdate.ConnectedClients.Count); + } + + [Fact] + public async Task ScheduledPublicationRepublishesChangesObservedInFlight() + { + var cancellationToken = TestContext.Current.CancellationToken; + var remoteSilo = Silo("127.0.0.1:222@100"); + var remoteDirectory = _remoteDirectories.GetOrAdd(remoteSilo, Substitute.For()); + var firstPublicationStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var secondPublicationStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var firstPublicationRelease = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var publicationCount = 0; + var firstClient = Client("local1"); + var secondClient = Client("local2"); + ImmutableDictionary ConnectedClients, long Version)>? followUpUpdate = null; + remoteDirectory.OnUpdateClientRoutes(default!).ReturnsForAnyArgs(info => + { + var update = info.ArgAt ConnectedClients, long Version)>>(0); + return Interlocked.Increment(ref publicationCount) switch + { + 1 => SignalAndReturn(firstPublicationStarted, firstPublicationRelease.Task), + 2 => CaptureFollowUpAndReturn(secondPublicationStarted, update), + _ => throw new InvalidOperationException("Unexpected publication"), + }; + }); + + _clusterMembershipService.UpdateSiloStatus(remoteSilo, SiloStatus.Active, "remoteSilo"); + _testAccessor.SchedulePublishUpdate = _testAccessor.SchedulePublishUpdates; + SetLocalClients([firstClient]); + Assert.True(_directory.TryLocalLookup(firstClient, out _)); + + _testAccessor.SchedulePublishUpdates(); + await firstPublicationStarted.Task.WaitAsync(cancellationToken); + + SetLocalClients([firstClient, secondClient]); + await _directory.OnUpdateClientRoutes( + ImmutableDictionary, long)>.Empty); + + firstPublicationRelease.TrySetResult(true); + await secondPublicationStarted.Task.WaitAsync(cancellationToken); + await _testAccessor.DrainScheduler().WaitAsync(cancellationToken); + + Assert.NotNull(followUpUpdate); + Assert.True(followUpUpdate.TryGetValue(_localSilo, out var localUpdate)); + Assert.Contains(secondClient, localUpdate.ConnectedClients); + _ = remoteDirectory.Received(2).OnUpdateClientRoutes( + Arg.Any, long)>>()); + + static Task SignalAndReturn(TaskCompletionSource started, Task task) + { + started.TrySetResult(true); + return task; + } + + Task CaptureFollowUpAndReturn( + TaskCompletionSource started, + ImmutableDictionary ConnectedClients, long Version)> update) + { + followUpUpdate = update; + started.TrySetResult(true); + return Task.CompletedTask; + } + } + + [Fact] + public async Task ScheduledPublicationIgnoresDuplicateInFlightTrigger() + { + var cancellationToken = TestContext.Current.CancellationToken; + var remoteSilo = Silo("127.0.0.1:222@100"); + var remoteDirectory = _remoteDirectories.GetOrAdd(remoteSilo, Substitute.For()); + var publicationStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var publicationRelease = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + remoteDirectory.OnUpdateClientRoutes(default!).ReturnsForAnyArgs(_ => + { + publicationStarted.TrySetResult(true); + return publicationRelease.Task; + }); + + _clusterMembershipService.UpdateSiloStatus(remoteSilo, SiloStatus.Active, "remoteSilo"); + _testAccessor.SchedulePublishUpdate = _testAccessor.SchedulePublishUpdates; + var localClient = Client("local"); + SetLocalClients([localClient]); + Assert.True(_directory.TryLocalLookup(localClient, out _)); + + _testAccessor.SchedulePublishUpdates(); + await publicationStarted.Task.WaitAsync(cancellationToken); + _testAccessor.SchedulePublishUpdates(); + + publicationRelease.TrySetResult(true); + await _testAccessor.DrainScheduler().WaitAsync(cancellationToken); + await _testAccessor.DrainScheduler().WaitAsync(cancellationToken); + + _ = remoteDirectory.Received(1).OnUpdateClientRoutes( + Arg.Any, long)>>()); + } + + [Fact] + public async Task ScheduledPublicationFailureRepublishesInFlightChangesOnce() + { + var cancellationToken = TestContext.Current.CancellationToken; + var remoteSilo = Silo("127.0.0.1:222@100"); + var remoteDirectory = _remoteDirectories.GetOrAdd(remoteSilo, Substitute.For()); + var firstPublicationStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var secondPublicationStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var firstPublicationRelease = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var publicationCount = 0; + remoteDirectory.OnUpdateClientRoutes(default!).ReturnsForAnyArgs(_ => + { + return Interlocked.Increment(ref publicationCount) switch + { + 1 => SignalAndReturn(firstPublicationStarted, firstPublicationRelease.Task), + 2 => SignalAndReturn(secondPublicationStarted, Task.CompletedTask), + _ => throw new InvalidOperationException("Unexpected publication"), + }; + }); + + _clusterMembershipService.UpdateSiloStatus(remoteSilo, SiloStatus.Active, "remoteSilo"); + _testAccessor.SchedulePublishUpdate = _testAccessor.SchedulePublishUpdates; + var localClient = Client("local"); + + _testAccessor.SchedulePublishUpdates(); + await firstPublicationStarted.Task.WaitAsync(cancellationToken); + + SetLocalClients([localClient]); + await _directory.OnUpdateClientRoutes( + ImmutableDictionary, long)>.Empty); + + firstPublicationRelease.TrySetException(new TimeoutException("Unable")); + await secondPublicationStarted.Task.WaitAsync(cancellationToken); + await _testAccessor.DrainScheduler().WaitAsync(cancellationToken); + _ = remoteDirectory.Received(2).OnUpdateClientRoutes( + Arg.Any, long)>>()); + + static Task SignalAndReturn(TaskCompletionSource started, Task task) + { + started.TrySetResult(true); + return task; + } + } + + [Fact] + public async Task ScheduledPublicationFailureWaitsForNextTrigger() + { + var cancellationToken = TestContext.Current.CancellationToken; + var remoteSilo = Silo("127.0.0.1:222@100"); + var remoteDirectory = _remoteDirectories.GetOrAdd(remoteSilo, Substitute.For()); + remoteDirectory.OnUpdateClientRoutes(default!).ReturnsForAnyArgs(_ => throw new TimeoutException("Unable")); + + _clusterMembershipService.UpdateSiloStatus(remoteSilo, SiloStatus.Active, "remoteSilo"); + _testAccessor.SchedulePublishUpdate = _testAccessor.SchedulePublishUpdates; + + _testAccessor.SchedulePublishUpdates(); + await _testAccessor.DrainScheduler().WaitAsync(cancellationToken); + await _testAccessor.DrainScheduler().WaitAsync(cancellationToken); + + _ = remoteDirectory.Received(1).OnUpdateClientRoutes( + Arg.Any, long)>>()); + } + + [Fact] + public async Task QuiescenceTracksPublicationRegisteredBeforeRpcStarts() + { + var cancellationToken = TestContext.Current.CancellationToken; + var remoteSilo = Silo("127.0.0.1:222@100"); + var remoteDirectory = _remoteDirectories.GetOrAdd(remoteSilo, Substitute.For()); + var publicationRegistered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var publicationRelease = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + _testAccessor.OnPublishRegistered = async () => + { + publicationRegistered.TrySetResult(true); + await publicationRelease.Task; + }; + + _clusterMembershipService.UpdateSiloStatus(remoteSilo, SiloStatus.Active, "remoteSilo"); + var localClient = Client("local"); + SetLocalClients([localClient]); + Assert.True(_directory.TryLocalLookup(localClient, out _)); + + try + { + _testAccessor.SchedulePublishUpdates(); + await publicationRegistered.Task.WaitAsync(cancellationToken); + + var quiesce = _testAccessor.Quiesce(cancellationToken); + publicationRelease.TrySetResult(true); + await quiesce; + + _ = remoteDirectory.DidNotReceive().OnUpdateClientRoutes( + Arg.Any, long)>>()); + } + finally + { + publicationRelease.TrySetResult(true); + } + } + + [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()); + var publicationStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var publicationRelease = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + remoteDirectory.OnUpdateClientRoutes(default!).ReturnsForAnyArgs(_ => + { + publicationStarted.TrySetResult(true); + return publicationRelease.Task; + }); + + _clusterMembershipService.UpdateSiloStatus(remoteSilo, SiloStatus.Active, "remoteSilo"); + _testAccessor.SchedulePublishUpdate = _testAccessor.SchedulePublishUpdates; + ((ILifecycleParticipant)_directory).Participate(_lifecycle); + + var membershipShutdownStarted = false; + var publicationStoppedBeforeMembershipShutdown = false; + _lifecycle.Subscribe( + "MembershipShutdownObserver", + ServiceLifecycleStage.BecomeActive, + static _ => Task.CompletedTask, + _ => + { + membershipShutdownStarted = true; + 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); + + var quiescing = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var registration = _testAccessor.StoppingToken.Register( + static state => ((TaskCompletionSource)state!).TrySetResult(true), + quiescing); + var lifecycleStop = _lifecycle.OnStop(cancellationToken); + await quiescing.Task.WaitAsync(cancellationToken); + await _testAccessor.DrainScheduler().WaitAsync(cancellationToken); + + Assert.False(lifecycleStop.IsCompleted); + Assert.False(membershipShutdownStarted); + + publicationRelease.TrySetResult(true); + await lifecycleStop; + lifecycleStopped = true; + + Assert.True( + publicationStoppedBeforeMembershipShutdown, + "Client route publication remained active when the membership shutdown stage began."); + Assert.True(_testAccessor.PublishTasksCompleted); + + var update = ImmutableDictionary, long)>.Empty.Add( + remoteSilo, + (ImmutableHashSet.Create(Client("remote")), 2)); + await _directory.OnUpdateClientRoutes(update); + _testAccessor.SchedulePublishUpdates(); + + _ = remoteDirectory.Received(1).OnUpdateClientRoutes( + Arg.Any, long)>>()); + } + finally + { + publicationRelease.TrySetResult(true); + if (!lifecycleStopped) + { + await _lifecycle.OnStop(CancellationToken.None); + } + } } private static SiloAddress Silo(string value) => SiloAddress.FromParsableString(value);