[lucene] 01/04: SOLR-15130: WIP.

[dw...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dweiss pushed a commit to branch jira/solr-15130
in repository https://gitbox.apache.org/repos/asf/lucene.git

commit 51aa6c331e08d8d7e71b65c9baee17c2a3ba5638
Author: Andrzej Bialecki <ab...@apache.org>
AuthorDate: Mon Feb 8 15:25:28 2021 +0100

    SOLR-15130: WIP.
---
 .../placement/plugins/AffinityPlacementConfig.java |  3 ++
 .../plugins/AffinityPlacementFactory.java          | 40 +++++++++++++++++++---
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/solr/core/src/java/org/apache/solr/cluster/placement/plugins/AffinityPlacementConfig.java b/solr/core/src/java/org/apache/solr/cluster/placement/plugins/AffinityPlacementConfig.java
index d9579bc..062750a 100644
--- a/solr/core/src/java/org/apache/solr/cluster/placement/plugins/AffinityPlacementConfig.java
+++ b/solr/core/src/java/org/apache/solr/cluster/placement/plugins/AffinityPlacementConfig.java
@@ -19,6 +19,7 @@ package org.apache.solr.cluster.placement.plugins;
 
 import org.apache.solr.cluster.placement.PlacementPluginConfig;
 import org.apache.solr.common.annotation.JsonProperty;
+import org.apache.solr.common.params.CollectionAdminParams;
 
 import java.util.Map;
 import java.util.Objects;
@@ -28,6 +29,8 @@ import java.util.Objects;
  */
 public class AffinityPlacementConfig implements PlacementPluginConfig {
 
+  public static final String COLLECTION_NODE_TYPE_PROPERTY = CollectionAdminParams.PROPERTY_PREFIX + "placement.affinity.node_type";
+
   public static final long DEFAULT_MINIMAL_FREE_DISK_GB = 20L;
   public static final long DEFAULT_PRIORITIZED_FREE_DISK_GB = 100L;
 
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
index eaec4ab..64aa9d5 100644
--- 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
@@ -23,6 +23,7 @@ import org.apache.solr.cluster.*;
 import org.apache.solr.cluster.placement.*;
 import org.apache.solr.cluster.placement.impl.NodeMetricImpl;
 import org.apache.solr.common.util.Pair;
+import org.apache.solr.common.util.StrUtils;
 import org.apache.solr.common.util.SuppressForbidden;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -131,6 +132,8 @@ public class AffinityPlacementFactory implements PlacementPluginFactory<Affinity
    */
   public static final String REPLICA_TYPE_SYSPROP = "replica_type";
 
+  public static final String NODE_TYPE_SYSPROP = "node_type";
+
   /**
    * This is the "AZ" name for nodes that do not define an AZ. Should not match a real AZ name (I think we're safe)
    */
@@ -210,17 +213,23 @@ public class AffinityPlacementFactory implements PlacementPluginFactory<Affinity
       Set<Node> nodes = request.getTargetNodes();
       SolrCollection solrCollection = request.getCollection();
 
-      nodes = filterNodesWithCollection(placementContext.getCluster(), request, nodes);
-
       // Request all needed attributes
       AttributeFetcher attributeFetcher = placementContext.getAttributeFetcher();
-      attributeFetcher.requestNodeSystemProperty(AVAILABILITY_ZONE_SYSPROP).requestNodeSystemProperty(REPLICA_TYPE_SYSPROP);
+      attributeFetcher
+          .requestNodeSystemProperty(AVAILABILITY_ZONE_SYSPROP)
+          .requestNodeSystemProperty(NODE_TYPE_SYSPROP)
+          .requestNodeSystemProperty(REPLICA_TYPE_SYSPROP);
       attributeFetcher
           .requestNodeMetric(NodeMetricImpl.NUM_CORES)
           .requestNodeMetric(NodeMetricImpl.FREE_DISK_GB);
       attributeFetcher.fetchFrom(nodes);
       final AttributeValues attrValues = attributeFetcher.fetchAttributes();
 
+      // filter out nodes that don't meet the `withCollection` constraint
+      nodes = filterNodesWithCollection(placementContext.getCluster(), request, attrValues, nodes);
+      // filter out nodes that don't match the "node types" specified in the collection props
+      nodes = filterNodesByNodeType(placementContext.getCluster(), request, attrValues, nodes);
+
       // Split the set of nodes into 3 sets of nodes accepting each replica type (sets can overlap if nodes accept multiple replica types)
       // These subsets sets are actually maps, because we capture the number of cores (of any replica type) present on each node.
       // Also get the number of currently existing cores per node, so we can keep update as we place new cores to not end up
@@ -632,7 +641,7 @@ public class AffinityPlacementFactory implements PlacementPluginFactory<Affinity
       }
     }
 
-    private Set<Node> filterNodesWithCollection(Cluster cluster, PlacementRequest request, Set<Node> initialNodes) throws PlacementException {
+    private Set<Node> filterNodesWithCollection(Cluster cluster, PlacementRequest request, AttributeValues attributeValues, Set<Node> initialNodes) throws PlacementException {
       // if there's a `withCollection` constraint for this collection then remove nodes
       // that are not eligible
       String withCollectionName = withCollections.get(request.getCollection().getName());
@@ -658,6 +667,29 @@ public class AffinityPlacementFactory implements PlacementPluginFactory<Affinity
       return filteredNodes;
     }
 
+    private Set<Node> filterNodesByNodeType(Cluster cluster, PlacementRequest request, AttributeValues attributeValues, Set<Node> initialNodes) throws PlacementException {
+      String collProperty = request.getCollection().getCustomProperty(AffinityPlacementConfig.COLLECTION_NODE_TYPE_PROPERTY);
+      if (collProperty == null) {
+        // no filtering by node type
+        return initialNodes;
+      }
+      Set<String> collNodeTypes = Set.copyOf(StrUtils.splitSmart(collProperty, ','));
+      Set<Node> filteredNodes = initialNodes.stream()
+          .filter(n -> {
+            Optional<String> nodePropOpt = attributeValues.getSystemProperty(n, AffinityPlacementConfig.COLLECTION_NODE_TYPE_PROPERTY);
+            if (!nodePropOpt.isPresent()) {
+              return false;
+            }
+            Set<String> nodeTypes = Set.copyOf(StrUtils.splitSmart(nodePropOpt.get(), ','));
+            nodeTypes.retainAll(collNodeTypes);
+            return !nodeTypes.isEmpty();
+          }).collect(Collectors.toSet());
+      if (filteredNodes.isEmpty()) {
+        throw new PlacementException("There are no nodes with types: " + collNodeTypes + " expected by collection " + request.getCollection().getName());
+      }
+      return filteredNodes;
+    }
+
     /**
      * Comparator implementing the placement strategy based on free space and number of cores: we want to place new replicas
      * on nodes with the less number of cores, but only if they do have enough disk space (expressed as a threshold value).