Fix delay allocation (#95921)

It is possible that desired balance is computed before delayed allocations are expired but reconciled after that. 
As a result delay markers are removed but shards are not allocated as their desired locations are not computed yet 
and follow up allocation is not scheduled as all markers might be expired and removed.
This change does not remove delay markers during desired balance reconciliation.

docs/changelog/95921.yaml

@@ -0,0 +1,5 @@
+pr: 95921
+summary: Attempt to fix delay allocation
+area: Allocation
+type: bug
+issues: []

server/src/internalClusterTest/java/org/elasticsearch/cluster/routing/DelayedAllocationIT.java

@@ -13,7 +13,6 @@ import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.ESIntegTestCase;
-import org.elasticsearch.test.junit.annotations.TestLogging;
 
 import java.util.List;
 
@@ -89,10 +88,6 @@ public class DelayedAllocationIT extends ESIntegTestCase {
      * allocation to a very small value, it kicks the allocation of the unassigned shard
      * even though the node it was hosted on will not come back.
      */
-    @TestLogging(
-        value = "org.elasticsearch.cluster.routing.allocation.allocator:TRACE",
-        reason = "https://github.com/elastic/elasticsearch/issues/93470"
-    )
     public void testDelayedAllocationChangeWithSettingTo100ms() throws Exception {
         internalCluster().startNodes(3);
         prepareCreate("test").setSettings(

server/src/main/java/org/elasticsearch/cluster/ClusterModule.java

@@ -29,8 +29,8 @@ import org.elasticsearch.cluster.routing.DelayedAllocationService;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.ShardRoutingRoleStrategy;
 import org.elasticsearch.cluster.routing.allocation.AllocationService;
+import org.elasticsearch.cluster.routing.allocation.AllocationService.RerouteStrategy;
 import org.elasticsearch.cluster.routing.allocation.ExistingShardsAllocator;
-import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.WriteLoadForecaster;
 import org.elasticsearch.cluster.routing.allocation.allocator.BalancedShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator;
@@ -89,7 +89,6 @@ import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -177,8 +176,8 @@ public class ClusterModule extends AbstractModule {
         };
     }
 
-    private ClusterState reconcile(ClusterState clusterState, Consumer<RoutingAllocation> routingAllocationConsumer) {
-        return allocationService.executeWithRoutingAllocation(clusterState, "reconcile-desired-balance", routingAllocationConsumer);
+    private ClusterState reconcile(ClusterState clusterState, RerouteStrategy rerouteStrategy) {
+        return allocationService.executeWithRoutingAllocation(clusterState, "reconcile-desired-balance", rerouteStrategy);
     }
 
     public static List<Entry> getNamedWriteables() {

server/src/main/java/org/elasticsearch/cluster/routing/DelayedAllocationService.java

@@ -41,7 +41,7 @@ import static org.elasticsearch.cluster.routing.allocation.allocator.AllocationA
  * This class is responsible for choosing the next (closest) delay expiration of a
  * delayed shard to schedule a reroute to remove the delay marker.
  * The actual removal of the delay marker happens in
- * {@link AllocationService#removeDelayMarkers(RoutingAllocation)}, triggering yet
+ * {@link AllocationService.RerouteStrategy#removeDelayMarkers(RoutingAllocation)}, triggering yet
  * another cluster change event.
  */
 public class DelayedAllocationService extends AbstractLifecycleComponent implements ClusterStateListener {

server/src/main/java/org/elasticsearch/cluster/routing/allocation/AllocationService.java

@@ -53,7 +53,6 @@ import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -334,43 +333,6 @@ public class AllocationService {
         }
     }
 
-    /**
-     * Removes delay markers from unassigned shards based on current time stamp.
-     */
-    private void removeDelayMarkers(RoutingAllocation allocation) {
-        final RoutingNodes.UnassignedShards.UnassignedIterator unassignedIterator = allocation.routingNodes().unassigned().iterator();
-        final Metadata metadata = allocation.metadata();
-        while (unassignedIterator.hasNext()) {
-            ShardRouting shardRouting = unassignedIterator.next();
-            UnassignedInfo unassignedInfo = shardRouting.unassignedInfo();
-            if (unassignedInfo.isDelayed()) {
-                final long newComputedLeftDelayNanos = unassignedInfo.getRemainingDelay(
-                    allocation.getCurrentNanoTime(),
-                    metadata.getIndexSafe(shardRouting.index()).getSettings(),
-                    metadata.nodeShutdowns()
-                );
-                if (newComputedLeftDelayNanos == 0) {
-                    unassignedIterator.updateUnassigned(
-                        new UnassignedInfo(
-                            unassignedInfo.getReason(),
-                            unassignedInfo.getMessage(),
-                            unassignedInfo.getFailure(),
-                            unassignedInfo.getNumFailedAllocations(),
-                            unassignedInfo.getUnassignedTimeInNanos(),
-                            unassignedInfo.getUnassignedTimeInMillis(),
-                            false,
-                            unassignedInfo.getLastAllocationStatus(),
-                            unassignedInfo.getFailedNodeIds(),
-                            unassignedInfo.getLastAllocatedNodeId()
-                        ),
-                        shardRouting.recoverySource(),
-                        allocation.changes()
-                    );
-                }
-            }
-        }
-    }
-
     /**
      * Internal helper to cap the number of elements in a potentially long list for logging.
      *
@@ -443,20 +405,62 @@ public class AllocationService {
      *
      * @return an updated cluster state, or the same instance that was passed as an argument if no changes were made.
      */
-    public ClusterState executeWithRoutingAllocation(
-        ClusterState clusterState,
-        String reason,
-        Consumer<RoutingAllocation> routingAllocationConsumer
-    ) {
+    public ClusterState executeWithRoutingAllocation(ClusterState clusterState, String reason, RerouteStrategy rerouteStrategy) {
         ClusterState fixedClusterState = adaptAutoExpandReplicas(clusterState);
         RoutingAllocation allocation = createRoutingAllocation(fixedClusterState, currentNanoTime());
-        reroute(allocation, routingAllocationConsumer);
+        reroute(allocation, rerouteStrategy);
         if (fixedClusterState == clusterState && allocation.routingNodesChanged() == false) {
             return clusterState;
         }
         return buildResultAndLogHealthChange(clusterState, allocation, reason);
     }
 
+    @FunctionalInterface
+    public interface RerouteStrategy {
+
+        /**
+         * Removes delay markers from unassigned shards based on current time stamp.
+         */
+        default void removeDelayMarkers(RoutingAllocation allocation) {
+            final RoutingNodes.UnassignedShards.UnassignedIterator unassignedIterator = allocation.routingNodes().unassigned().iterator();
+            final Metadata metadata = allocation.metadata();
+            while (unassignedIterator.hasNext()) {
+                ShardRouting shardRouting = unassignedIterator.next();
+                UnassignedInfo unassignedInfo = shardRouting.unassignedInfo();
+                if (unassignedInfo.isDelayed()) {
+                    final long newComputedLeftDelayNanos = unassignedInfo.getRemainingDelay(
+                        allocation.getCurrentNanoTime(),
+                        metadata.getIndexSafe(shardRouting.index()).getSettings(),
+                        metadata.nodeShutdowns()
+                    );
+                    if (newComputedLeftDelayNanos == 0) {
+                        unassignedIterator.updateUnassigned(
+                            new UnassignedInfo(
+                                unassignedInfo.getReason(),
+                                unassignedInfo.getMessage(),
+                                unassignedInfo.getFailure(),
+                                unassignedInfo.getNumFailedAllocations(),
+                                unassignedInfo.getUnassignedTimeInNanos(),
+                                unassignedInfo.getUnassignedTimeInMillis(),
+                                false,
+                                unassignedInfo.getLastAllocationStatus(),
+                                unassignedInfo.getFailedNodeIds(),
+                                unassignedInfo.getLastAllocatedNodeId()
+                            ),
+                            shardRouting.recoverySource(),
+                            allocation.changes()
+                        );
+                    }
+                }
+            }
+        }
+
+        /**
+         * Generic action to be executed on preconfigured allocation
+         */
+        void execute(RoutingAllocation allocation);
+    }
+
     private static void logClusterHealthStateChange(final ClusterState previousState, final ClusterState newState, String reason) {
         ClusterHealthStatus previousHealth = getHealthStatus(previousState);
         ClusterHealthStatus currentHealth = getHealthStatus(newState);
@@ -515,16 +519,14 @@ public class AllocationService {
         return false;
     }
 
-    private void reroute(RoutingAllocation allocation, Consumer<RoutingAllocation> routingAllocationConsumer) {
+    private void reroute(RoutingAllocation allocation, RerouteStrategy rerouteStrategy) {
         assert hasDeadNodes(allocation) == false : "dead nodes should be explicitly cleaned up. See disassociateDeadNodes";
         assert AutoExpandReplicas.getAutoExpandReplicaChanges(allocation.metadata(), () -> allocation).isEmpty()
             : "auto-expand replicas out of sync with number of nodes in the cluster";
         assert assertInitialized();
-
-        removeDelayMarkers(allocation);
-
-        allocateExistingUnassignedShards(allocation);  // try to allocate existing shard copies first
-        routingAllocationConsumer.accept(allocation);
+        rerouteStrategy.removeDelayMarkers(allocation);
+        allocateExistingUnassignedShards(allocation); // try to allocate existing shard copies first
+        rerouteStrategy.execute(allocation);
         assert RoutingNodes.assertShardStats(allocation.routingNodes());
     }
 

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceShardsAllocator.java

@@ -15,6 +15,7 @@ import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ClusterStateTaskExecutor;
 import org.elasticsearch.cluster.ClusterStateTaskListener;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.allocation.AllocationService.RerouteStrategy;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.RoutingExplanations;
 import org.elasticsearch.cluster.routing.allocation.ShardAllocationDecision;
@@ -36,7 +37,6 @@ import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.function.Consumer;
 
 /**
  * A {@link ShardsAllocator} which asynchronously refreshes the desired balance held by the {@link DesiredBalanceComputer} and then takes
@@ -70,7 +70,7 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
 
     @FunctionalInterface
     public interface DesiredBalanceReconcilerAction {
-        ClusterState apply(ClusterState clusterState, Consumer<RoutingAllocation> routingAllocationAction);
+        ClusterState apply(ClusterState clusterState, RerouteStrategy rerouteStrategy);
     }
 
     public DesiredBalanceShardsAllocator(
@@ -242,6 +242,24 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
         }
     }
 
+    private RerouteStrategy createReconcileAllocationAction(DesiredBalance desiredBalance) {
+        return new RerouteStrategy() {
+            @Override
+            public void removeDelayMarkers(RoutingAllocation allocation) {
+                // it is possible that desired balance is computed before some delayed allocations are expired but reconciled after.
+                // If delayed markers are removed during reconciliation then
+                // * shards are not assigned anyway as balance is not computed for them
+                // * followup reroute is not scheduled to allocate them
+                // for this reason we should keep delay markers during reconciliation
+            }
+
+            @Override
+            public void execute(RoutingAllocation allocation) {
+                reconcile(desiredBalance, allocation);
+            }
+        };
+    }
+
     public DesiredBalance getDesiredBalance() {
         return currentDesiredBalance;
     }
@@ -312,7 +330,7 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
             try (var ignored = batchExecutionContext.dropHeadersContext()) {
                 var newState = reconciler.apply(
                     batchExecutionContext.initialState(),
-                    routingAllocation -> reconcile(latest.getTask().desiredBalance, routingAllocation)
+                    createReconcileAllocationAction(latest.getTask().desiredBalance)
                 );
                 latest.success(() -> queue.complete(latest.getTask().desiredBalance.lastConvergedIndex()));
                 return newState;

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceShardsAllocatorTests.java

@@ -22,10 +22,13 @@ import org.elasticsearch.cluster.block.ClusterBlocks;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.Metadata;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.cluster.routing.IndexRoutingTable;
 import org.elasticsearch.cluster.routing.RoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.UnassignedInfo;
 import org.elasticsearch.cluster.routing.allocation.AllocationService;
+import org.elasticsearch.cluster.routing.allocation.AllocationService.RerouteStrategy;
 import org.elasticsearch.cluster.routing.allocation.ExistingShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.ShardAllocationDecision;
@@ -39,6 +42,7 @@ import org.elasticsearch.common.UUIDs;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.DeterministicTaskQueue;
 import org.elasticsearch.common.util.concurrent.PrioritizedEsThreadPoolExecutor;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.gateway.GatewayAllocator;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.snapshots.SnapshotShardSizeInfo;
@@ -48,6 +52,7 @@ import org.elasticsearch.threadpool.TestThreadPool;
 import java.util.List;
 import java.util.Map;
 import java.util.Queue;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -57,6 +62,9 @@ import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 
+import static org.elasticsearch.cluster.routing.AllocationId.newInitializing;
+import static org.elasticsearch.cluster.routing.TestShardRouting.newShardRouting;
+import static org.elasticsearch.cluster.routing.UnassignedInfo.INDEX_DELAYED_NODE_LEFT_TIMEOUT_SETTING;
 import static org.elasticsearch.common.settings.ClusterSettings.createBuiltInClusterSettings;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItem;
@@ -118,9 +126,9 @@ public class DesiredBalanceShardsAllocatorTests extends ESAllocationTestCase {
         var clusterSettings = createBuiltInClusterSettings(settings);
         var clusterService = new ClusterService(
             settings,
-            createBuiltInClusterSettings(settings),
+            clusterSettings,
             new FakeThreadPoolMasterService(LOCAL_NODE_ID, threadPool, deterministicTaskQueue::scheduleNow),
-            new ClusterApplierService(LOCAL_NODE_ID, Settings.EMPTY, clusterSettings, threadPool) {
+            new ClusterApplierService(LOCAL_NODE_ID, settings, clusterSettings, threadPool) {
                 @Override
                 protected PrioritizedEsThreadPoolExecutor createThreadPoolExecutor() {
                     return deterministicTaskQueue.getPrioritizedEsThreadPoolExecutor();
@@ -137,7 +145,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESAllocationTestCase {
         var allocationServiceRef = new SetOnce<AllocationService>();
         var reconcileAction = new DesiredBalanceReconcilerAction() {
             @Override
-            public ClusterState apply(ClusterState clusterState, Consumer<RoutingAllocation> routingAllocationAction) {
+            public ClusterState apply(ClusterState clusterState, RerouteStrategy routingAllocationAction) {
                 return allocationServiceRef.get().executeWithRoutingAllocation(clusterState, "reconcile", routingAllocationAction);
             }
         };
@@ -194,6 +202,112 @@ public class DesiredBalanceShardsAllocatorTests extends ESAllocationTestCase {
         }
     }
 
+    public void testShouldNotRemoveAllocationDelayMarkersOnReconcile() {
+
+        var localNode = newNode(LOCAL_NODE_ID);
+        var otherNode = newNode(OTHER_NODE_ID);
+
+        var unassignedTimeNanos = System.nanoTime();
+        var computationTime = unassignedTimeNanos;
+        var reconciliationTime = unassignedTimeNanos + INDEX_DELAYED_NODE_LEFT_TIMEOUT_SETTING.get(Settings.EMPTY).nanos();
+
+        var delayedUnasssignedInfo = new UnassignedInfo(
+            UnassignedInfo.Reason.NODE_RESTARTING,
+            null,
+            null,
+            0,
+            unassignedTimeNanos,
+            TimeValue.nsecToMSec(unassignedTimeNanos),
+            true,
+            UnassignedInfo.AllocationStatus.NO_ATTEMPT,
+            Set.of(),
+            "node-3"
+        );
+
+        var inSyncAllocationId = UUIDs.randomBase64UUID();
+        var index = IndexMetadata.builder("test")
+            .settings(indexSettings(Version.CURRENT, 1, 1))
+            .putInSyncAllocationIds(0, Set.of(inSyncAllocationId))
+            .build();
+        var shardId = new ShardId(index.getIndex(), 0);
+        var indexRoutingTable = IndexRoutingTable.builder(index.getIndex())
+            .addShard(newShardRouting(shardId, LOCAL_NODE_ID, null, true, ShardRoutingState.STARTED, newInitializing(inSyncAllocationId)))
+            .addShard(newShardRouting(shardId, null, null, false, ShardRoutingState.UNASSIGNED, delayedUnasssignedInfo))
+            .build();
+
+        var initialState = ClusterState.builder(new ClusterName(ClusterServiceUtils.class.getSimpleName()))
+            .nodes(DiscoveryNodes.builder().add(localNode).add(otherNode).localNodeId(localNode.getId()).masterNodeId(localNode.getId()))
+            .metadata(Metadata.builder().put(index, false).build())
+            .routingTable(RoutingTable.builder(TestShardRoutingRoleStrategies.DEFAULT_ROLE_ONLY).add(indexRoutingTable).build())
+            .blocks(ClusterBlocks.EMPTY_CLUSTER_BLOCK)
+            .build();
+
+        var threadPool = new TestThreadPool(getTestName());
+        var clusterService = ClusterServiceUtils.createClusterService(initialState, threadPool);
+        var allocationServiceRef = new SetOnce<AllocationService>();
+        var reconciledStateRef = new AtomicReference<ClusterState>();
+        var reconcileAction = new DesiredBalanceReconcilerAction() {
+            @Override
+            public ClusterState apply(ClusterState clusterState, RerouteStrategy routingAllocationAction) {
+                ClusterState reconciled = allocationServiceRef.get()
+                    .executeWithRoutingAllocation(clusterState, "reconcile", routingAllocationAction);
+                reconciledStateRef.set(reconciled);
+                return reconciled;
+            }
+        };
+
+        var clusterSettings = createBuiltInClusterSettings();
+        var desiredBalanceShardsAllocator = new DesiredBalanceShardsAllocator(
+            clusterSettings,
+            createShardsAllocator(),
+            threadPool,
+            clusterService,
+            reconcileAction
+        );
+        var allocationService = new AllocationService(
+            new AllocationDeciders(List.of()),
+            createGatewayAllocator(
+                (allocation, unassignedAllocationHandler) -> unassignedAllocationHandler.removeAndIgnore(
+                    UnassignedInfo.AllocationStatus.NO_ATTEMPT,
+                    allocation.changes()
+                )
+            ),
+            desiredBalanceShardsAllocator,
+            () -> ClusterInfo.EMPTY,
+            () -> SnapshotShardSizeInfo.EMPTY,
+            TestShardRoutingRoleStrategies.DEFAULT_ROLE_ONLY
+        ) {
+
+            int call = 0;
+            long[] timeToReturn = new long[] { computationTime, reconciliationTime };
+
+            @Override
+            protected long currentNanoTime() {
+                return timeToReturn[call++];
+            }
+        };
+        allocationServiceRef.set(allocationService);
+
+        try {
+            rerouteAndWait(allocationService, initialState, "test");
+
+            var reconciledState = reconciledStateRef.get();
+
+            // Desired balance computation could be performed _before_ delayed allocation is expired
+            // while corresponding reconciliation might happen _after_.
+            // In such case reconciliation will not allocate delayed shard as its balance is not computed yet
+            // and must NOT clear delayed flag so that a followup reroute is scheduled for the shard.
+            var unassigned = reconciledState.getRoutingNodes().unassigned();
+            assertThat(unassigned.size(), equalTo(1));
+            var unassignedShard = unassigned.iterator().next();
+            assertThat(unassignedShard.unassignedInfo().isDelayed(), equalTo(true));
+
+        } finally {
+            clusterService.close();
+            terminate(threadPool);
+        }
+    }
+
     public void testCallListenersOnlyAfterProducingFreshInput() throws InterruptedException {
 
         var reconciliations = new AtomicInteger(0);
@@ -210,7 +324,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESAllocationTestCase {
         var allocationServiceRef = new SetOnce<AllocationService>();
         var reconcileAction = new DesiredBalanceReconcilerAction() {
             @Override
-            public ClusterState apply(ClusterState clusterState, Consumer<RoutingAllocation> routingAllocationAction) {
+            public ClusterState apply(ClusterState clusterState, RerouteStrategy routingAllocationAction) {
                 reconciliations.incrementAndGet();
                 return allocationServiceRef.get().executeWithRoutingAllocation(clusterState, "reconcile", routingAllocationAction);
             }
@@ -312,7 +426,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESAllocationTestCase {
         var allocationServiceRef = new SetOnce<AllocationService>();
         var reconcileAction = new DesiredBalanceReconcilerAction() {
             @Override
-            public ClusterState apply(ClusterState clusterState, Consumer<RoutingAllocation> routingAllocationAction) {
+            public ClusterState apply(ClusterState clusterState, RerouteStrategy routingAllocationAction) {
                 return allocationServiceRef.get().executeWithRoutingAllocation(clusterState, "reconcile", routingAllocationAction);
             }
         };
@@ -430,7 +544,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESAllocationTestCase {
             threadPool,
             clusterService,
             desiredBalanceComputer,
-            (reconcilerClusterState, routingAllocationAction) -> reconcilerClusterState
+            (reconcilerClusterState, rerouteStrategy) -> reconcilerClusterState
         );
 
         var service = createAllocationService(desiredBalanceShardsAllocator, createGatewayAllocator());
@@ -482,7 +596,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESAllocationTestCase {
             threadPool,
             clusterService,
             desiredBalanceComputer,
-            (reconcilerClusterState, routingAllocationAction) -> reconcilerClusterState
+            (reconcilerClusterState, rerouteStrategy) -> reconcilerClusterState
         );
 
         var service = createAllocationService(desiredBalanceShardsAllocator, createGatewayAllocator());