diff --git a/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancerMonitor.cs b/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancerMonitor.cs index 47b96f39aa4..81889138bd6 100644 --- a/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancerMonitor.cs +++ b/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancerMonitor.cs @@ -61,6 +61,13 @@ public void Participate(ISiloLifecycle observer) OnStart, _ => Task.CompletedTask); + // The rebalancer must migrate before the catalog deactivates ordinary grains during shutdown. + observer.Subscribe( + $"{nameof(ActivationRebalancerMonitor)}.Migration", + ServiceLifecycleStage.GrainDeactivation + 1, + _ => Task.CompletedTask, + MigrateRebalancerOnStop); + observer.Subscribe( nameof(ActivationRepartitioner), ServiceLifecycleStage.ApplicationServices, @@ -101,24 +108,26 @@ await this.RunOrQueueTask(() => }); } - private async Task OnStop(CancellationToken cancellationToken) + private async Task MigrateRebalancerOnStop(CancellationToken cancellationToken) { - await this.RunOrQueueTask(() => + await this.RunOrQueueTask(async () => { - if (_latestReport is { } report && Silo.IsSameLogicalSilo(report.Host)) + if (!cancellationToken.IsCancellationRequested && + _activationDirectory.FindTarget(_rebalancerGrain.GetGrainId()) is { } activation) { - if (_activationDirectory.FindTarget(_rebalancerGrain.GetGrainId()) is { } activation) - { - LogMigratingRebalancer(Silo); - activation.Migrate(null, cancellationToken); // migrate it anywhere else - } + LogMigratingRebalancer(Silo); + activation.Migrate(null, cancellationToken); // migrate it anywhere else + await activation.Deactivated.WaitAsync(cancellationToken); } - - _monitorTimer?.Dispose(); - return Task.CompletedTask; }); } + private Task OnStop(CancellationToken cancellationToken) => this.RunOrQueueTask(() => + { + _monitorTimer?.Dispose(); + return Task.CompletedTask; + }); + public Task ResumeRebalancing() => _rebalancerGrain.ResumeRebalancing(); public Task SuspendRebalancing(TimeSpan? duration) => _rebalancerGrain.SuspendRebalancing(duration); diff --git a/test/Orleans.Placement.Tests/ActivationRebalancingTests/RebalancingTestBase.cs b/test/Orleans.Placement.Tests/ActivationRebalancingTests/RebalancingTestBase.cs index 0d23c6ac8e3..246f9c8ea2f 100644 --- a/test/Orleans.Placement.Tests/ActivationRebalancingTests/RebalancingTestBase.cs +++ b/test/Orleans.Placement.Tests/ActivationRebalancingTests/RebalancingTestBase.cs @@ -40,9 +40,16 @@ protected static int GetActivationCount(DetailedGrainStatistic[] stats, SiloAddr protected void AddTestActivations(List tasks, SiloAddress silo, int count) { RequestContext.Set(IPlacementDirector.PlacementHintKey, silo); - for (var i = 0; i < count; i++) + try { - tasks.Add(GrainFactory.GetGrain(Guid.NewGuid()).Ping()); + for (var i = 0; i < count; i++) + { + tasks.Add(GrainFactory.GetGrain(Guid.NewGuid()).Ping()); + } + } + finally + { + RequestContext.Remove(IPlacementDirector.PlacementHintKey); } } diff --git a/test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs b/test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs index be73eea653c..8f8134e28f8 100644 --- a/test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs +++ b/test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs @@ -23,6 +23,8 @@ namespace UnitTests.ActivationRebalancingTests; public class StatePreservationRebalancingTests(SPFixture fixture, ITestOutputHelper output) : RebalancingTestBase(fixture, output), IClassFixture { + private static readonly TimeSpan WaitTimeout = TimeSpan.FromSeconds(30); + private const string ErrorMessage = "The rebalancer was not found in any of the 4 silos. " + "Either you have added more silos and not updated this code, " + @@ -32,12 +34,13 @@ public class StatePreservationRebalancingTests(SPFixture fixture, ITestOutputHel public async Task Should_Migrate_And_Preserve_State_When_Hosting_Silo_Dies() { var tasks = new List(); + using var rebalancerEvents = RebalancerDiagnosticObserver.Create(); + var rebalancer = Cluster.Client!.GetGrain(0); + var targetHost = Cluster.Silos[1].SiloAddress; // Move the rebalancer to the first secondary silo, since we will stop it later and we cannot stop // the primary in this test setup. - RequestContext.Set(IPlacementDirector.PlacementHintKey, Cluster.Silos[1].SiloAddress); - await Cluster.Client!.GetGrain(0).Cast().MigrateOnIdle(); - RequestContext.Remove(IPlacementDirector.PlacementHintKey); + await MoveRebalancerToSilo(rebalancer, targetHost); AddTestActivations(tasks, Silo1, 300); AddTestActivations(tasks, Silo2, 30); @@ -46,6 +49,9 @@ public async Task Should_Migrate_And_Preserve_State_When_Hosting_Silo_Dies() await Task.WhenAll(tasks); + rebalancerEvents.Clear(); + await rebalancerEvents.WaitForCycleCountAsync(targetHost, 3, WaitTimeout); + var stats = await MgmtGrain.GetDetailedGrainStatistics(); var initialSilo1Activations = GetActivationCount(stats, Silo1); @@ -60,95 +66,90 @@ public async Task Should_Migrate_And_Preserve_State_When_Hosting_Silo_Dies() $"Silo3: {initialSilo3Activations}\n" + $"Silo4: {initialSilo4Activations}\n"); - var silo1Activations = initialSilo1Activations; - var silo2Activations = initialSilo2Activations; - var silo3Activations = initialSilo3Activations; - var silo4Activations = initialSilo4Activations; - - var rebalancerHostNum = 0; - var index = 0; - - while (index < 6) - { - if (index == 3) - { - (var rebalancerHost, rebalancerHostNum) = await FindRebalancerHost(Silo1); - - OutputHelper.WriteLine($"Cycle {index}: Now stopping Silo{rebalancerHostNum}, which is the host of the rebalancer\n"); + (var rebalancerHost, var rebalancerHostNum) = await FindRebalancerHost(Silo1); + var reportBeforeStop = await rebalancer.GetReport(); + + OutputHelper.WriteLine($"Now stopping Silo{rebalancerHostNum}, which is the host of the rebalancer\n"); + + Assert.Equal(targetHost, rebalancerHost); + Assert.NotEqual(rebalancerHost, Cluster.Silos[0].SiloAddress); + + await Cluster.StopSiloAsync(Cluster.Silos.First(x => x.SiloAddress.Equals(rebalancerHost))); + + var reportAfterStop = await rebalancer.GetReport(); + var newHost = reportAfterStop.Host; + Assert.NotEqual(rebalancerHost, newHost); + Assert.Equal(reportBeforeStop.ClusterImbalance, reportAfterStop.ClusterImbalance); + Assert.Equal(reportBeforeStop.Status, reportAfterStop.Status); + Assert.Equal( + reportBeforeStop.Statistics + .Single(statistic => statistic.SiloAddress.Equals(newHost)) + .AcquiredActivations, + reportAfterStop.Statistics + .Single(statistic => statistic.SiloAddress.Equals(newHost)) + .AcquiredActivations); + Assert.Equal( + reportBeforeStop.Statistics + .Where(statistic => !statistic.SiloAddress.Equals(rebalancerHost)) + .OrderBy(statistic => statistic.SiloAddress.ToString(), StringComparer.Ordinal) + .Select(statistic => ( + statistic.TimeStamp, + statistic.SiloAddress, + statistic.DispersedActivations, + statistic.AcquiredActivations)), + reportAfterStop.Statistics + .Where(statistic => !statistic.SiloAddress.Equals(rebalancerHost)) + .OrderBy(statistic => statistic.SiloAddress.ToString(), StringComparer.Ordinal) + .Select(statistic => ( + statistic.TimeStamp, + statistic.SiloAddress, + statistic.DispersedActivations, + statistic.AcquiredActivations))); + + rebalancerEvents.Clear(); + await rebalancerEvents.WaitForCycleCountAsync(newHost, 3, WaitTimeout); + + stats = await MgmtGrain.GetDetailedGrainStatistics(); + Assert.DoesNotContain(stats, statistic => statistic.SiloAddress.Equals(rebalancerHost)); + + var silo1Activations = GetActivationCount(stats, Silo1); + var silo2Activations = GetActivationCount(stats, Silo2); + var silo3Activations = GetActivationCount(stats, Silo3); + var silo4Activations = GetActivationCount(stats, Silo4); - Assert.NotEqual(rebalancerHost, Cluster.Silos[0].SiloAddress); - await Cluster.StopSiloAsync(Cluster.Silos.First(x => x.SiloAddress.Equals(rebalancerHost))); - } - - await Task.Delay(SPFixture.SessionCyclePeriod); - stats = await MgmtGrain.GetDetailedGrainStatistics(); + OutputHelper.WriteLine( + $"Post-rebalancing activations:\n" + + $"Silo1: {(rebalancerHostNum == 1 ? "DEAD" : silo1Activations)}\n" + + $"Silo2: {(rebalancerHostNum == 2 ? "DEAD" : silo2Activations)}\n" + + $"Silo3: {(rebalancerHostNum == 3 ? "DEAD" : silo3Activations)}\n" + + $"Silo4: {(rebalancerHostNum == 4 ? "DEAD" : silo4Activations)}\n"); - silo1Activations = GetActivationCount(stats, Silo1); - silo2Activations = GetActivationCount(stats, Silo2); - silo3Activations = GetActivationCount(stats, Silo3); - silo4Activations = GetActivationCount(stats, Silo4); + (var finalHost, rebalancerHostNum) = await FindRebalancerHost(newHost); - index++; - } + Assert.Equal(newHost, finalHost); + OutputHelper.WriteLine($"The rebalancer is hosted by Silo{rebalancerHostNum} now"); + } - if (rebalancerHostNum == 1) + private async Task MoveRebalancerToSilo( + IActivationRebalancerWorker rebalancer, + SiloAddress targetHost) + { + if ((await rebalancer.GetReport()).Host.Equals(targetHost)) { - Assert.True(silo2Activations > initialSilo2Activations, - $"Did not expect Silo2 to have less activations than what it started with: " + - $"[{initialSilo2Activations} -> {silo2Activations}]"); - - Assert.True(silo3Activations < initialSilo3Activations, - $"Did not expect Silo3 to have more activations than what it started with: " + - $"[{initialSilo3Activations} -> {silo3Activations}]"); + return; } - else if (rebalancerHostNum == 2) - { - Assert.True(silo3Activations < initialSilo3Activations, - $"Did not expect Silo3 to have more activations than what it started with: " + - $"[{initialSilo3Activations} -> {silo3Activations}]"); - Assert.True(silo4Activations > initialSilo4Activations, - $"Did not expect Silo4 to have less activations than what it started with: " + - $"[{initialSilo4Activations} -> {silo4Activations}]"); - } - else if (rebalancerHostNum == 3) + RequestContext.Set(IPlacementDirector.PlacementHintKey, targetHost); + try { - Assert.True(silo1Activations < initialSilo1Activations, - $"Did not expect Silo1 to have more activations than what it started with: " + - $"[{initialSilo1Activations} -> {silo1Activations}]"); - - Assert.True(silo2Activations > initialSilo2Activations, - $"Did not expect Silo2 to have less activations than what it started with: " + - $"[{initialSilo2Activations} -> {silo2Activations}]"); + await rebalancer.Cast().MigrateOnIdle(); } - else if (rebalancerHostNum == 4) + finally { - Assert.True(silo1Activations < initialSilo1Activations, - $"Did not expect Silo1 to have more activations than what it started with: " + - $"[{initialSilo1Activations} -> {silo1Activations}]"); - - Assert.True(silo2Activations > initialSilo2Activations, - $"Did not expect Silo2 to have less activations than what it started with: " + - $"[{initialSilo2Activations} -> {silo2Activations}]"); + RequestContext.Remove(IPlacementDirector.PlacementHintKey); } - OutputHelper.WriteLine( - $"Post-rebalancing activations ({index} cycles):\n" + - $"Silo1: {(rebalancerHostNum == 1 ? "DEAD" : silo1Activations)}\n" + - $"Silo2: {(rebalancerHostNum == 2 ? "DEAD" : silo2Activations)}\n" + - $"Silo3: {(rebalancerHostNum == 3 ? "DEAD" : silo3Activations)}\n" + - $"Silo4: {(rebalancerHostNum == 4 ? "DEAD" : silo4Activations)}\n"); - - (_, rebalancerHostNum) = await FindRebalancerHost(rebalancerHostNum switch - { - 1 => Silo2, - 2 => Silo3, - 3 => Silo4, - 4 => Silo1, - _ => throw new InvalidOperationException(ErrorMessage) - }); - - OutputHelper.WriteLine($"The rebalancer is hosted by Silo{rebalancerHostNum} now"); + Assert.Equal(targetHost, (await rebalancer.GetReport()).Host); } private async Task<(SiloAddress, int)> FindRebalancerHost(SiloAddress target)