SOLR-17534: Add ClusterState.getCollectionNames (#2826)

dsmiley · web-flow · commit cd681c073c3c · 2024-11-04T13:03:05.000-05:00
Refactoring to introduce getCollectionNames(). The motivation is to reduce callers of getCollectionsMap and getCollectionStates, which will go away soon.
diff --git a/solr/CHANGES.txt b/solr/CHANGES.txt
@@ -217,6 +217,8 @@ led to the suppression of exceptions. (Andrey Bozhko)
 
 * SOLR-11318: Introduce unit testing for AssertTool. (Eric Pugh, Jason Gerlowski)
 
+* SOLR-17534: Introduce ClusterState.getCollectionNames, a convenience method (David Smiley)
+
 ==================  9.7.1 ==================
 Bug Fixes
 ---------------------
diff --git a/solr/core/src/java/org/apache/solr/cli/DeleteTool.java b/solr/core/src/java/org/apache/solr/cli/DeleteTool.java
@@ -20,6 +20,7 @@
 
 import java.io.PrintStream;
 import java.lang.invoke.MethodHandles;
+import java.util.Collection;
 import java.util.List;
 import java.util.Locale;
 import java.util.Optional;
@@ -155,7 +156,7 @@ protected void deleteCollection(CloudSolrClient cloudSolrClient, CommandLine cli
             configName);
       } else {
         // need to scan all Collections to see if any are using the config
-        Set<String> collections = zkStateReader.getClusterState().getCollectionsMap().keySet();
+        Collection<String> collections = zkStateReader.getClusterState().getCollectionNames();
 
         // give a little note to the user if there are many collections in case it takes a while
         if (collections.size() > 50)
diff --git a/solr/core/src/java/org/apache/solr/handler/admin/ColStatus.java b/solr/core/src/java/org/apache/solr/handler/admin/ColStatus.java
@@ -21,7 +21,6 @@
 import java.lang.invoke.MethodHandles;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeMap;
@@ -76,7 +75,7 @@ public void getColStatus(NamedList<Object> results) {
     Collection<String> collections;
     String col = props.getStr(ZkStateReader.COLLECTION_PROP);
     if (col == null) {
-      collections = new HashSet<>(clusterState.getCollectionStates().keySet());
+      collections = clusterState.getCollectionNames();
     } else {
       collections = Collections.singleton(col);
     }
diff --git a/solr/core/src/test/org/apache/solr/cloud/ClusterStateMockUtilTest.java b/solr/core/src/test/org/apache/solr/cloud/ClusterStateMockUtilTest.java
@@ -36,7 +36,7 @@ public void testBuildClusterState_Simple() {
         ClusterStateMockUtil.buildClusterState("csr", "baseUrl1:8983_")) {
       ClusterState clusterState = zkStateReader.getClusterState();
       assertNotNull(clusterState);
-      assertEquals(1, clusterState.getCollectionStates().size());
+      assertEquals(1, clusterState.size());
       DocCollection collection1 = clusterState.getCollectionOrNull("collection1");
       assertNotNull(collection1);
       assertEquals(DocRouter.DEFAULT, collection1.getRouter());
@@ -62,7 +62,7 @@ public void testBuildClusterState_ReplicaTypes() {
         ClusterStateMockUtil.buildClusterState("csntp", "baseUrl1:8983_")) {
       ClusterState clusterState = zkStateReader.getClusterState();
       assertNotNull(clusterState);
-      assertEquals(1, clusterState.getCollectionStates().size());
+      assertEquals(1, clusterState.size());
       DocCollection collection1 = clusterState.getCollectionOrNull("collection1");
       assertNotNull(collection1);
       assertEquals(DocRouter.DEFAULT, collection1.getRouter());
@@ -83,7 +83,7 @@ public void testBuildClusterState_ReplicaStateAndType() {
         ClusterStateMockUtil.buildClusterState("csrStRpDnF", "baseUrl1:8983_")) {
       ClusterState clusterState = zkStateReader.getClusterState();
       assertNotNull(clusterState);
-      assertEquals(1, clusterState.getCollectionStates().size());
+      assertEquals(1, clusterState.size());
       DocCollection collection1 = clusterState.getCollectionOrNull("collection1");
       assertNotNull(collection1);
       assertEquals(DocRouter.DEFAULT, collection1.getRouter());
diff --git a/solr/core/src/test/org/apache/solr/cloud/ClusterStateTest.java b/solr/core/src/test/org/apache/solr/cloud/ClusterStateTest.java
@@ -81,7 +81,7 @@ public void testStoreAndRead() {
 
     assertEquals(
         "Provided liveNodes not used properly", 2, loadedClusterState.getLiveNodes().size());
-    assertEquals("No collections found", 2, loadedClusterState.getCollectionsMap().size());
+    assertEquals("No collections found", 2, loadedClusterState.size());
     assertEquals(
         "Properties not copied properly",
         replica.getStr("prop1"),
@@ -109,13 +109,13 @@ public void testStoreAndRead() {
 
     assertEquals(
         "Provided liveNodes not used properly", 2, loadedClusterState.getLiveNodes().size());
-    assertEquals("Should not have collections", 0, loadedClusterState.getCollectionsMap().size());
+    assertEquals("Should not have collections", 0, loadedClusterState.size());
 
     loadedClusterState =
         ClusterState.createFromJson(-1, (byte[]) null, liveNodes, Instant.now(), null);
 
     assertEquals(
         "Provided liveNodes not used properly", 2, loadedClusterState.getLiveNodes().size());
-    assertEquals("Should not have collections", 0, loadedClusterState.getCollectionsMap().size());
+    assertEquals("Should not have collections", 0, loadedClusterState.size());
   }
 }
diff --git a/solr/core/src/test/org/apache/solr/cloud/OverseerTest.java b/solr/core/src/test/org/apache/solr/cloud/OverseerTest.java
@@ -758,7 +758,7 @@ private void waitForCollections(ZkStateReader stateReader, String... collections
     while (0 < maxIterations--) {
 
       final ClusterState state = stateReader.getClusterState();
-      Set<String> availableCollections = state.getCollectionsMap().keySet();
+      Set<String> availableCollections = (Set<String>) state.getCollectionNames();
       int availableCount = 0;
       for (String requiredCollection : collections) {
         stateReader.waitForState(
diff --git a/solr/core/src/test/org/apache/solr/cloud/ReindexCollectionTest.java b/solr/core/src/test/org/apache/solr/cloud/ReindexCollectionTest.java
@@ -192,7 +192,7 @@ private void doTestSameTargetReindexing(boolean sourceRemove, boolean followAlia
     String prefix = ReindexCollectionCmd.TARGET_COL_PREFIX + targetCollection;
     while (!timeOut.hasTimedOut()) {
       timeOut.sleep(500);
-      for (String name : cloudManager.getClusterState().getCollectionsMap().keySet()) {
+      for (String name : cloudManager.getClusterState().getCollectionNames()) {
         if (name.startsWith(prefix)) {
           realTargetCollection = name;
           break;
diff --git a/solr/modules/sql/src/java/org/apache/solr/handler/sql/SolrSchema.java b/solr/modules/sql/src/java/org/apache/solr/handler/sql/SolrSchema.java
@@ -18,6 +18,7 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.Date;
 import java.util.EnumSet;
 import java.util.HashMap;
@@ -87,23 +88,12 @@ protected Map<String, Table> getTableMap() {
     String zk = this.properties.getProperty("zk");
     CloudSolrClient cloudSolrClient = solrClientCache.getCloudSolrClient(zk);
     ClusterState clusterState = cloudSolrClient.getClusterState();
-
-    final Map<String, Table> builder = new HashMap<>();
-
-    Set<String> collections = clusterState.getCollectionsMap().keySet();
-    for (String collection : collections) {
-      builder.put(collection, new SolrTable(this, collection));
-    }
-
     Aliases aliases = ZkStateReader.from(cloudSolrClient).getAliases();
-    for (String alias : aliases.getCollectionAliasListMap().keySet()) {
-      // don't create duplicate entries
-      if (!collections.contains(alias)) {
-        builder.put(alias, new SolrTable(this, alias));
-      }
-    }
 
-    return Map.copyOf(builder);
+    Collection<String> collectionNames = clusterState.getCollectionNames();
+    Set<String> aliasNames = aliases.getCollectionAliasListMap().keySet();
+    return Stream.concat(collectionNames.stream(), aliasNames.stream())
+        .collect(Collectors.toMap(n -> n, n -> new SolrTable(this, n), (t1, t2) -> t1));
   }
 
   private Map<String, LukeResponse.FieldInfo> getFieldInfo(final String collection) {
diff --git a/solr/prometheus-exporter/src/java/org/apache/solr/prometheus/scraper/SolrCloudScraper.java b/solr/prometheus-exporter/src/java/org/apache/solr/prometheus/scraper/SolrCloudScraper.java
@@ -19,6 +19,7 @@
 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -137,8 +138,8 @@ private Set<String> getBaseUrls() throws IOException {
         .collect(Collectors.toSet());
   }
 
-  private Set<String> getCollections() throws IOException {
-    return solrClient.getClusterState().getCollectionStates().keySet();
+  private Collection<String> getCollections() throws IOException {
+    return solrClient.getClusterState().getCollectionNames();
   }
 
   @Override
diff --git a/solr/solrj-streaming/src/test/org/apache/solr/client/solrj/io/sql/JdbcTest.java b/solr/solrj-streaming/src/test/org/apache/solr/client/solrj/io/sql/JdbcTest.java
@@ -29,8 +29,6 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Properties;
-import java.util.Set;
-import java.util.SortedSet;
 import java.util.TreeSet;
 import org.apache.lucene.tests.util.LuceneTestCase;
 import org.apache.solr.client.solrj.impl.CloudSolrClient;
@@ -652,8 +650,7 @@ private void testJDBCMethods(
       solrClient.connect();
       ZkStateReader zkStateReader = ZkStateReader.from(solrClient);
 
-      Set<String> collectionsSet = zkStateReader.getClusterState().getCollectionsMap().keySet();
-      SortedSet<String> tables = new TreeSet<>(collectionsSet);
+      var tables = new TreeSet<String>(zkStateReader.getClusterState().getCollectionNames());
 
       Aliases aliases = zkStateReader.getAliases();
       tables.addAll(aliases.getCollectionAliasListMap().keySet());
diff --git a/solr/solrj-zookeeper/src/java/org/apache/solr/common/cloud/ZkStateReader.java b/solr/solrj-zookeeper/src/java/org/apache/solr/common/cloud/ZkStateReader.java
@@ -625,7 +625,7 @@ private void constructState(Set<String> changedCollections) {
           collectionWatches.watchedCollections().size(),
           collectionWatches.activeCollectionCount(),
           lazyCollectionStates.keySet().size(),
-          clusterState.getCollectionStates().size());
+          clusterState.size());
     }
 
     if (log.isTraceEnabled()) {
diff --git a/solr/solrj/src/java/org/apache/solr/common/cloud/ClusterState.java b/solr/solrj/src/java/org/apache/solr/common/cloud/ClusterState.java
@@ -154,6 +154,16 @@ public DocCollection getCollectionOrNull(String collectionName, boolean allowCac
     return ref == null ? null : ref.get(allowCached);
   }
 
+  /**
+   * The collection names. Like a Set but might not implement it. Immutable; non-null. Some names
+   * returned might not resolve via {@link #getCollectionOrNull(String)}, so consider this a close
+   * approximation.
+   */
+  public Collection<String> getCollectionNames() {
+    // should we document we are sorted too?  Albeit that ties our hands.
+    return immutableCollectionStates.keySet();
+  }
+
   /**
    * Get a map of collection name vs DocCollection objects
    *
@@ -471,6 +481,7 @@ public String toString() {
     }
   }
 
+  /** The approximate number of collections. */
   public int size() {
     return collectionStates.size();
   }
diff --git a/solr/test-framework/src/java/org/apache/solr/cloud/AbstractFullDistribZkTestBase.java b/solr/test-framework/src/java/org/apache/solr/cloud/AbstractFullDistribZkTestBase.java
@@ -1813,7 +1813,7 @@ protected void assertDocCounts(boolean verbose) throws Exception {
           "Could not find collection "
               + DEFAULT_COLLECTION
               + " in "
-              + clusterState.getCollectionsMap().keySet());
+              + clusterState.getCollectionNames());
     }
 
     for (CloudJettyRunner cjetty : cloudJettys) {
diff --git a/solr/test-framework/src/java/org/apache/solr/cloud/MiniSolrCloudCluster.java b/solr/test-framework/src/java/org/apache/solr/cloud/MiniSolrCloudCluster.java
@@ -585,7 +585,7 @@ public void deleteAllCollections() throws Exception {
 
       reader.createClusterStateWatchersAndUpdate(); // up to date aliases & collections
       reader.aliasesManager.applyModificationAndExportToZk(aliases -> Aliases.EMPTY);
-      for (String collection : reader.getClusterState().getCollectionStates().keySet()) {
+      for (String collection : reader.getClusterState().getCollectionNames()) {
         CollectionAdminRequest.deleteCollection(collection).process(solrClient);
       }
 
@@ -595,7 +595,7 @@ public void deleteAllCollections() throws Exception {
             "Still waiting to see all collections removed from clusterstate.");
       }
 
-      for (String collection : reader.getClusterState().getCollectionStates().keySet()) {
+      for (String collection : reader.getClusterState().getCollectionNames()) {
         reader.waitForState(collection, 15, TimeUnit.SECONDS, Objects::isNull);
       }
     }
