From b400d4722240b5b13a78b96f1c0042f28e34b04a Mon Sep 17 00:00:00 2001
From: guohao1 <guohao1@360.cn>
Date: Fri, 17 May 2024 11:45:42 +0800
Subject: [PATCH 1/2] HDDS-10871. ContainerBalancerSelectionCriteria memory
 occupancy optimization

---
 .../ContainerBalancerSelectionCriteria.java   | 19 ++++++++++---------
 .../balancer/ContainerBalancerTask.java       |  5 ++---
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java b/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java
index da1b8741cfd3..e41c850f24ca 100644
--- a/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java
+++ b/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java
@@ -52,7 +52,7 @@ public class ContainerBalancerSelectionCriteria {
   private NodeManager nodeManager;
   private ReplicationManager replicationManager;
   private ContainerManager containerManager;
-  private Set<ContainerID> selectedContainers;
+  private Map<ContainerID, DatanodeDetails> containerToSourceMap;
   private Set<ContainerID> excludeContainers;
   private Set<ContainerID> excludeContainersDueToFailure;
   private FindSourceStrategy findSourceStrategy;
@@ -63,18 +63,23 @@ public ContainerBalancerSelectionCriteria(
       NodeManager nodeManager,
       ReplicationManager replicationManager,
       ContainerManager containerManager,
-      FindSourceStrategy findSourceStrategy) {
+      FindSourceStrategy findSourceStrategy,
+      Map<ContainerID, DatanodeDetails> containerToSourceMap) {
     this.balancerConfiguration = balancerConfiguration;
     this.nodeManager = nodeManager;
     this.replicationManager = replicationManager;
     this.containerManager = containerManager;
-    selectedContainers = new HashSet<>();
+    this.containerToSourceMap = containerToSourceMap;
     excludeContainersDueToFailure = new HashSet<>();
     excludeContainers = balancerConfiguration.getExcludeContainers();
     this.findSourceStrategy = findSourceStrategy;
     this.setMap = new HashMap<>();
   }
 
+  public Set<ContainerID> getSelectedContainers() {
+    return containerToSourceMap.keySet();
+  }
+
   /**
    * Checks whether container is currently undergoing replication or deletion.
    *
@@ -177,7 +182,7 @@ public boolean shouldBeExcluded(ContainerID containerID,
       return true;
     }
     return excludeContainers.contains(containerID) || excludeContainersDueToFailure.contains(containerID) ||
-        selectedContainers.contains(containerID) ||
+        getSelectedContainers().contains(containerID) ||
         !isContainerClosed(container, node) || isECContainerAndLegacyRMEnabled(container) ||
         isContainerReplicatingOrDeleting(containerID) ||
         !findSourceStrategy.canSizeLeaveSource(node, container.getUsedBytes())
@@ -240,11 +245,6 @@ public void setExcludeContainers(
     this.excludeContainers = excludeContainers;
   }
 
-  public void setSelectedContainers(
-      Set<ContainerID> selectedContainers) {
-    this.selectedContainers = selectedContainers;
-  }
-
   public void addToExcludeDueToFailContainers(ContainerID container) {
     this.excludeContainersDueToFailure.add(container);
   }
@@ -261,6 +261,7 @@ private NavigableSet<ContainerID> getCandidateContainers(DatanodeDetails node) {
       if (excludeContainersDueToFailure != null) {
         idSet.removeAll(excludeContainersDueToFailure);
       }
+      Set<ContainerID> selectedContainers = getSelectedContainers();
       if (selectedContainers != null) {
         idSet.removeAll(selectedContainers);
       }
diff --git a/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerTask.java b/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerTask.java
index 0bfedd43960f..a4072c8d4f9f 100644
--- a/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerTask.java
+++ b/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerTask.java
@@ -447,7 +447,8 @@ private boolean initializeIteration() {
     }
 
     selectionCriteria = new ContainerBalancerSelectionCriteria(config,
-        nodeManager, replicationManager, containerManager, findSourceStrategy);
+        nodeManager, replicationManager, containerManager, findSourceStrategy,
+        containerToSourceMap);
     return true;
   }
 
@@ -927,8 +928,6 @@ private void updateTargetsAndSelectionCriteria(
     containerToTargetMap.put(containerID, target);
     selectedTargets.add(target);
     selectedSources.add(source);
-    selectionCriteria.setSelectedContainers(
-        new HashSet<>(containerToSourceMap.keySet()));
   }
 
   /**

From 95e5bfca1d21a267fbf7f4faad108096316d8dec Mon Sep 17 00:00:00 2001
From: guohao1 <guohao1@360.cn>
Date: Thu, 30 May 2024 15:34:28 +0800
Subject: [PATCH 2/2] code review

---
 .../balancer/ContainerBalancerSelectionCriteria.java  | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java b/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java
index e41c850f24ca..3965a1593238 100644
--- a/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java
+++ b/hadoop-hdds/server-scm/src/main/java/org/apache/hadoop/hdds/scm/container/balancer/ContainerBalancerSelectionCriteria.java
@@ -76,10 +76,6 @@ public ContainerBalancerSelectionCriteria(
     this.setMap = new HashMap<>();
   }
 
-  public Set<ContainerID> getSelectedContainers() {
-    return containerToSourceMap.keySet();
-  }
-
   /**
    * Checks whether container is currently undergoing replication or deletion.
    *
@@ -182,7 +178,7 @@ public boolean shouldBeExcluded(ContainerID containerID,
       return true;
     }
     return excludeContainers.contains(containerID) || excludeContainersDueToFailure.contains(containerID) ||
-        getSelectedContainers().contains(containerID) ||
+        containerToSourceMap.containsKey(containerID) ||
         !isContainerClosed(container, node) || isECContainerAndLegacyRMEnabled(container) ||
         isContainerReplicatingOrDeleting(containerID) ||
         !findSourceStrategy.canSizeLeaveSource(node, container.getUsedBytes())
@@ -261,10 +257,7 @@ private NavigableSet<ContainerID> getCandidateContainers(DatanodeDetails node) {
       if (excludeContainersDueToFailure != null) {
         idSet.removeAll(excludeContainersDueToFailure);
       }
-      Set<ContainerID> selectedContainers = getSelectedContainers();
-      if (selectedContainers != null) {
-        idSet.removeAll(selectedContainers);
-      }
+      idSet.removeAll(containerToSourceMap.keySet());
       newSet.addAll(idSet);
       return newSet;
     } catch (NodeNotFoundException e) {