SOLR-17582 Stream CLUSTERSTATUS API response

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

@@ -183,6 +183,12 @@ private void fetchClusterStatusForCollOrAlias(
     String routeKey = solrParams.get(ShardParams._ROUTE_);
     String shard = solrParams.get(ZkStateReader.SHARD_ID_PROP);
 
+    Set<String> requestedShards = new HashSet<>();
+    if (shard != null) {
+      String[] paramShards = shard.split(",");
+      requestedShards.addAll(Arrays.asList(paramShards));
+    }
+
     Stream<DocCollection> collectionStream;
     if (collection == null) {
       collectionStream = clusterState.collectionStream();
@@ -210,12 +216,20 @@ private void fetchClusterStatusForCollOrAlias(
 
     MapWriter collectionPropsWriter =
         ew -> {
-          SolrCollectionProperiesIterator collectionProps =
-              new SolrCollectionProperiesIterator(
-                  collectionStream.iterator(), collectionVsAliases, routeKey, liveNodes, shard);
-          while (collectionProps.hasNext()) {
-            Map<String, Object> collectionPropsMap = (collectionProps.next().asMap());
-            collectionPropsMap.forEach(
+          Iterator<Map<String, Object>> it =
+              collectionStream
+                  .map(
+                      (collectionState) ->
+                          collectionPropsResponse(
+                              collectionState,
+                              collectionVsAliases,
+                              routeKey,
+                              liveNodes,
+                              requestedShards))
+                  .iterator();
+          while (it.hasNext()) {
+            Map<String, Object> props = it.next();
+            props.forEach(
                 (key, value) -> {
                   try {
                     ew.put(key, value);
@@ -280,8 +294,8 @@ private Map<String, Object> getCollectionStatus(
    */
   @SuppressWarnings("unchecked")
   protected void crossCheckReplicaStateWithLiveNodes(
-      List<String> liveNodes, NamedList<Object> collectionProps) {
-    for (Map.Entry<String, Object> next : collectionProps) {
+      List<String> liveNodes, Map<String, Object> collectionProps) {
+    for (Map.Entry<String, Object> next : collectionProps.entrySet()) {
       Map<String, Object> collMap = (Map<String, Object>) next.getValue();
       Map<String, Object> shards = (Map<String, Object>) collMap.get("shards");
       for (Object nextShard : shards.values()) {
@@ -342,76 +356,47 @@ public static Map<String, Object> postProcessCollectionJSON(Map<String, Object>
     return collection;
   }
 
-  private class SolrCollectionProperiesIterator implements Iterator<NamedList<Object>> {
-
-    final Iterator<DocCollection> it;
-    Map<String, List<String>> collectionVsAliases;
-    String routeKey;
-    List<String> liveNodes;
-    String shard;
-
-    public SolrCollectionProperiesIterator(
-        Iterator<DocCollection> it,
-        Map<String, List<String>> collectionVsAliases,
-        String routeKey,
-        List<String> liveNodes,
-        String shard) {
-      this.it = it;
-      this.collectionVsAliases = collectionVsAliases;
-      this.routeKey = routeKey;
-      this.liveNodes = liveNodes;
-      this.shard = shard;
-    }
-
-    @Override
-    public boolean hasNext() {
-      return it.hasNext();
-    }
-
-    @Override
-    public NamedList<Object> next() {
-      NamedList<Object> collectionProps = new SimpleOrderedMap<>();
-      DocCollection clusterStateCollection = it.next();
-      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));
+  private Map<String, Object> collectionPropsResponse(
+      DocCollection clusterStateCollection,
+      Map<String, List<String>> collectionVsAliases,
+      String routeKey,
+      List<String> liveNodes,
+      Set<String> requestedShards) {
+    Map<String, Object> collectionProps = new HashMap<>();
+    Map<String, Object> collectionStatus;
+    String name = clusterStateCollection.getName();
+
+    if (routeKey != null) {
+      DocRouter router = clusterStateCollection.getRouter();
+      Collection<Slice> slices = router.getSearchSlices(routeKey, null, clusterStateCollection);
+      for (Slice slice : slices) {
+        requestedShards.add(slice.getName());
       }
+    }
 
-      byte[] bytes = Utils.toJSON(clusterStateCollection);
-      @SuppressWarnings("unchecked")
-      Map<String, Object> docCollection = (Map<String, Object>) Utils.fromJSON(bytes);
-      collectionStatus = getCollectionStatus(docCollection, name, requestedShards);
+    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());
+    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);
-      return collectionProps;
+    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.put(name, collectionStatus);
+
+    // now we need to walk the collectionProps tree to cross-check replica state with live
+    // nodes
+    crossCheckReplicaStateWithLiveNodes(liveNodes, collectionProps);
+    return collectionProps;
   }
 }

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

@@ -139,11 +139,11 @@ private ClusterState fetchClusterState(SolrClient client)
       liveNodesTimestamp = System.nanoTime();
     }
 
-    var collectionsNl = (Map<String, Map<String, Object>>) cluster.get("collections");
+    var collectionsMap = (Map<String, Map<String, Object>>) cluster.get("collections");
 
     Map<String, DocCollection> collStateByName =
-        CollectionUtil.newLinkedHashMap(collectionsNl.size());
-    for (Entry<String, Map<String, Object>> entry : collectionsNl.entrySet()) {
+        CollectionUtil.newLinkedHashMap(collectionsMap.size());
+    for (Entry<String, Map<String, Object>> entry : collectionsMap.entrySet()) {
       collStateByName.put(
           entry.getKey(), getDocCollectionFromObjects(entry.getKey(), entry.getValue()));
     }
@@ -184,13 +184,12 @@ private DocCollection fetchCollectionState(SolrClient client, String collection)
     SimpleOrderedMap<?> cluster =
         submitClusterStateRequest(client, collection, ClusterStateRequestType.FETCH_COLLECTION);
 
-    var collState = (Map<String, Object>) cluster.findRecursive("collections");
-    var collStateMap = (Map<String, Object>) collState.get(collection);
+    var collState = (Map<String, Object>) cluster.findRecursive("collections", collection);
 
-    if (collStateMap == null) {
+    if (collState == null) {
       throw new NotACollectionException(); // probably an alias
     }
-    return getDocCollectionFromObjects(collection, collStateMap);
+    return getDocCollectionFromObjects(collection, collState);
   }
 
   private SimpleOrderedMap<?> submitClusterStateRequest(

solr/solrj/src/java/org/apache/solr/common/util/NamedList.java

@@ -319,7 +319,7 @@ private void killAll(String name) {
   /**
    * Recursively parses the NamedList structure to arrive at a specific element. As you descend the
    * NamedList tree, the last element can be any type, including NamedList, but the previous
-   * elements MUST be NamedList objects themselves. A null value is returned if the indicated
+   * elements MUST be NamedList or Map objects themselves. A null value is returned if the indicated
    * hierarchy doesn't exist, but NamedList allows null values so that could be the actual value at
    * the end of the path.
    *
@@ -348,9 +348,9 @@ public Object findRecursive(String... args) {
        * it to this list. Then we retrieve the first key from this list and
        * assign it to value.
        *
-       * On the next loop, we check whether the retrieved value is a NamedList.
-       * If it is, then we drop to that NamedList, grab the value of the
-       * next key, and start the loop over. If it is not a NamedList, then we
+       * On the next loop, we check whether the retrieved value is a NamedList or Map.
+       * If it is, then we drop to that NamedList/Map, grab the value of the
+       * next key, and start the loop over. If it is not a NamedList/Map, then we
        * assign the value to null and break out of the loop.
        *
        * Assigning the value to null and then breaking out of the loop seems
@@ -363,6 +363,8 @@ public Object findRecursive(String... args) {
       } else {
         if (value instanceof NamedList) {
           currentList = (NamedList<?>) value;
+        } else if (value instanceof Map) {
+          currentList = new NamedList<>((Map<String, Object>) value);
         } else {
           value = null;
           break;

solr/solrj/src/test/org/apache/solr/common/util/NamedListTest.java

@@ -17,6 +17,7 @@
 package org.apache.solr.common.util;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import org.apache.solr.SolrTestCase;
@@ -209,4 +210,36 @@ public void testShallowMap() {
     assertEquals("Val2", nl.get("key2"));
     assertEquals("Val2", m.get("key2"));
   }
+
+  @Test
+  @SuppressWarnings("unchecked")
+  public void testFindRecursiveWithMap() {
+
+    // key1 (NamedList)
+    // - key2 (Map)
+    // -- key3 (NamedList)
+    // --- key4 (Map)
+
+    NamedList<Object> nl1 = new NamedList<>();
+    Map<String, Object> map2 = new HashMap<>();
+    NamedList<Object> nl3 = new NamedList<>();
+    Map<String, Object> map4 = new HashMap<>();
+
+    map4.put("key4", "value4");
+    nl3.add("key3", map4);
+    map2.put("key2", nl3);
+    nl1.add("key1", map2);
+
+    Map<String, Object> test1 = (Map<String, Object>) nl1.findRecursive("key1");
+    assertNotNull(test1);
+
+    NamedList<Object> test2 = (NamedList<Object>) nl1.findRecursive("key1", "key2");
+    assertNotNull(test2);
+
+    Map<String, Object> test3 = (Map<String, Object>) nl1.findRecursive("key1", "key2", "key3");
+    assertNotNull(test3);
+
+    String test4 = (String) nl1.findRecursive("key1", "key2", "key3", "key4");
+    assertEquals("value4", test4);
+  }
 }