Small cleanups to Allocation Performance (#89378)

Two fixes:
1. Looking up `Custom` values over and over for every shard incurs a measurable cost.
This removes that cost for desired nodes and node shutdown metadata.
2. Node shutdown metadata logic wasn't inlining nicely because of the wrapped map.
No need to be as complicated as we were in many spots, use a simple immutable map
for all operations and remove a bunch of branching.

Author: original-brownbear

server/src/main/java/org/elasticsearch/cluster/metadata/Metadata.java

@@ -1024,9 +1024,7 @@ public Map<String, DataStreamAlias> dataStreamAliases() {
     }
 
     public Map<String, SingleNodeShutdownMetadata> nodeShutdowns() {
-        return Optional.ofNullable((NodesShutdownMetadata) this.custom(NodesShutdownMetadata.TYPE))
-            .map(NodesShutdownMetadata::getAllNodeMetadataMap)
-            .orElse(Collections.emptyMap());
+        return this.custom(NodesShutdownMetadata.TYPE, NodesShutdownMetadata.EMPTY).getAllNodeMetadataMap();
     }
 
     public Map<String, Custom> customs() {

server/src/main/java/org/elasticsearch/cluster/metadata/NodesShutdownMetadata.java

@@ -22,7 +22,6 @@
 import org.elasticsearch.xcontent.XContentParser;
 
 import java.io.IOException;
-import java.util.Collections;
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.List;
@@ -93,11 +92,11 @@ public static NodesShutdownMetadata getShutdownsOrEmpty(final ClusterState state
     private final Map<String, SingleNodeShutdownMetadata> nodes;
 
     public NodesShutdownMetadata(Map<String, SingleNodeShutdownMetadata> nodes) {
-        this.nodes = Collections.unmodifiableMap(nodes);
+        this.nodes = Map.copyOf(nodes);
     }
 
     public NodesShutdownMetadata(StreamInput in) throws IOException {
-        this(in.readMap(StreamInput::readString, SingleNodeShutdownMetadata::new));
+        this(in.readImmutableMap(StreamInput::readString, SingleNodeShutdownMetadata::new));
     }
 
     @Override

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

@@ -26,7 +26,6 @@
 import org.elasticsearch.snapshots.RestoreService.RestoreInProgressUpdater;
 import org.elasticsearch.snapshots.SnapshotShardSizeInfo;
 
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -73,6 +72,11 @@ public class RoutingAllocation {
 
     private final Map<String, SingleNodeShutdownMetadata> nodeReplacementTargets;
 
+    private final Map<String, SingleNodeShutdownMetadata> nodeShutdowns;
+
+    @Nullable
+    private final DesiredNodes desiredNodes;
+
     public RoutingAllocation(
         AllocationDeciders deciders,
         ClusterState clusterState,
@@ -106,13 +110,15 @@ public RoutingAllocation(
         this.clusterInfo = clusterInfo;
         this.shardSizeInfo = shardSizeInfo;
         this.currentNanoTime = currentNanoTime;
+        this.nodeShutdowns = clusterState.metadata().nodeShutdowns();
         Map<String, SingleNodeShutdownMetadata> targetNameToShutdown = new HashMap<>();
-        for (SingleNodeShutdownMetadata shutdown : clusterState.metadata().nodeShutdowns().values()) {
+        for (SingleNodeShutdownMetadata shutdown : nodeShutdowns.values()) {
             if (shutdown.getType() == SingleNodeShutdownMetadata.Type.REPLACE) {
                 targetNameToShutdown.put(shutdown.getTargetNodeName(), shutdown);
             }
         }
-        this.nodeReplacementTargets = Collections.unmodifiableMap(targetNameToShutdown);
+        this.nodeReplacementTargets = Map.copyOf(targetNameToShutdown);
+        this.desiredNodes = DesiredNodes.latestFromClusterState(clusterState);
     }
 
     /** returns the nano time captured at the beginning of the allocation. used to make sure all time based decisions are aligned */
@@ -173,14 +179,14 @@ public SnapshotShardSizeInfo snapshotShardSizeInfo() {
 
     @Nullable
     public DesiredNodes desiredNodes() {
-        return DesiredNodes.latestFromClusterState(clusterState);
+        return desiredNodes;
     }
 
     /**
      * Returns the map of node id to shutdown metadata currently in the cluster
      */
     public Map<String, SingleNodeShutdownMetadata> nodeShutdowns() {
-        return metadata().nodeShutdowns();
+        return nodeShutdowns;
     }
 
     /**

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/NodeShutdownAllocationDecider.java

@@ -9,14 +9,11 @@
 package org.elasticsearch.cluster.routing.allocation.decider;
 
 import org.elasticsearch.cluster.metadata.IndexMetadata;
-import org.elasticsearch.cluster.metadata.Metadata;
-import org.elasticsearch.cluster.metadata.NodesShutdownMetadata;
 import org.elasticsearch.cluster.metadata.SingleNodeShutdownMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.RoutingNode;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
-import org.elasticsearch.core.Nullable;
 
 /**
  * An allocation decider that prevents shards from being allocated to a
@@ -35,7 +32,7 @@ public class NodeShutdownAllocationDecider extends AllocationDecider {
      */
     @Override
     public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation) {
-        final SingleNodeShutdownMetadata thisNodeShutdownMetadata = getNodeShutdownMetadata(allocation.metadata(), node.nodeId());
+        final SingleNodeShutdownMetadata thisNodeShutdownMetadata = allocation.nodeShutdowns().get(node.nodeId());
 
         if (thisNodeShutdownMetadata == null) {
             // There's no shutdown metadata for this node, return yes.
@@ -73,7 +70,7 @@ public Decision canRemain(IndexMetadata indexMetadata, ShardRouting shardRouting
      */
     @Override
     public Decision shouldAutoExpandToNode(IndexMetadata indexMetadata, DiscoveryNode node, RoutingAllocation allocation) {
-        SingleNodeShutdownMetadata thisNodeShutdownMetadata = getNodeShutdownMetadata(allocation.metadata(), node.getId());
+        SingleNodeShutdownMetadata thisNodeShutdownMetadata = allocation.nodeShutdowns().get(node.getId());
 
         if (thisNodeShutdownMetadata == null) {
             return allocation.decision(Decision.YES, NAME, "node [%s] is not preparing for removal from the cluster", node.getId());
@@ -95,14 +92,4 @@ public Decision shouldAutoExpandToNode(IndexMetadata indexMetadata, DiscoveryNod
         };
     }
 
-    @Nullable
-    private static SingleNodeShutdownMetadata getNodeShutdownMetadata(Metadata metadata, String nodeId) {
-        NodesShutdownMetadata nodesShutdownMetadata = metadata.custom(NodesShutdownMetadata.TYPE);
-        if (nodesShutdownMetadata == null || nodesShutdownMetadata.getAllNodeMetadataMap() == null) {
-            // There are no nodes in the process of shutting down, return null.
-            return null;
-        }
-
-        return nodesShutdownMetadata.getAllNodeMetadataMap().get(nodeId);
-    }
 }