HDDS-10160. Cache sort results in ContainerBalancerSelectionCriteria

hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/AbstractFindTargetGreedy.java

@@ -96,7 +96,7 @@ private void setConfiguration(ContainerBalancerConfiguration conf) {
    * Find a {@link ContainerMoveSelection} consisting of a target and
    * container to move for a source datanode. Favours more under-utilized nodes.
    * @param source Datanode to find a target for
-   * @param candidateContainers Set of candidate containers satisfying
+   * @param container candidate container satisfying
    *                            selection criteria
    *                            {@link ContainerBalancerSelectionCriteria}
    * (DatanodeDetails, Long) method returns true if the size specified in the
@@ -105,29 +105,27 @@ private void setConfiguration(ContainerBalancerConfiguration conf) {
    */
   @Override
   public ContainerMoveSelection findTargetForContainerMove(
-      DatanodeDetails source, Set<ContainerID> candidateContainers) {
+      DatanodeDetails source, ContainerID container) {
     sortTargetForSource(source);
     for (DatanodeUsageInfo targetInfo : potentialTargets) {
       DatanodeDetails target = targetInfo.getDatanodeDetails();
-      for (ContainerID container : candidateContainers) {
-        Set<ContainerReplica> replicas;
-        ContainerInfo containerInfo;
-        try {
-          replicas = containerManager.getContainerReplicas(container);
-          containerInfo = containerManager.getContainer(container);
-        } catch (ContainerNotFoundException e) {
-          logger.warn("Could not get Container {} from Container Manager for " +
-              "obtaining replicas in Container Balancer.", container, e);
-          continue;
-        }
+      Set<ContainerReplica> replicas;
+      ContainerInfo containerInfo;
+      try {
+        replicas = containerManager.getContainerReplicas(container);
+        containerInfo = containerManager.getContainer(container);
+      } catch (ContainerNotFoundException e) {
+        logger.warn("Could not get Container {} from Container Manager for " +
+            "obtaining replicas in Container Balancer.", container, e);
+        return null;
+      }
 
-        if (replicas.stream().noneMatch(
-            replica -> replica.getDatanodeDetails().equals(target)) &&
-            containerMoveSatisfiesPlacementPolicy(container, replicas, source,
-            target) &&
-            canSizeEnterTarget(target, containerInfo.getUsedBytes())) {
-          return new ContainerMoveSelection(target, container);
-        }
+      if (replicas.stream().noneMatch(
+          replica -> replica.getDatanodeDetails().equals(target)) &&
+          containerMoveSatisfiesPlacementPolicy(container, replicas, source,
+          target) &&
+          canSizeEnterTarget(target, containerInfo.getUsedBytes())) {
+        return new ContainerMoveSelection(target, container);
       }
     }
     logger.info("Container Balancer could not find a target for " +

hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java

@@ -31,8 +31,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.NavigableSet;
 import java.util.Set;
 import java.util.TreeSet;
@@ -52,6 +55,7 @@ public class ContainerBalancerSelectionCriteria {
   private Set<ContainerID> selectedContainers;
   private Set<ContainerID> excludeContainers;
   private FindSourceStrategy findSourceStrategy;
+  private Map<DatanodeDetails, NavigableSet<ContainerID>> setMap;
 
   public ContainerBalancerSelectionCriteria(
       ContainerBalancerConfiguration balancerConfiguration,
@@ -66,6 +70,7 @@ public ContainerBalancerSelectionCriteria(
     selectedContainers = new HashSet<>();
     excludeContainers = balancerConfiguration.getExcludeContainers();
     this.findSourceStrategy = findSourceStrategy;
+    this.setMap = new HashMap<>();
   }
 
   /**
@@ -79,38 +84,37 @@ private boolean isContainerReplicatingOrDeleting(ContainerID containerID) {
   }
 
   /**
-   * Gets containers that are suitable for moving based on the following
-   * required criteria:
-   * 1. Container must not be undergoing replication.
-   * 2. Container must not already be selected for balancing.
-   * 3. Container size should be closer to 5GB.
-   * 4. Container must not be in the configured exclude containers list.
-   * 5. Container should be closed.
-   * 6. If the {@link LegacyReplicationManager} is enabled, then the container should not be an EC container.
-   * @param node DatanodeDetails for which to find candidate containers.
-   * @return NavigableSet of candidate containers that satisfy the criteria.
+   * Get ContainerID Set for the Datanode, it will be returned as NavigableSet
+   * Since sorting will be time-consuming, the Set will be cached.
+   *
+   * @param node source datanode
+   * @return cached Navigable ContainerID Set
    */
-  public NavigableSet<ContainerID> getCandidateContainers(
-      DatanodeDetails node, long sizeMovedAlready) {
-    NavigableSet<ContainerID> containerIDSet =
-        new TreeSet<>(orderContainersByUsedBytes().reversed());
-    try {
-      containerIDSet.addAll(nodeManager.getContainers(node));
-    } catch (NodeNotFoundException e) {
-      LOG.warn("Could not find Datanode {} while selecting candidate " +
-          "containers for Container Balancer.", node.toString(), e);
-      return containerIDSet;
-    }
-    if (excludeContainers != null) {
-      containerIDSet.removeAll(excludeContainers);
+  public Set<ContainerID> getContainerIDSet(DatanodeDetails node) {
+    // Check if the node is registered at the beginning
+    if (!nodeManager.isNodeRegistered(node)) {
+      return Collections.emptySet();
     }
-    if (selectedContainers != null) {
-      containerIDSet.removeAll(selectedContainers);
+
+    try {
+      // Initialize containerSet for node
+      setMap.computeIfAbsent(node, n -> {
+        try {
+          addNodeToSetMap(n);
+          return setMap.get(n);
+        } catch (NodeNotFoundException e) {
+          LOG.warn("Could not find Datanode {} while selecting candidate " +
+              "containers for Container Balancer.", n.toString(), e);
+          return null;
+        }
+      });
+    } catch (Exception e) {
+      LOG.error("An unexpected error occurred while processing the node.", e);
+      setMap.remove(node);
+      return Collections.emptySet();
     }
 
-    containerIDSet.removeIf(
-        containerID -> shouldBeExcluded(containerID, node, sizeMovedAlready));
-    return containerIDSet;
+    return setMap.get(node);
   }
 
   /**
@@ -165,7 +169,21 @@ private boolean isECContainerAndLegacyRMEnabled(ContainerInfo container) {
         && replicationManager.getConfig().isLegacyEnabled();
   }
 
-  private boolean shouldBeExcluded(ContainerID containerID,
+  /**
+   * Gets containers that are suitable for moving based on the following
+   * required criteria:
+   * 1. Container must not be in ExcludedContainers.
+   * 2. Container must not be in SelectedContainers.
+   * 3. Container must not be undergoing replication.
+   * 4. Container must not already be selected for balancing.
+   * 5. Container size should be closer to 5GB.
+   * 6. Container must not be in the configured exclude containers list.
+   * 7. Container should be closed.
+   * 8. If the {@link LegacyReplicationManager} is enabled, then the container should not be an EC container.
+   * @param node DatanodeDetails for which to find candidate containers.
+   * @return Set of candidate containers that satisfy the criteria.
+   */
+  public boolean shouldBeExcluded(ContainerID containerID,
       DatanodeDetails node, long sizeMovedAlready) {
     ContainerInfo container;
     try {
@@ -175,7 +193,8 @@ private boolean shouldBeExcluded(ContainerID containerID,
           "candidate container. Excluding it.", containerID);
       return true;
     }
-    return !isContainerClosed(container, node) || isECContainerAndLegacyRMEnabled(container) ||
+    return excludeContainers.contains(containerID) || selectedContainers.contains(containerID) ||
+        !isContainerClosed(container, node) || isECContainerAndLegacyRMEnabled(container) ||
         isContainerReplicatingOrDeleting(containerID) ||
         !findSourceStrategy.canSizeLeaveSource(node, container.getUsedBytes())
         || breaksMaxSizeToMoveLimit(container.containerID(),
@@ -242,4 +261,19 @@ public void setSelectedContainers(
     this.selectedContainers = selectedContainers;
   }
 
+
+  private void addNodeToSetMap(DatanodeDetails node)
+      throws NodeNotFoundException {
+    NavigableSet<ContainerID> newSet =
+        new TreeSet<>(orderContainersByUsedBytes().reversed());
+    Set<ContainerID> idSet = nodeManager.getContainers(node);
+    if (excludeContainers != null) {
+      idSet.removeAll(excludeContainers);
+    }
+    if (selectedContainers != null) {
+      idSet.removeAll(selectedContainers);
+    }
+    newSet.addAll(idSet);
+    setMap.put(node, newSet);
+  }
 }

hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerTask.java

@@ -50,7 +50,6 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.NavigableSet;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -692,11 +691,10 @@ private long cancelMovesThatExceedTimeoutDuration() {
    * @return ContainerMoveSelection containing the selected target and container
    */
   private ContainerMoveSelection matchSourceWithTarget(DatanodeDetails source) {
-    NavigableSet<ContainerID> candidateContainers =
-        selectionCriteria.getCandidateContainers(source,
-            sizeScheduledForMoveInLatestIteration);
+    Set<ContainerID> sourceContainerIDSet =
+        selectionCriteria.getContainerIDSet(source);
 
-    if (candidateContainers.isEmpty()) {
+    if (sourceContainerIDSet.isEmpty()) {
       if (LOG.isDebugEnabled()) {
         LOG.debug("ContainerBalancer could not find any candidate containers " +
             "for datanode {}", source.getUuidString());
@@ -708,9 +706,23 @@ private ContainerMoveSelection matchSourceWithTarget(DatanodeDetails source) {
       LOG.debug("ContainerBalancer is finding suitable target for source " +
           "datanode {}", source.getUuidString());
     }
-    ContainerMoveSelection moveSelection =
-        findTargetStrategy.findTargetForContainerMove(
-            source, candidateContainers);
+
+    ContainerMoveSelection moveSelection = null;
+    Set<ContainerID> toRemoveContainerIds = new HashSet<>();
+    for (ContainerID containerId: sourceContainerIDSet) {
+      if (selectionCriteria.shouldBeExcluded(containerId, source,
+          sizeScheduledForMoveInLatestIteration)) {
+        toRemoveContainerIds.add(containerId);
+        continue;
+      }
+      moveSelection = findTargetStrategy.findTargetForContainerMove(source,
+          containerId);
+      if (moveSelection != null) {
+        break;
+      }
+    }
+    // Update cached containerIDSet in setMap
+    sourceContainerIDSet.removeAll(toRemoveContainerIds);
 
     if (moveSelection == null) {
       if (LOG.isDebugEnabled()) {

hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/FindTargetStrategy.java

@@ -25,7 +25,6 @@
 import javax.annotation.Nonnull;
 import java.util.Collection;
 import java.util.List;
-import java.util.Set;
 
 /**
  * This interface can be used to implement strategies to find a target for a
@@ -40,7 +39,7 @@ public interface FindTargetStrategy {
    * enter a potential target.
    *
    * @param source Datanode to find a target for
-   * @param candidateContainers Set of candidate containers satisfying
+   * @param candidateContainer candidate containers satisfying
    *                            selection criteria
    *                            {@link ContainerBalancerSelectionCriteria}
    * (DatanodeDetails, Long) method returns true if the size specified in the
@@ -49,7 +48,7 @@ public interface FindTargetStrategy {
    * selected container
    */
   ContainerMoveSelection findTargetForContainerMove(
-      DatanodeDetails source, Set<ContainerID> candidateContainers);
+      DatanodeDetails source, ContainerID candidateContainer);
 
   /**
    * increase the Entering size of a candidate target data node.

hadoop-hdds/server-scm/src/test/java/org/apache/hadoop/hdds/scm/container/MockNodeManager.java

@@ -786,7 +786,7 @@ public List<SCMCommand> processHeartbeat(DatanodeDetails datanodeDetails,
   @Override
   public Boolean isNodeRegistered(
       DatanodeDetails datanodeDetails) {
-    return false;
+    return healthyNodes.contains(datanodeDetails);
   }
 
   @Override