fix shard movements

ywangd · Nov 18, 2024 · 0b03f4e · 0b03f4e
1 parent 2c8feed
commit 0b03f4e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/...org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceShardsAllocator.java b/...org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceShardsAllocator.java
@@ -283,8 +283,10 @@ private static List<MoveAllocationCommand> getMoveCommands(AllocationCommands co
 
     private void setCurrentDesiredBalance(DesiredBalance newDesiredBalance) {
         // Update current desired balance if the old value has not been changed already by master fail-over
+        final AtomicReference<DesiredBalance> oldDesiredBalance = new AtomicReference<>();
         final DesiredBalance updatedDesiredBalance = currentDesiredBalanceRef.updateAndGet(current -> {
             if (current != DesiredBalance.NOT_MASTER) {
+                oldDesiredBalance.set(current);
                 return newDesiredBalance;
             } else {
                 return current;
@@ -300,7 +302,8 @@ private void setCurrentDesiredBalance(DesiredBalance newDesiredBalance) {
             } else {
                 logger.debug("Desired balance updated for [{}]", newDesiredBalance.lastConvergedIndex());
             }
-            computedShardMovements.inc(DesiredBalance.shardMovements(updatedDesiredBalance, newDesiredBalance));
+            assert oldDesiredBalance.get() != null;
+            computedShardMovements.inc(DesiredBalance.shardMovements(oldDesiredBalance.get(), newDesiredBalance));
         } else {
             logger.debug("discard desired balance for [{}]", newDesiredBalance.lastConvergedIndex());
         }