Revert "Remove DiscoveryNodes#getAllNodes (#83538) (#83574)"

This reverts commit e6603b2.

Author: arteam

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

@@ -24,6 +24,7 @@
 import java.io.IOException;
 import java.util.AbstractCollection;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -165,6 +166,14 @@ public ImmutableOpenMap<String, DiscoveryNode> getCoordinatingOnlyNodes() {
         return nodes.build();
     }
 
+    /**
+     * Return all the nodes as a collection
+     * @return
+     */
+    public Collection<DiscoveryNode> getAllNodes() {
+        return this;
+    }
+
     /**
      * Returns a stream of all nodes, with master nodes at the front
      */

server/src/main/java/org/elasticsearch/persistent/PersistentTasksClusterService.java

@@ -336,6 +336,7 @@ private <Params extends PersistentTaskParams> Assignment createAssignment(
         // want to assign a persistent task to a node that will shortly be
         // leaving the cluster
         final List<DiscoveryNode> candidateNodes = currentState.nodes()
+            .getAllNodes()
             .stream()
             .filter(dn -> isNodeShuttingDown(currentState, dn.getId()) == false)
             .collect(Collectors.toList());

server/src/test/java/org/elasticsearch/action/search/TransportSearchActionTests.java

@@ -1289,7 +1289,7 @@ public void testLocalShardIteratorFromPointInTime() {
             } else {
                 // relocated or no longer assigned
                 relocatedContexts.add(new ShardId(indexMetadata.getIndex(), shardId));
-                targetNode = randomFrom(clusterState.nodes()).getId();
+                targetNode = randomFrom(clusterState.nodes().getAllNodes()).getId();
             }
             contexts.put(
                 new ShardId(indexMetadata.getIndex(), shardId),
@@ -1343,7 +1343,7 @@ public void testLocalShardIteratorFromPointInTime() {
             anotherShardId,
             new SearchContextIdForNode(
                 null,
-                randomFrom(clusterState.nodes()).getId(),
+                randomFrom(clusterState.nodes().getAllNodes()).getId(),
                 new ShardSearchContextId(UUIDs.randomBase64UUID(), randomNonNegativeLong(), null)
             )
         );

server/src/test/java/org/elasticsearch/persistent/PersistentTasksClusterServiceTests.java

@@ -627,7 +627,7 @@ public void testTasksNotAssignedToShuttingDownNodes() {
         // Now that we have a bunch of tasks that need to be assigned, let's
         // mark half the nodes as shut down and make sure they do not have any
         // tasks assigned
-        Collection<DiscoveryNode> allNodes = clusterState.nodes();
+        Collection<DiscoveryNode> allNodes = clusterState.nodes().getAllNodes();
         Map<String, SingleNodeShutdownMetadata> shutdownMetadataMap = new HashMap<>();
         allNodes.stream()
             .limit(Math.floorDiv(allNodes.size(), 2))
@@ -816,6 +816,10 @@ private Assignment randomNodeAssignment(Collection<DiscoveryNode> nodes) {
         return Optional.ofNullable(randomFrom(nodes)).map(node -> new Assignment(node.getId(), "test assignment")).orElse(NO_NODE_FOUND);
     }
 
+    private Assignment randomNodeAssignment(DiscoveryNodes nodes) {
+        return randomNodeAssignment(nodes.getAllNodes());
+    }
+
     private String dumpEvent(ClusterChangedEvent event) {
         return "nodes_changed: "
             + event.nodesChanged()

x-pack/plugin/ccr/src/test/java/org/elasticsearch/xpack/ccr/action/ShardFollowTasksExecutorAssignmentTests.java

@@ -92,7 +92,7 @@ private void runAssignmentTest(
         clusterStateBuilder.nodes(nodesBuilder);
         final Assignment assignment = executor.getAssignment(
             mock(ShardFollowTask.class),
-            clusterStateBuilder.nodes(),
+            clusterStateBuilder.nodes().getAllNodes(),
             clusterStateBuilder.build()
         );
         consumer.accept(theSpecial, assignment);

x-pack/plugin/enrich/src/main/java/org/elasticsearch/xpack/enrich/action/InternalExecutePolicyAction.java

@@ -160,7 +160,8 @@ DiscoveryNode selectNodeForPolicyExecution(DiscoveryNodes discoNodes) {
                 return discoNodes.getLocalNode();
             }
 
-            final var nodes = discoNodes.stream()
+            final var nodes = discoNodes.getAllNodes()
+                .stream()
                 // filter out elected master node (which is the local node)
                 .filter(discoNode -> discoNode.getId().equals(discoNodes.getMasterNodeId()) == false)
                 // filter out dedicated master nodes

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportGetDeploymentStatsAction.java

@@ -143,6 +143,7 @@ protected void doExecute(
             ClusterState latestState = clusterService.state();
             Set<String> nodesShuttingDown = TransportStartTrainedModelDeploymentAction.nodesShuttingDown(latestState);
             List<DiscoveryNode> nodes = latestState.getNodes()
+                .getAllNodes()
                 .stream()
                 .filter(d -> nodesShuttingDown.contains(d.getId()) == false)
                 .filter(StartTrainedModelDeploymentAction.TaskParams::mayAllocateToNode)

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/action/TransportStartTrainedModelDeploymentAction.java

@@ -441,6 +441,7 @@ public boolean test(ClusterState clusterState) {
             }
             Set<String> nodesShuttingDown = nodesShuttingDown(clusterState);
             List<DiscoveryNode> nodes = clusterState.nodes()
+                .getAllNodes()
                 .stream()
                 .filter(d -> nodesShuttingDown.contains(d.getId()) == false)
                 .filter(TaskParams::mayAllocateToNode)

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/inference/allocation/TrainedModelAllocationClusterService.java

@@ -253,7 +253,7 @@ ClusterState createModelAllocation(ClusterState currentState, StartTrainedModelD
 
         Set<String> shuttingDownNodes = nodesShuttingDown(currentState);
         Map<String, String> nodeToReason = new TreeMap<>();
-        for (DiscoveryNode node : currentState.getNodes()) {
+        for (DiscoveryNode node : currentState.getNodes().getAllNodes()) {
             if (StartTrainedModelDeploymentAction.TaskParams.mayAllocateToNode(node) && shuttingDownNodes.contains(node.getId()) == false) {
                 Optional<String> maybeError = nodeHasCapacity(currentState, params, node);
                 if (maybeError.isPresent()) {
@@ -357,6 +357,7 @@ ClusterState addRemoveAllocationNodes(ClusterState currentState) {
         final TrainedModelAllocationMetadata.Builder builder = TrainedModelAllocationMetadata.builder(currentState);
         Set<String> shuttingDownNodes = nodesShuttingDown(currentState);
         Map<String, DiscoveryNode> currentEligibleNodes = currentState.getNodes()
+            .getAllNodes()
             .stream()
             // TODO: Change when we update `mayAllocateToNode`
             .filter(

x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/job/task/OpenJobPersistentTasksExecutor.java

@@ -219,7 +219,7 @@ public void validate(OpenJobAction.JobParams params, ClusterState clusterState)
         validateJobAndId(jobId, job);
         // If we already know that we can't find an ml node because all ml nodes are running at capacity or
         // simply because there are no ml nodes in the cluster then we fail quickly here:
-        PersistentTasksCustomMetadata.Assignment assignment = getAssignment(params, clusterState.nodes(), clusterState);
+        PersistentTasksCustomMetadata.Assignment assignment = getAssignment(params, clusterState.nodes().getAllNodes(), clusterState);
         if (assignment.equals(AWAITING_UPGRADE)) {
             throw makeCurrentlyBeingUpgradedException(logger, params.getJobId());
         }