Undesired allocation tracking improvements (#137940)

Co-authored-by: Yang Wang <[email protected]>

Author: nicktindall

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

@@ -9,6 +9,8 @@
 
 package org.elasticsearch.cluster.routing.allocation.allocator;
 
+import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
@@ -17,6 +19,7 @@
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.time.TimeProvider;
 import org.elasticsearch.common.util.FeatureFlag;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.logging.LogManager;
@@ -82,6 +85,7 @@ public class UndesiredAllocationsTracker {
     private final FrequencyCappedAction undesiredAllocationDurationLogInterval;
     private volatile TimeValue undesiredAllocationDurationLoggingThreshold;
     private volatile int maxUndesiredAllocationsToTrack;
+    private boolean missingAllocationAssertionsEnabled = true;
 
     UndesiredAllocationsTracker(ClusterSettings clusterSettings, TimeProvider timeProvider) {
         this.timeProvider = timeProvider;
@@ -161,6 +165,14 @@ public void maybeLogUndesiredShardsWarning(
         }
     }
 
+    private boolean shardTierMatchesNodeTier(ShardRouting shardRouting, DiscoveryNode discoveryNode) {
+        return switch (shardRouting.role()) {
+            case INDEX_ONLY -> discoveryNode.getRoles().contains(DiscoveryNodeRole.INDEX_ROLE);
+            case SEARCH_ONLY -> discoveryNode.getRoles().contains(DiscoveryNodeRole.SEARCH_ROLE);
+            default -> true;
+        };
+    }
+
     private void logDecisionsForUndesiredShardsOverThreshold(
         RoutingNodes routingNodes,
         RoutingAllocation routingAllocation,
@@ -199,10 +211,21 @@ private void logUndesiredShardDetails(
         allocation.setDebugMode(RoutingAllocation.DebugMode.EXCLUDE_YES_DECISIONS);
         try {
             final var assignment = desiredBalance.getAssignment(shardRouting.shardId());
-            logger.warn("Shard {} has been in an undesired allocation for {}", shardRouting.shardId(), undesiredDuration);
-            for (final var nodeId : assignment.nodeIds()) {
-                final var decision = allocation.deciders().canAllocate(shardRouting, routingNodes.node(nodeId), allocation);
-                logger.warn("Shard {} allocation decision for node [{}]: {}", shardRouting.shardId(), nodeId, decision);
+            if (assignment != null) {
+                logger.warn("Shard {} has been in an undesired allocation for {}", shardRouting.shardId(), undesiredDuration);
+                for (final var nodeId : assignment.nodeIds()) {
+                    if (allocation.nodes().nodeExists(nodeId)) {
+                        if (shardTierMatchesNodeTier(shardRouting, allocation.nodes().get(nodeId))) {
+                            final var decision = allocation.deciders().canAllocate(shardRouting, routingNodes.node(nodeId), allocation);
+                            logger.warn("Shard {} allocation decision for node [{}]: {}", shardRouting.shardId(), nodeId, decision);
+                        }
+                    } else {
+                        logger.warn("Shard {} desired node [{}] has left the cluster", shardRouting.shardId(), nodeId);
+                    }
+                }
+            } else {
+                assert missingAllocationAssertionsEnabled == false
+                    : "Shard " + shardRouting + " was missing an assignment, this shouldn't be possible. " + desiredBalance;
             }
         } finally {
             allocation.setDebugMode(originalDebugMode);
@@ -239,4 +262,10 @@ Map<String, UndesiredAllocation> getUndesiredAllocations() {
      * @param undesiredSince The timestamp when the shard was first observed in an undesired allocation
      */
     record UndesiredAllocation(ShardId shardId, long undesiredSince) {}
+
+    // Exposed for testing
+    public Releasable disableMissingAllocationAssertions() {
+        missingAllocationAssertionsEnabled = false;
+        return () -> missingAllocationAssertionsEnabled = true;
+    }
 }