From 377f16de5e465277209d1f791d95416132584b16 Mon Sep 17 00:00:00 2001
From: Noble Paul <noble@apache.org>
Date: Thu, 14 Jun 2018 00:07:30 +1000
Subject: [PATCH] SOLR-12208: added a testcase for fetching arbitrary replica
 info

---
 .../solr/cloud/autoscaling/TestPolicyCloud.java    | 14 ++++++++++++--
 .../solrj/impl/SolrClientNodeStateProvider.java    | 13 ++++++++++---
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/solr/core/src/test/org/apache/solr/cloud/autoscaling/TestPolicyCloud.java b/solr/core/src/test/org/apache/solr/cloud/autoscaling/TestPolicyCloud.java
index cde23d08a54..19324e73c88 100644
--- a/solr/core/src/test/org/apache/solr/cloud/autoscaling/TestPolicyCloud.java
+++ b/solr/core/src/test/org/apache/solr/cloud/autoscaling/TestPolicyCloud.java
@@ -33,16 +33,17 @@ import org.apache.lucene.util.LuceneTestCase;
 import org.apache.solr.client.solrj.SolrRequest;
 import org.apache.solr.client.solrj.SolrServerException;
 import org.apache.solr.client.solrj.cloud.DistributedQueueFactory;
+import org.apache.solr.client.solrj.cloud.SolrCloudManager;
 import org.apache.solr.client.solrj.cloud.autoscaling.AutoScalingConfig;
 import org.apache.solr.client.solrj.cloud.autoscaling.Policy;
 import org.apache.solr.client.solrj.cloud.autoscaling.ReplicaInfo;
 import org.apache.solr.client.solrj.cloud.autoscaling.Row;
-import org.apache.solr.client.solrj.cloud.SolrCloudManager;
 import org.apache.solr.client.solrj.cloud.autoscaling.Suggestion;
 import org.apache.solr.client.solrj.embedded.JettySolrRunner;
 import org.apache.solr.client.solrj.impl.CloudSolrClient;
 import org.apache.solr.client.solrj.impl.HttpSolrClient;
 import org.apache.solr.client.solrj.impl.SolrClientCloudManager;
+import org.apache.solr.client.solrj.impl.SolrClientNodeStateProvider;
 import org.apache.solr.client.solrj.request.CollectionAdminRequest;
 import org.apache.solr.cloud.OverseerTaskProcessor;
 import org.apache.solr.cloud.SolrCloudTestCase;
@@ -106,10 +107,19 @@ public class TestPolicyCloud extends SolrCloudTestCase {
         "  }" +
         "}";
     AutoScalingConfig config = new AutoScalingConfig((Map<String, Object>) Utils.fromJSONString(autoScaleJson));
+    AtomicInteger count = new AtomicInteger(0);
     SolrCloudManager cloudManager = new SolrClientCloudManager(new ZkDistributedQueueFactory(cluster.getZkClient()), cluster.getSolrClient());
+    String nodeName = cloudManager.getClusterStateProvider().getLiveNodes().iterator().next();
+    SolrClientNodeStateProvider nodeStateProvider = (SolrClientNodeStateProvider) cloudManager.getNodeStateProvider();
+    Map<String, Map<String, List<ReplicaInfo>>> result = nodeStateProvider.getReplicaInfo(nodeName, Collections.singleton("UPDATE./update.requests"));
+    nodeStateProvider.forEachReplica(nodeName, replicaInfo -> {
+      if (replicaInfo.getVariables().containsKey("UPDATE./update.requests")) count.incrementAndGet();
+    });
+    assertTrue(count.get() > 0);
+
     Policy.Session session = config.getPolicy().createSession(cloudManager);
 
-    AtomicInteger count = new AtomicInteger(0);
+    count .set(0);
     for (Row row : session.getSorted()) {
       row.collectionVsShardVsReplicas.forEach((c, shardVsReplicas) -> shardVsReplicas.forEach((s, replicaInfos) -> {
         for (ReplicaInfo replicaInfo : replicaInfos) {
diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/impl/SolrClientNodeStateProvider.java b/solr/solrj/src/java/org/apache/solr/client/solrj/impl/SolrClientNodeStateProvider.java
index c331c73befe..b65d42ab3ff 100644
--- a/solr/solrj/src/java/org/apache/solr/client/solrj/impl/SolrClientNodeStateProvider.java
+++ b/solr/solrj/src/java/org/apache/solr/client/solrj/impl/SolrClientNodeStateProvider.java
@@ -27,6 +27,7 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -113,6 +114,10 @@ public class SolrClientNodeStateProvider implements NodeStateProvider, MapWriter
     return ctx.getTags();
   }
 
+  public void forEachReplica(String node, Consumer<ReplicaInfo> consumer){
+    Row.forEachReplica(nodeVsCollectionVsShardVsReplicaInfo.get(node), consumer);
+  }
+
 
   @Override
   public Map<String, Map<String, List<ReplicaInfo>>> getReplicaInfo(String node, Collection<String> keys) {
@@ -124,9 +129,11 @@ public class SolrClientNodeStateProvider implements NodeStateProvider, MapWriter
           if (r.getVariables().containsKey(key)) continue;
           String perReplicaAttrKeyPrefix = "solr.core." + r.getCollection() + "." + r.getShard() + "." + Utils.parseMetricsReplicaName(r.getCollection(), r.getCore()) + ":";
           Suggestion.ConditionType tagType = Suggestion.getTagType(key);
-          if(tagType == null) continue;
-          String perReplicaValue = tagType.metricsAttribute;
-          perReplicaValue = perReplicaValue == null ? key : perReplicaValue;
+          String perReplicaValue = key;
+          if (tagType != null) {
+            perReplicaValue = tagType.metricsAttribute;
+            perReplicaValue = perReplicaValue == null ? key : perReplicaValue;
+          }
           perReplicaAttrKeyPrefix += perReplicaValue;
           keyVsReplica.put(perReplicaAttrKeyPrefix, new Pair<>(key, r));
         }