From c103aa3a5eb76f39c8b8938796165d5e75795eb3 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Wed, 12 Aug 2026 21:20:50 -0700 Subject: [PATCH 1/2] fix(runtime): preserve rebalancer state on shutdown Migrate and await the activation rebalancer worker before ordinary grain deactivation begins during graceful silo shutdown. Synchronize the regression test on migration and cycle completion and verify preserved worker state. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../ActivationRebalancerMonitor.cs | 31 ++-- .../RebalancingTestBase.cs | 11 +- .../StatePreservationRebalancingTests.cs | 160 +++++++++--------- 3 files changed, 109 insertions(+), 93 deletions(-) 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..6a051f5ec7d 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,8 @@ public async Task Should_Migrate_And_Preserve_State_When_Hosting_Silo_Dies() await Task.WhenAll(tasks); + await rebalancerEvents.WaitForCycleCountAsync(targetHost, 3, WaitTimeout); + var stats = await MgmtGrain.GetDetailedGrainStatistics(); var initialSilo1Activations = GetActivationCount(stats, Silo1); @@ -60,95 +65,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) From 105d30a07f610a801e64bcf67e41da98416c7424 Mon Sep 17 00:00:00 2001 From: Reuben Bond Date: Thu, 13 Aug 2026 07:08:48 -0700 Subject: [PATCH 2/2] test(rebalancing): baseline cycle observations Clear previously observed cycle events after the test load is created so the initial barrier only counts cycles which process that load. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 8fc23ff9-c5c7-435c-91c5-e23f9b6dc452 --- .../StatePreservationRebalancingTests.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs b/test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs index 6a051f5ec7d..8f8134e28f8 100644 --- a/test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs +++ b/test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs @@ -49,6 +49,7 @@ 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();