From 169db61b2b61010dce80b8f7391fbc2d98169a25 Mon Sep 17 00:00:00 2001
From: ReubenBond <rebond@microsoft.com>
Date: Mon, 20 May 2024 21:33:59 -0700
Subject: [PATCH] wip

---
 .../Options/ActiveRebalancingOptions.cs       | 10 +---------
 .../Rebalancing/ActivationRebalancer.cs       | 20 ++++++++++---------
 .../Placement/Rebalancing/MaxHeap.cs          | 12 +++++++++++
 3 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/src/Orleans.Runtime/Configuration/Options/ActiveRebalancingOptions.cs b/src/Orleans.Runtime/Configuration/Options/ActiveRebalancingOptions.cs
index 3c3596a129..9491249554 100644
--- a/src/Orleans.Runtime/Configuration/Options/ActiveRebalancingOptions.cs
+++ b/src/Orleans.Runtime/Configuration/Options/ActiveRebalancingOptions.cs
@@ -20,15 +20,7 @@ public sealed class ActiveRebalancingOptions
     /// In order to preserve memory, the most heaviest links are recorded in a probabilistic way, so there is an inherent error associated with that.
     /// That error is inversely proportional to this value, so values under 100 are not recommended. If you notice that the system is not converging fast enough, do consider increasing this number.
     /// </remarks>
-    public uint MaxEdgeCount { get; set; } =
-
-
-
-        10 * 
-
-
-
-        DEFAULT_MAX_EDGE_COUNT;
+    public uint MaxEdgeCount { get; set; } = DEFAULT_MAX_EDGE_COUNT;
 
     /// <summary>
     /// The default value of <see cref="MaxEdgeCount"/>.
diff --git a/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancer.cs b/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancer.cs
index 78a9a5d106..a483a0496f 100644
--- a/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancer.cs
+++ b/src/Orleans.Runtime/Placement/Rebalancing/ActivationRebalancer.cs
@@ -211,7 +211,8 @@ public async ValueTask<AcceptExchangeResponse> AcceptExchangeRequest(AcceptExcha
             var remoteActivations = request.ActivationCountSnapshot;
             var localActivations = GetLocalActivationCount();
 
-            var imbalance = CalculateImbalance(remoteActivations, localActivations);
+            var initialImbalance =  CalculateImbalance(remoteActivations, localActivations);
+            int imbalance = initialImbalance;
             _logger.LogInformation("Imbalance is {Imbalance} (remote: {RemoteCount} vs local {LocalCount})", imbalance, remoteActivations, localActivations);
 
             var (localHeap, remoteHeap) = CreateCandidateHeaps(localSet, remoteSet);
@@ -318,22 +319,23 @@ bool TryMigrateRemoteToLocal()
 
             bool TryMigrateCore(MaxHeap<CandidateVertexHeapElement> sourceHeap, int localDelta, int remoteDelta, [NotNullWhen(true)] out CandidateVertexHeapElement? chosenVertex)
             {
-                chosenVertex = null;
-                if (sourceHeap.Count == 0)
+                var anticipatedImbalance = CalculateImbalance(localActivations + localDelta, remoteActivations + remoteDelta);
+                if (anticipatedImbalance >= initialImbalance && !_toleranceRule.IsSatisfiedBy((uint)anticipatedImbalance))
                 {
+                    // Taking from this heap would not improve imbalance.
+                    chosenVertex = null;
                     return false;
                 }
 
-                var anticipatedImbalance = CalculateImbalance(localActivations + localDelta, remoteActivations + remoteDelta);
-                if (anticipatedImbalance > imbalance && !_toleranceRule.IsSatisfiedBy((uint)anticipatedImbalance))
+                if (!sourceHeap.TryPop(out chosenVertex))
                 {
+                    // Heap is empty.
                     return false;
                 }
 
-                chosenVertex = sourceHeap.Pop();
-                if (chosenVertex.AccumulatedTransferScore <= 0)
+                if (chosenVertex.AccumulatedTransferScore < 0)
                 {
-                    // If it got affected by a previous run, and the score is not positive, simply pop and ignore it.
+                    // If it got affected by a previous run, and the score is negative, simply pop and ignore it.
                     return false;
                 }
 
@@ -355,7 +357,7 @@ bool TryMigrateCore(MaxHeap<CandidateVertexHeapElement> sourceHeap, int localDel
         }
     }
 
-    private static int CalculateImbalance(int left, int right) => Math.Abs(Math.Abs(left) - Math.Abs(right));
+    private static int CalculateImbalance(int left, int right) => (int)Math.Abs(Math.Abs((long)left) - Math.Abs((long)right));
     private static (MaxHeap<CandidateVertexHeapElement> Local, MaxHeap<CandidateVertexHeapElement> Remote) CreateCandidateHeaps(List<CandidateVertex> local, ImmutableArray<CandidateVertex> remote)
     {
         Dictionary<GrainId, CandidateVertex> sourceIndex = new(local.Count + remote.Length);
diff --git a/src/Orleans.Runtime/Placement/Rebalancing/MaxHeap.cs b/src/Orleans.Runtime/Placement/Rebalancing/MaxHeap.cs
index c88bc8c19e..270dac03e2 100644
--- a/src/Orleans.Runtime/Placement/Rebalancing/MaxHeap.cs
+++ b/src/Orleans.Runtime/Placement/Rebalancing/MaxHeap.cs
@@ -143,6 +143,18 @@ public TElement Peek()
         return _nodes[0]!;
     }
 
+    public bool TryPop([NotNullWhen(true)] out TElement value)
+    {
+        if (_size > 0)
+        {
+            value = Pop();
+            return true;
+        }
+
+        value = default!;
+        return false;
+    }
+
     /// <summary>
     ///  Removes and returns the maximal element from the <see cref="MaxHeap{TElement}"/>.
     /// </summary>