Skip to content

fix(runtime): preserve rebalancer state on shutdown #10573

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ReubenBond merged 2 commits into from
Aug 13, 2026
Merged

fix(runtime): preserve rebalancer state on shutdown #10573

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c103aa3
fix(runtime): preserve rebalancer state on shutdown
ReubenBond Aug 13, 2026
105d30a
test(rebalancing): baseline cycle observations
ReubenBond Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 20 additions & 11 deletions src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancerMonitor.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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);

Expand Down
11 changes: 9 additions & 2 deletions test/Orleans.Placement.Tests/ActivationRebalancingTests/RebalancingTestBase.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,16 @@ protected static int GetActivationCount(DetailedGrainStatistic[] stats, SiloAddr
protected void AddTestActivations(List<Task> tasks, SiloAddress silo, int count)
{
RequestContext.Set(IPlacementDirector.PlacementHintKey, silo);
for (var i = 0; i < count; i++)
try
{
tasks.Add(GrainFactory.GetGrain<IRebalancingTestGrain>(Guid.NewGuid()).Ping());
for (var i = 0; i < count; i++)
{
tasks.Add(GrainFactory.GetGrain<IRebalancingTestGrain>(Guid.NewGuid()).Ping());
}
}
finally
{
RequestContext.Remove(IPlacementDirector.PlacementHintKey);
}
}

Expand Down
161 changes: 81 additions & 80 deletions test/Orleans.Placement.Tests/ActivationRebalancingTests/StatePreservationRebalancingTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ namespace UnitTests.ActivationRebalancingTests;
public class StatePreservationRebalancingTests(SPFixture fixture, ITestOutputHelper output)
: RebalancingTestBase<SPFixture>(fixture, output), IClassFixture<SPFixture>
{
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, " +
Expand All @@ -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<Task>();
using var rebalancerEvents = RebalancerDiagnosticObserver.Create();
var rebalancer = Cluster.Client!.GetGrain<IActivationRebalancerWorker>(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<IActivationRebalancerWorker>(0).Cast<IGrainManagementExtension>().MigrateOnIdle();
RequestContext.Remove(IPlacementDirector.PlacementHintKey);
await MoveRebalancerToSilo(rebalancer, targetHost);

AddTestActivations(tasks, Silo1, 300);
AddTestActivations(tasks, Silo2, 30);
Expand All @@ -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);

Comment thread
ReubenBond marked this conversation as resolved.
var stats = await MgmtGrain.GetDetailedGrainStatistics();

var initialSilo1Activations = GetActivationCount(stats, Silo1);
Expand All @@ -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<IGrainManagementExtension>().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)
Expand Down
Loading