SOLR-17289: OrderedNodePlacementPlugin: optimize don't loop collections

apache · May 14, 2024 · ee1286e · ee1286e
1 parent d141875
commit ee1286e
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
diff --git a/solr/core/src/java/org/apache/solr/cluster/placement/plugins/AffinityPlacementFactory.java b/solr/core/src/java/org/apache/solr/cluster/placement/plugins/AffinityPlacementFactory.java
@@ -256,6 +256,28 @@ private static final class AffinityPlacementContext {
       private boolean doSpreadAcrossDomains;
     }
 
+    @Override
+    protected Map<Node, WeightedNode> getWeightedNodes(
+        PlacementContext placementContext,
+        Set<Node> nodes,
+        Iterable<SolrCollection> relevantCollections,
+        boolean skipNodesWithErrors)
+        throws PlacementException {
+      Map<Node, WeightedNode> weightedNodes =
+          getBaseWeightedNodes(placementContext, nodes, relevantCollections, skipNodesWithErrors);
+
+      Iterable<SolrCollection> initCollections;
+      if (withCollections.isEmpty() && withCollectionShards.isEmpty()) {
+        initCollections = relevantCollections;
+      } else {
+        // TODO this doesn't scale to many collections
+        initCollections = placementContext.getCluster().collections();
+      }
+      initWeightedNodeReplicas(weightedNodes, initCollections);
+
+      return weightedNodes;
+    }
+
     @Override
     protected Map<Node, WeightedNode> getBaseWeightedNodes(
         PlacementContext placementContext,

diff --git a/solr/core/src/java/org/apache/solr/cluster/placement/plugins/OrderedNodePlacementPlugin.java b/solr/core/src/java/org/apache/solr/cluster/placement/plugins/OrderedNodePlacementPlugin.java
@@ -83,6 +83,7 @@ public List<PlacementPlan> computePlacements(
 
     Collection<WeightedNode> weightedNodes =
         getWeightedNodes(placementContext, allNodes, allCollections, true).values();
+
     while (!pendingRequests.isEmpty()) {
       PendingPlacementRequest request = pendingRequests.poll();
       if (!request.isPending()) {
@@ -301,6 +302,7 @@ public BalancePlan computeBalancing(
         .createBalancePlan(balanceRequest, replicaMovements);
   }
 
+  /** Compute weighted nodes, to include init'ing each replica. */
   protected Map<Node, WeightedNode> getWeightedNodes(
       PlacementContext placementContext,
       Set<Node> nodes,
@@ -310,7 +312,14 @@ protected Map<Node, WeightedNode> getWeightedNodes(
     Map<Node, WeightedNode> weightedNodes =
         getBaseWeightedNodes(placementContext, nodes, relevantCollections, skipNodesWithErrors);
 
-    for (SolrCollection collection : placementContext.getCluster().collections()) {
+    initWeightedNodeReplicas(weightedNodes, relevantCollections);
+
+    return weightedNodes;
+  }
+
+  protected static void initWeightedNodeReplicas(
+      Map<Node, WeightedNode> weightedNodes, Iterable<SolrCollection> collections) {
+    for (SolrCollection collection : collections) {
       for (Shard shard : collection.shards()) {
         for (Replica replica : shard.replicas()) {
           WeightedNode weightedNode = weightedNodes.get(replica.getNode());
@@ -320,8 +329,6 @@ protected Map<Node, WeightedNode> getWeightedNodes(
         }
       }
     }
-
-    return weightedNodes;
   }
 
   protected abstract Map<Node, WeightedNode> getBaseWeightedNodes(