SOLR-17535: Migrate ClusterState.getCollectionsMap callers to Stream (#2846)

Deprecated getCollectionsMap.

A refactoring, albeit some additional points:
* DeleteCollection: avoid NPE for non-existent collection (trivial change)
* ClusterStateUtil (test-framework): optimize waiting to use ZkStateReader.waitForState for single-collection

Author: dsmiley

solr/core/src/java/org/apache/solr/cloud/ConfigSetCmds.java

@@ -31,7 +31,6 @@
 import java.util.HashMap;
 import java.util.Map;
 import org.apache.solr.common.SolrException;
-import org.apache.solr.common.cloud.DocCollection;
 import org.apache.solr.common.cloud.ZkNodeProps;
 import org.apache.solr.common.cloud.ZkStateReader;
 import org.apache.solr.common.params.ConfigSetParams;
@@ -184,16 +183,18 @@ private static void deleteConfigSet(
       String configSetName, boolean force, CoreContainer coreContainer) throws IOException {
     ZkStateReader zkStateReader = coreContainer.getZkController().getZkStateReader();
 
-    for (Map.Entry<String, DocCollection> entry :
-        zkStateReader.getClusterState().getCollectionsMap().entrySet()) {
-      String configName = entry.getValue().getConfigName();
-      if (configSetName.equals(configName))
-        throw new SolrException(
-            SolrException.ErrorCode.BAD_REQUEST,
-            "Can not delete ConfigSet as it is currently being used by collection ["
-                + entry.getKey()
-                + "]");
-    }
+    zkStateReader
+        .getClusterState()
+        .forEachCollection(
+            state -> {
+              String configName = state.getConfigName();
+              if (configSetName.equals(configName))
+                throw new SolrException(
+                    SolrException.ErrorCode.BAD_REQUEST,
+                    "Can not delete ConfigSet as it is currently being used by collection ["
+                        + state.getName()
+                        + "]");
+            });
 
     String propertyPath = ConfigSetProperties.DEFAULT_FILENAME;
     NamedList<Object> properties =

solr/core/src/java/org/apache/solr/cloud/DistributedClusterStateUpdater.java

@@ -975,7 +975,7 @@ public void computeUpdates(ClusterState clusterState, SolrZkClient client) {
       final DocCollection docCollection = clusterState.getCollectionOrNull(collectionName);
       Optional<ZkWriteCommand> result =
           docCollection != null
-              ? NodeMutator.computeCollectionUpdate(nodeName, collectionName, docCollection, client)
+              ? NodeMutator.computeCollectionUpdate(nodeName, docCollection, client)
               : Optional.empty();
 
       if (docCollection == null) {

solr/core/src/java/org/apache/solr/cloud/api/collections/DeleteCollectionCmd.java

@@ -168,21 +168,19 @@ public void call(ClusterState state, ZkNodeProps message, NamedList<Object> resu
       }
 
       // delete related config set iff: it is auto generated AND not related to any other collection
-      String configSetName = coll.getConfigName();
+      String configSetName = coll == null ? null : coll.getConfigName();
 
       if (ConfigSetsHandler.isAutoGeneratedConfigSet(configSetName)) {
         boolean configSetIsUsedByOtherCollection = false;
 
         // make sure the configSet is not shared with other collections
         // Similar to what happens in: ConfigSetCmds::deleteConfigSet
-        for (Map.Entry<String, DocCollection> entry :
-            zkStateReader.getClusterState().getCollectionsMap().entrySet()) {
-          String otherConfigSetName = entry.getValue().getConfigName();
-          if (configSetName.equals(otherConfigSetName)) {
-            configSetIsUsedByOtherCollection = true;
-            break;
-          }
-        }
+        configSetIsUsedByOtherCollection =
+            zkStateReader
+                .getClusterState()
+                .collectionStream()
+                .map(DocCollection::getConfigName)
+                .anyMatch(configSetName::equals);
 
         if (!configSetIsUsedByOtherCollection) {
           // delete the config set

solr/core/src/java/org/apache/solr/cloud/api/collections/ReplicaMigrationUtils.java

@@ -34,9 +34,7 @@
 import org.apache.solr.common.SolrCloseableLatch;
 import org.apache.solr.common.cloud.ClusterState;
 import org.apache.solr.common.cloud.CollectionStateWatcher;
-import org.apache.solr.common.cloud.DocCollection;
 import org.apache.solr.common.cloud.Replica;
-import org.apache.solr.common.cloud.Slice;
 import org.apache.solr.common.cloud.ZkNodeProps;
 import org.apache.solr.common.cloud.ZkStateReader;
 import org.apache.solr.common.params.CoreAdminParams;
@@ -305,30 +303,20 @@ static boolean cleanupReplicas(
   }
 
   static List<Replica> getReplicasOfNodes(Collection<String> nodeNames, ClusterState state) {
-    List<Replica> sourceReplicas = new ArrayList<>();
-    for (Map.Entry<String, DocCollection> e : state.getCollectionsMap().entrySet()) {
-      for (Slice slice : e.getValue().getSlices()) {
-        for (Replica replica : slice.getReplicas()) {
-          if (nodeNames.contains(replica.getNodeName())) {
-            sourceReplicas.add(replica);
-          }
-        }
-      }
-    }
-    return sourceReplicas;
+    return state
+        .collectionStream()
+        .flatMap(dc -> dc.getSlices().stream())
+        .flatMap(s -> s.getReplicas().stream())
+        .filter(r -> nodeNames.contains(r.getNodeName()))
+        .toList();
   }
 
   static List<Replica> getReplicasOfNode(String nodeName, ClusterState state) {
-    List<Replica> sourceReplicas = new ArrayList<>();
-    for (Map.Entry<String, DocCollection> e : state.getCollectionsMap().entrySet()) {
-      for (Slice slice : e.getValue().getSlices()) {
-        for (Replica replica : slice.getReplicas()) {
-          if (nodeName.equals(replica.getNodeName())) {
-            sourceReplicas.add(replica);
-          }
-        }
-      }
-    }
-    return sourceReplicas;
+    return state
+        .collectionStream()
+        .flatMap(dc -> dc.getSlices().stream())
+        .flatMap(s -> s.getReplicas().stream())
+        .filter(r -> nodeName.equals(r.getNodeName()))
+        .toList();
   }
 }

solr/core/src/java/org/apache/solr/cloud/overseer/NodeMutator.java

@@ -49,23 +49,12 @@ public List<ZkWriteCommand> downNode(ClusterState clusterState, ZkNodeProps mess
 
     log.debug("DownNode state invoked for node: {}", nodeName);
 
-    List<ZkWriteCommand> zkWriteCommands = new ArrayList<>();
-
-    Map<String, DocCollection> collections = clusterState.getCollectionsMap();
-    for (Map.Entry<String, DocCollection> entry : collections.entrySet()) {
-      String collectionName = entry.getKey();
-      DocCollection docCollection = entry.getValue();
-      if (docCollection.isPerReplicaState()) continue;
-
-      Optional<ZkWriteCommand> zkWriteCommand =
-          computeCollectionUpdate(nodeName, collectionName, docCollection, zkClient);
-
-      if (zkWriteCommand.isPresent()) {
-        zkWriteCommands.add(zkWriteCommand.get());
-      }
-    }
-
-    return zkWriteCommands;
+    return clusterState
+        .collectionStream()
+        .filter(entry -> !entry.isPerReplicaState())
+        .map(docCollection -> computeCollectionUpdate(nodeName, docCollection, zkClient))
+        .flatMap(Optional::stream)
+        .toList();
   }
 
   /**
@@ -77,7 +66,7 @@ public List<ZkWriteCommand> downNode(ClusterState clusterState, ZkNodeProps mess
    *     for an update to state.json, depending on the configuration of the collection.
    */
   public static Optional<ZkWriteCommand> computeCollectionUpdate(
-      String nodeName, String collectionName, DocCollection docCollection, SolrZkClient client) {
+      String nodeName, DocCollection docCollection, SolrZkClient client) {
     boolean needToUpdateCollection = false;
     List<String> downedReplicas = new ArrayList<>();
     final Map<String, Slice> slicesCopy = new LinkedHashMap<>(docCollection.getSlicesMap());
@@ -107,13 +96,13 @@ public static Optional<ZkWriteCommand> computeCollectionUpdate(
 
         return Optional.of(
             new ZkWriteCommand(
-                collectionName,
+                docCollection.getName(),
                 docCollection.copyWithSlices(slicesCopy),
                 PerReplicaStatesOps.downReplicas(downedReplicas, prs),
                 false));
       } else {
         return Optional.of(
-            new ZkWriteCommand(collectionName, docCollection.copyWithSlices(slicesCopy)));
+            new ZkWriteCommand(docCollection.getName(), docCollection.copyWithSlices(slicesCopy)));
       }
     } else {
       // No update needed for this collection

solr/core/src/java/org/apache/solr/cluster/maintenance/InactiveShardRemover.java

@@ -143,7 +143,8 @@ public void stop() {
   void deleteInactiveSlices() {
     final ClusterState clusterState = coreContainer.getZkController().getClusterState();
     Collection<Slice> inactiveSlices =
-        clusterState.getCollectionsMap().values().stream()
+        clusterState
+            .collectionStream()
             .flatMap(v -> collectInactiveSlices(v).stream())
             .collect(Collectors.toSet());
 

solr/core/src/java/org/apache/solr/cluster/placement/impl/SimpleClusterAbstractionsImpl.java

@@ -93,7 +93,8 @@ public SolrCollection getCollection(String collectionName) {
 
     @Override
     public Iterator<SolrCollection> iterator() {
-      return clusterState.getCollectionsMap().values().stream()
+      return clusterState
+          .collectionStream()
           .map(SolrCollectionImpl::fromDocCollection)
           .collect(Collectors.toSet())
           .iterator();

solr/core/src/java/org/apache/solr/handler/admin/ClusterStatus.java

@@ -24,8 +24,9 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
-import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import org.apache.solr.common.SolrException;
 import org.apache.solr.common.cloud.Aliases;
 import org.apache.solr.common.cloud.ClusterState;
@@ -43,6 +44,7 @@
 import org.apache.zookeeper.KeeperException;
 
 public class ClusterStatus {
+
   private final ZkStateReader zkStateReader;
   private final SolrParams solrParams;
   private final String collection; // maybe null
@@ -178,79 +180,74 @@ private void fetchClusterStatusForCollOrAlias(
     String routeKey = solrParams.get(ShardParams._ROUTE_);
     String shard = solrParams.get(ZkStateReader.SHARD_ID_PROP);
 
-    Map<String, DocCollection> collectionsMap = null;
+    Stream<DocCollection> collectionStream;
     if (collection == null) {
-      collectionsMap = clusterState.getCollectionsMap();
+      collectionStream = clusterState.collectionStream();
     } else {
-      collectionsMap =
-          Collections.singletonMap(collection, clusterState.getCollectionOrNull(collection));
-    }
-
-    boolean isAlias = aliasVsCollections.containsKey(collection);
-    boolean didNotFindCollection = collectionsMap.get(collection) == null;
-
-    if (didNotFindCollection && isAlias) {
-      // In this case this.collection is an alias name not a collection
-      // get all collections and filter out collections not in the alias
-      // clusterState.getCollectionsMap() should be replaced with an inexpensive call
-      collectionsMap =
-          clusterState.getCollectionsMap().entrySet().stream()
-              .filter((entry) -> aliasVsCollections.get(collection).contains(entry.getKey()))
-              .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
-    }
-
-    NamedList<Object> collectionProps = new SimpleOrderedMap<>();
-
-    for (Map.Entry<String, DocCollection> entry : collectionsMap.entrySet()) {
-      Map<String, Object> collectionStatus;
-      String name = entry.getKey();
-      DocCollection clusterStateCollection = entry.getValue();
-      if (clusterStateCollection == null) {
-        if (collection != null) {
+      DocCollection collState = clusterState.getCollectionOrNull(collection);
+      if (collState != null) {
+        collectionStream = Stream.of(collState);
+      } else { // couldn't find collection
+        // hopefully an alias...
+        if (!aliasVsCollections.containsKey(collection)) { // not an alias either
           SolrException solrException =
               new SolrException(
-                  SolrException.ErrorCode.BAD_REQUEST, "Collection: " + name + " not found");
+                  SolrException.ErrorCode.BAD_REQUEST, "Collection: " + collection + " not found");
           solrException.setMetadata("CLUSTERSTATUS", "NOT_FOUND");
           throw solrException;
-        } else {
-          // collection might have got deleted at the same time
-          continue;
         }
+        // In this case this.collection is an alias name not a collection
+        // Resolve them (not recursively but maybe should?).
+        collectionStream =
+            aliasVsCollections.get(collection).stream()
+                .map(clusterState::getCollectionOrNull)
+                .filter(Objects::nonNull);
       }
+    }
 
-      Set<String> requestedShards = new HashSet<>();
-      if (routeKey != null) {
-        DocRouter router = clusterStateCollection.getRouter();
-        Collection<Slice> slices = router.getSearchSlices(routeKey, null, clusterStateCollection);
-        for (Slice slice : slices) {
-          requestedShards.add(slice.getName());
-        }
-      }
-      if (shard != null) {
-        String[] paramShards = shard.split(",");
-        requestedShards.addAll(Arrays.asList(paramShards));
-      }
+    // TODO use an Iterable to stream the data to the client instead of gathering it all in mem
 
-      byte[] bytes = Utils.toJSON(clusterStateCollection);
-      @SuppressWarnings("unchecked")
-      Map<String, Object> docCollection = (Map<String, Object>) Utils.fromJSON(bytes);
-      collectionStatus = getCollectionStatus(docCollection, name, requestedShards);
+    NamedList<Object> collectionProps = new SimpleOrderedMap<>();
 
-      collectionStatus.put("znodeVersion", clusterStateCollection.getZNodeVersion());
-      collectionStatus.put(
-          "creationTimeMillis", clusterStateCollection.getCreationTime().toEpochMilli());
+    collectionStream.forEach(
+        clusterStateCollection -> {
+          Map<String, Object> collectionStatus;
+          String name = clusterStateCollection.getName();
+
+          Set<String> requestedShards = new HashSet<>();
+          if (routeKey != null) {
+            DocRouter router = clusterStateCollection.getRouter();
+            Collection<Slice> slices =
+                router.getSearchSlices(routeKey, null, clusterStateCollection);
+            for (Slice slice : slices) {
+              requestedShards.add(slice.getName());
+            }
+          }
+          if (shard != null) {
+            String[] paramShards = shard.split(",");
+            requestedShards.addAll(Arrays.asList(paramShards));
+          }
 
-      if (collectionVsAliases.containsKey(name) && !collectionVsAliases.get(name).isEmpty()) {
-        collectionStatus.put("aliases", collectionVsAliases.get(name));
-      }
-      String configName = clusterStateCollection.getConfigName();
-      collectionStatus.put("configName", configName);
-      if (solrParams.getBool("prs", false) && clusterStateCollection.isPerReplicaState()) {
-        PerReplicaStates prs = clusterStateCollection.getPerReplicaStates();
-        collectionStatus.put("PRS", prs);
-      }
-      collectionProps.add(name, collectionStatus);
-    }
+          byte[] bytes = Utils.toJSON(clusterStateCollection);
+          @SuppressWarnings("unchecked")
+          Map<String, Object> docCollection = (Map<String, Object>) Utils.fromJSON(bytes);
+          collectionStatus = getCollectionStatus(docCollection, name, requestedShards);
+
+          collectionStatus.put("znodeVersion", clusterStateCollection.getZNodeVersion());
+          collectionStatus.put(
+              "creationTimeMillis", clusterStateCollection.getCreationTime().toEpochMilli());
+
+          if (collectionVsAliases.containsKey(name) && !collectionVsAliases.get(name).isEmpty()) {
+            collectionStatus.put("aliases", collectionVsAliases.get(name));
+          }
+          String configName = clusterStateCollection.getConfigName();
+          collectionStatus.put("configName", configName);
+          if (solrParams.getBool("prs", false) && clusterStateCollection.isPerReplicaState()) {
+            PerReplicaStates prs = clusterStateCollection.getPerReplicaStates();
+            collectionStatus.put("PRS", prs);
+          }
+          collectionProps.add(name, collectionStatus);
+        });
 
     // now we need to walk the collectionProps tree to cross-check replica state with live nodes
     crossCheckReplicaStateWithLiveNodes(liveNodes, collectionProps);

solr/core/src/java/org/apache/solr/handler/admin/api/ListCollections.java

@@ -20,10 +20,7 @@
 import static org.apache.solr.security.PermissionNameProvider.Name.COLL_READ_PERM;
 
 import jakarta.inject.Inject;
-import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 import org.apache.solr.client.api.endpoint.ListCollectionsApi;
 import org.apache.solr.client.api.model.ListCollectionsResponse;
 import org.apache.solr.common.cloud.DocCollection;
@@ -51,10 +48,16 @@ public ListCollectionsResponse listCollections() {
         instantiateJerseyResponse(ListCollectionsResponse.class);
     validateZooKeeperAwareCoreContainer(coreContainer);
 
-    Map<String, DocCollection> collections =
-        coreContainer.getZkController().getZkStateReader().getClusterState().getCollectionsMap();
-    List<String> collectionList = new ArrayList<>(collections.keySet());
-    Collections.sort(collectionList);
+    // resolve each name to ensure it exists.
+    //  TODO https://issues.apache.org/jira/browse/SOLR-16909 to go direct to ZK?
+    List<String> collectionList =
+        coreContainer
+            .getZkController()
+            .getClusterState()
+            .collectionStream()
+            .map(DocCollection::getName)
+            .sorted()
+            .toList();
     // XXX should we add aliases here?
     response.collections = collectionList;