SOLR-17381 SolrJ fix to fetch entire ClusterState if asked

solr/solrj/src/java/org/apache/solr/client/solrj/impl/BaseHttpClusterStateProvider.java

@@ -98,8 +98,8 @@ public ClusterState.CollectionRef getState(String collection) {
     for (String nodeName : liveNodes) {
       String baseUrl = Utils.getBaseUrlForNodeName(nodeName, urlScheme);
       try (SolrClient client = getSolrClient(baseUrl)) {
-        ClusterState cs = fetchClusterState(client, collection, null);
-        return cs.getCollectionRef(collection);
+        DocCollection docCollection = fetchCollectionState(client, collection);
+        return new ClusterState.CollectionRef(docCollection);
       } catch (SolrServerException | IOException e) {
         log.warn(
             "Attempt to fetch cluster state from {} failed.",
@@ -128,30 +128,12 @@ public ClusterState.CollectionRef getState(String collection) {
   }
 
   @SuppressWarnings("unchecked")
-  private ClusterState fetchClusterState(
-      SolrClient client, String collection, Map<String, Object> clusterProperties)
+  private ClusterState fetchClusterState(SolrClient client, Map<String, Object> clusterProperties)
       throws SolrServerException, IOException, NotACollectionException {
     SimpleOrderedMap<?> cluster =
-        submitClusterStateRequest(client, collection, ClusterStateRequestType.FETCH_COLLECTION);
+        submitClusterStateRequest(client, null, ClusterStateRequestType.FETCH_CLUSTER_STATE);
 
-    Map<String, Object> collectionsMap;
-    if (collection != null) {
-      collectionsMap =
-          Collections.singletonMap(
-              collection, ((NamedList<?>) cluster.get("collections")).get(collection));
-    } else {
-      collectionsMap = ((NamedList<?>) cluster.get("collections")).asMap(10);
-    }
-    int znodeVersion;
-    Map<String, Object> collFromStatus = (Map<String, Object>) (collectionsMap).get(collection);
-    if (collection != null && collFromStatus == null) {
-      throw new NotACollectionException(); // probably an alias
-    }
-    if (collection != null) { // can be null if alias
-      znodeVersion = (int) collFromStatus.get("znodeVersion");
-    } else {
-      znodeVersion = -1;
-    }
+    Map<String, Object> collectionsMap = ((NamedList<?>) cluster.get("collections")).asMap(10);
 
     ClusterState cs = new ClusterState(this.liveNodes, new HashMap<>());
     List<String> liveNodesList = (List<String>) cluster.get("live_nodes");
@@ -163,14 +145,11 @@ private ClusterState fetchClusterState(
     }
 
     for (Map.Entry<String, Object> e : collectionsMap.entrySet()) {
-      @SuppressWarnings("rawtypes")
-      Map m = (Map) e.getValue();
-      Long creationTimeMillisFromClusterStatus = (Long) m.get("creationTimeMillis");
-      Instant creationTime =
-          creationTimeMillisFromClusterStatus == null
-              ? Instant.EPOCH
-              : Instant.ofEpochMilli(creationTimeMillisFromClusterStatus);
-      cs = cs.copyWith(e.getKey(), fillPrs(znodeVersion, e, creationTime, m));
+      String collectionName = e.getKey();
+      Map<String, Object> collStateMap = (Map<String, Object>) e.getValue();
+      cs =
+          cs.copyWith(
+              collectionName, getDocCollectionFromObjects(collectionName, collStateMap, -1));
     }
 
     if (clusterProperties != null) {
@@ -182,6 +161,37 @@ private ClusterState fetchClusterState(
     return cs;
   }
 
+  private DocCollection getDocCollectionFromObjects(
+      String collectionName, Map<String, Object> collStateMap, int zNodeVersion) {
+
+    Long creationTimeMillis = (Long) collStateMap.get("creationTimeMillis");
+    Instant creationTime =
+        creationTimeMillis == null
+            ? Instant.EPOCH
+            : Instant.ofEpochMilli(creationTimeMillis);
+    return fillPrs(collectionName, collStateMap, creationTime, zNodeVersion);
+  }
+
+  private DocCollection fetchCollectionState(SolrClient client, String collection)
+      throws SolrServerException, IOException, NotACollectionException {
+
+    SimpleOrderedMap<?> cluster =
+        submitClusterStateRequest(client, collection, ClusterStateRequestType.FETCH_COLLECTION);
+
+    Map<String, Object> collectionsMap =
+        Collections.singletonMap(
+            collection, ((NamedList<?>) cluster.get("collections")).get(collection));
+
+    int znodeVersion = -1;
+    @SuppressWarnings("unchecked")
+    Map<String, Object> collStateMap = (Map<String, Object>) (collectionsMap).get(collection);
+    if (collStateMap == null) {
+      throw new NotACollectionException(); // probably an alias
+    }
+    znodeVersion = (int) collStateMap.get("znodeVersion");
+    return getDocCollectionFromObjects(collection, collStateMap, znodeVersion);
+  }
+
   private SimpleOrderedMap<?> submitClusterStateRequest(
       SolrClient client, String collection, ClusterStateRequestType requestType)
       throws SolrServerException, IOException {
@@ -198,30 +208,35 @@ private SimpleOrderedMap<?> submitClusterStateRequest(
     } else if (requestType == ClusterStateRequestType.FETCH_NODE_ROLES) {
       params.set("roles", "true");
     }
-
-    params.set("includeAll", "false");
+    if (requestType == ClusterStateRequestType.FETCH_CLUSTER_STATE) {
+      params.set("includeAll", "true");
+    } else {
+      params.set("includeAll", "false");
+    }
     params.set("prs", "true");
     QueryRequest request = new QueryRequest(params);
     request.setPath("/admin/collections");
     return (SimpleOrderedMap<?>) client.request(request).get("cluster");
   }
 
-  @SuppressWarnings({"rawtypes", "unchecked"})
+  @SuppressWarnings({"unchecked"})
   private DocCollection fillPrs(
-      int znodeVersion, Map.Entry<String, Object> e, Instant creationTime, Map m) {
+      String collectionName,
+      Map<String, Object> collStateMap,
+      Instant creationTime,
+      int znodeVersion) {
     DocCollection.PrsSupplier prsSupplier = null;
-    if (m.containsKey("PRS")) {
-      Map prs = (Map) m.remove("PRS");
+    if (collStateMap.containsKey("PRS")) {
+      Map<String, Object> prs = (Map<String, Object>) collStateMap.remove("PRS");
       prsSupplier =
           () ->
               new PerReplicaStates(
                   (String) prs.get("path"),
                   (Integer) prs.get("cversion"),
                   (List<String>) prs.get("states"));
     }
-
     return ClusterState.collectionFromObjects(
-        e.getKey(), m, znodeVersion, creationTime, prsSupplier);
+        collectionName, collStateMap, znodeVersion, creationTime, prsSupplier);
   }
 
   @Override
@@ -346,8 +361,8 @@ public ClusterState getClusterState() {
     for (String nodeName : liveNodes) {
       String baseUrl = Utils.getBaseUrlForNodeName(nodeName, urlScheme);
       try (SolrClient client = getSolrClient(baseUrl)) {
-        return fetchClusterState(client, null, null);
-      } catch (SolrServerException | SolrClient.RemoteSolrException | IOException e) {
+        return fetchClusterState(client, null);
+      } catch (SolrServerException | RemoteSolrException | IOException e) {
         log.warn("Attempt to fetch cluster state from {} failed.", baseUrl, e);
       } catch (NotACollectionException e) {
         // not possible! (we passed in null for collection, so it can't be an alias)
@@ -376,7 +391,7 @@ public Map<String, Object> getClusterProperties() {
         SimpleOrderedMap<?> cluster =
             submitClusterStateRequest(client, null, ClusterStateRequestType.FETCH_CLUSTER_PROP);
         return (Map<String, Object>) cluster.get("properties");
-      } catch (SolrServerException | SolrClient.RemoteSolrException | IOException e) {
+      } catch (SolrServerException | RemoteSolrException | IOException e) {
         log.warn("Attempt to fetch cluster state from {} failed.", baseUrl, e);
       }
     }
@@ -428,6 +443,7 @@ private enum ClusterStateRequestType {
     FETCH_LIVE_NODES,
     FETCH_CLUSTER_PROP,
     FETCH_NODE_ROLES,
-    FETCH_COLLECTION
+    FETCH_COLLECTION,
+    FETCH_CLUSTER_STATE
   }
 }