Cache available tier names on DiscoveryNode and DiscoveryNodes (#88436)

We burn almost 2% of the runtime in many-shards benchmark bootstrapping on looping
through the roles just to compare names for the data tier allocation decider.
This change makes that step a lot cheaper by caching the role names in sets
that only need to be build rarely.

server/src/main/java/org/elasticsearch/cluster/node/DiscoveryNode.java

@@ -111,6 +111,8 @@ public class DiscoveryNode implements Writeable, ToXContentFragment {
     private final Map<String, String> attributes;
     private final Version version;
     private final SortedSet<DiscoveryNodeRole> roles;
+
+    private final Set<String> roleNames;
     private final String externalId;
 
     /**
@@ -314,7 +316,15 @@ public class DiscoveryNode implements Writeable, ToXContentFragment {
         this.attributes = Map.copyOf(attributes);
         assert DiscoveryNodeRole.roleNames().stream().noneMatch(attributes::containsKey)
             : "Node roles must not be provided as attributes but saw attributes " + attributes;
-        this.roles = Collections.unmodifiableSortedSet(new TreeSet<>(roles));
+        final TreeSet<DiscoveryNodeRole> sortedRoles = new TreeSet<>();
+        final String[] roleNames = new String[roles.size()];
+        int i = 0;
+        for (DiscoveryNodeRole role : roles) {
+            sortedRoles.add(role);
+            roleNames[i++] = role.roleName();
+        }
+        this.roles = Collections.unmodifiableSortedSet(sortedRoles);
+        this.roleNames = Set.of(roleNames);
         this.externalId = Objects.requireNonNullElse(externalId, this.nodeName);
     }
 
@@ -355,6 +365,7 @@ public class DiscoveryNode implements Writeable, ToXContentFragment {
         this.attributes = Collections.unmodifiableMap(in.readMap(readStringLiteral, readStringLiteral));
         int rolesSize = in.readVInt();
         final SortedSet<DiscoveryNodeRole> roles = new TreeSet<>();
+        final String[] roleNames = new String[rolesSize];
         for (int i = 0; i < rolesSize; i++) {
             final String roleName = in.readString();
             final String roleNameAbbreviation = in.readString();
@@ -369,6 +380,7 @@ public class DiscoveryNode implements Writeable, ToXContentFragment {
                     : "role name abbreviation [" + roleName + "] does not match role [" + role.roleNameAbbreviation() + "]";
                 roles.add(role);
             }
+            roleNames[i] = roleName;
         }
         this.roles = Collections.unmodifiableSortedSet(roles);
         this.version = Version.readVersion(in);
@@ -377,6 +389,17 @@ public class DiscoveryNode implements Writeable, ToXContentFragment {
         } else {
             this.externalId = nodeName;
         }
+        this.roleNames = Set.of(roleNames);
+    }
+
+    /**
+     * Check if node has the role with the given {@code roleName}.
+     *
+     * @param roleName role name to check
+     * @return true if node has the role of the given name
+     */
+    public boolean hasRole(String roleName) {
+        return this.roleNames.contains(roleName);
     }
 
     @Override

server/src/main/java/org/elasticsearch/cluster/node/DiscoveryNodes.java

@@ -60,6 +60,8 @@ public class DiscoveryNodes extends AbstractCollection<DiscoveryNode> implements
     private final Version maxNodeVersion;
     private final Version minNodeVersion;
 
+    private final Set<String> availableRoles;
+
     private DiscoveryNodes(
         Map<String, DiscoveryNode> nodes,
         Map<String, DiscoveryNode> dataNodes,
@@ -84,6 +86,11 @@ public class DiscoveryNodes extends AbstractCollection<DiscoveryNode> implements
         this.minNodeVersion = minNodeVersion;
         this.maxNodeVersion = maxNodeVersion;
         assert (localNodeId == null) == (localNode == null);
+        this.availableRoles = dataNodes.values()
+            .stream()
+            .flatMap(n -> n.getRoles().stream())
+            .map(DiscoveryNodeRole::roleName)
+            .collect(Collectors.toUnmodifiableSet());
     }
 
     @Override
@@ -107,6 +114,16 @@ public class DiscoveryNodes extends AbstractCollection<DiscoveryNode> implements
         return localNodeId.equals(masterNodeId);
     }
 
+    /**
+     * Checks if any node has the role with the given {@code roleName}.
+     *
+     * @param roleName name to check
+     * @return true if any node has the role of the given name
+     */
+    public boolean isRoleAvailable(String roleName) {
+        return availableRoles.contains(roleName);
+    }
+
     /**
      * Get the number of known nodes
      *

x-pack/plugin/autoscaling/src/main/java/org/elasticsearch/xpack/autoscaling/storage/ReactiveStorageDeciderService.java

@@ -408,14 +408,7 @@ public class ReactiveStorageDeciderService implements AutoscalingDeciderService
             IndexMetadata indexMetadata = indexMetadata(shard, allocation);
             Set<Decision.Type> decisionTypes = allocation.routingNodes()
                 .stream()
-                .map(
-                    node -> DataTierAllocationDecider.shouldFilter(
-                        indexMetadata,
-                        node.node().getRoles(),
-                        this::highestPreferenceTier,
-                        allocation
-                    )
-                )
+                .map(node -> DataTierAllocationDecider.shouldFilter(indexMetadata, node.node(), this::highestPreferenceTier, allocation))
                 .map(Decision::type)
                 .collect(Collectors.toSet());
             if (decisionTypes.contains(Decision.Type.NO)) {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/cluster/routing/allocation/DataTierAllocationDecider.java

@@ -46,7 +46,7 @@ public final class DataTierAllocationDecider extends AllocationDecider {
 
     @Override
     public Decision canAllocate(IndexMetadata indexMetadata, RoutingNode node, RoutingAllocation allocation) {
-        return shouldFilter(indexMetadata, node.node().getRoles(), allocation);
+        return shouldFilter(indexMetadata, node.node(), allocation);
     }
 
     @Override
@@ -56,15 +56,11 @@ public final class DataTierAllocationDecider extends AllocationDecider {
 
     @Override
     public Decision shouldAutoExpandToNode(IndexMetadata indexMetadata, DiscoveryNode node, RoutingAllocation allocation) {
-        return shouldFilter(indexMetadata, node.getRoles(), allocation);
-    }
-
-    private Decision shouldFilter(IndexMetadata indexMetadata, DiscoveryNode node, RoutingAllocation allocation) {
-        return shouldFilter(indexMetadata, node.getRoles(), allocation);
+        return shouldFilter(indexMetadata, node, allocation);
     }
 
-    private static Decision shouldFilter(IndexMetadata indexMd, Set<DiscoveryNodeRole> roles, RoutingAllocation allocation) {
-        return shouldFilter(indexMd, roles, DataTierAllocationDecider::preferredAvailableTier, allocation);
+    private static Decision shouldFilter(IndexMetadata indexMd, DiscoveryNode node, RoutingAllocation allocation) {
+        return shouldFilter(indexMd, node, DataTierAllocationDecider::preferredAvailableTier, allocation);
     }
 
     public interface PreferredTierFunction {
@@ -75,7 +71,7 @@ public final class DataTierAllocationDecider extends AllocationDecider {
 
     public static Decision shouldFilter(
         IndexMetadata indexMd,
-        Set<DiscoveryNodeRole> roles,
+        DiscoveryNode node,
         PreferredTierFunction preferredTierFunction,
         RoutingAllocation allocation
     ) {
@@ -86,7 +82,7 @@ public final class DataTierAllocationDecider extends AllocationDecider {
         Optional<String> tier = preferredTierFunction.apply(tierPreference, allocation.nodes(), allocation.desiredNodes());
         if (tier.isPresent()) {
             String tierName = tier.get();
-            if (allocationAllowed(tierName, roles)) {
+            if (allocationAllowed(tierName, node)) {
                 if (allocation.debugDecision()) {
                     return debugYesAllowed(allocation, tierPreference, tierName);
                 }
@@ -226,12 +222,12 @@ public final class DataTierAllocationDecider extends AllocationDecider {
     static boolean tierNodesPresent(String singleTier, DiscoveryNodes nodes) {
         assert singleTier.equals(DiscoveryNodeRole.DATA_ROLE.roleName()) || DataTier.validTierName(singleTier)
             : "tier " + singleTier + " is an invalid tier name";
-        for (DiscoveryNode node : nodes) {
-            if (allocationAllowed(singleTier, node.getRoles())) {
-                return true;
-            }
-        }
-        return false;
+        return nodes.isRoleAvailable(DiscoveryNodeRole.DATA_ROLE.roleName()) || nodes.isRoleAvailable(singleTier);
+    }
+
+    public static boolean allocationAllowed(String tierName, DiscoveryNode node) {
+        assert Strings.hasText(tierName) : "tierName must be not null and non-empty, but was [" + tierName + "]";
+        return node.hasRole(DiscoveryNodeRole.DATA_ROLE.roleName()) || node.hasRole(tierName);
     }
 
     public static boolean allocationAllowed(String tierName, Set<DiscoveryNodeRole> roles) {