fix

apache · Dec 13, 2024 · 9f127d5 · 9f127d5
1 parent 398af82
commit 9f127d5
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 70 deletions.
diff --git a/...main/java/org/apache/doris/nereids/trees/plans/commands/AbstractRebalanceDiskCommand.java b/...main/java/org/apache/doris/nereids/trees/plans/commands/AbstractRebalanceDiskCommand.java
diff --git a/...n/java/org/apache/doris/nereids/trees/plans/commands/AdminCancelRebalanceDiskCommand.java b/...n/java/org/apache/doris/nereids/trees/plans/commands/AdminCancelRebalanceDiskCommand.java
@@ -22,22 +22,27 @@
 import org.apache.doris.common.DdlException;
 import org.apache.doris.common.ErrorCode;
 import org.apache.doris.common.ErrorReport;
+import org.apache.doris.common.util.NetUtils;
 import org.apache.doris.mysql.privilege.PrivPredicate;
 import org.apache.doris.nereids.trees.plans.PlanType;
 import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.qe.StmtExecutor;
 import org.apache.doris.system.Backend;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
 import org.apache.log4j.LogManager;
 import org.apache.log4j.Logger;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * admin cancel rebalance disk
  */
-public class AdminCancelRebalanceDiskCommand extends AbstractRebalanceDiskCommand {
+public class AdminCancelRebalanceDiskCommand extends Command implements NoForward {
     private static final Logger LOG = LogManager.getLogger(AdminCancelRebalanceDiskCommand.class);
     private List<String> backends;
 
@@ -68,6 +73,28 @@ private void handleCancelRebalanceDisk() throws AnalysisException {
         Env.getCurrentEnv().getTabletScheduler().cancelRebalanceDisk(rebalanceDiskBackends);
     }
 
+    private List<Backend> getNeedRebalanceDiskBackends(List<String> backends) throws AnalysisException {
+        ImmutableMap<Long, Backend> backendsInfo = Env.getCurrentSystemInfo().getAllBackendsByAllCluster();
+        List<Backend> needRebalanceDiskBackends = Lists.newArrayList();
+        if (backends == null) {
+            needRebalanceDiskBackends.addAll(backendsInfo.values());
+        } else {
+            Map<String, Long> backendsID = new HashMap<>();
+            for (Backend backend : backendsInfo.values()) {
+                backendsID.put(
+                        NetUtils.getHostPortInAccessibleFormat(backend.getHost(), backend.getHeartbeatPort()),
+                        backend.getId());
+            }
+            for (String be : backends) {
+                if (backendsID.containsKey(be)) {
+                    needRebalanceDiskBackends.add(backendsInfo.get(backendsID.get(be)));
+                    backendsID.remove(be);
+                }
+            }
+        }
+        return needRebalanceDiskBackends;
+    }
+
     @Override
     protected void checkSupportedInCloudMode(ConnectContext ctx) throws DdlException {
         LOG.info("AdminCancelRebalanceDiskCommand not supported in cloud mode");

diff --git a/...rc/main/java/org/apache/doris/nereids/trees/plans/commands/AdminRebalanceDiskCommand.java b/...rc/main/java/org/apache/doris/nereids/trees/plans/commands/AdminRebalanceDiskCommand.java
@@ -22,22 +22,27 @@
 import org.apache.doris.common.DdlException;
 import org.apache.doris.common.ErrorCode;
 import org.apache.doris.common.ErrorReport;
+import org.apache.doris.common.util.NetUtils;
 import org.apache.doris.mysql.privilege.PrivPredicate;
 import org.apache.doris.nereids.trees.plans.PlanType;
 import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.qe.StmtExecutor;
 import org.apache.doris.system.Backend;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
 import org.apache.log4j.LogManager;
 import org.apache.log4j.Logger;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * admin rebalance disk
  */
-public class AdminRebalanceDiskCommand extends AbstractRebalanceDiskCommand {
+public class AdminRebalanceDiskCommand extends Command implements NoForward {
     private static final Logger LOG = LogManager.getLogger(AdminRebalanceDiskCommand.class);
     private final long timeoutS = 24 * 3600; // default 24 hours
     private List<String> backends;
@@ -69,6 +74,28 @@ private void handleRebalanceDisk() throws AnalysisException {
         Env.getCurrentEnv().getTabletScheduler().rebalanceDisk(rebalanceDiskBackends, timeoutS);
     }
 
+    private List<Backend> getNeedRebalanceDiskBackends(List<String> backends) throws AnalysisException {
+        ImmutableMap<Long, Backend> backendsInfo = Env.getCurrentSystemInfo().getAllBackendsByAllCluster();
+        List<Backend> needRebalanceDiskBackends = Lists.newArrayList();
+        if (backends == null) {
+            needRebalanceDiskBackends.addAll(backendsInfo.values());
+        } else {
+            Map<String, Long> backendsID = new HashMap<>();
+            for (Backend backend : backendsInfo.values()) {
+                backendsID.put(
+                        NetUtils.getHostPortInAccessibleFormat(backend.getHost(), backend.getHeartbeatPort()),
+                        backend.getId());
+            }
+            for (String be : backends) {
+                if (backendsID.containsKey(be)) {
+                    needRebalanceDiskBackends.add(backendsInfo.get(backendsID.get(be)));
+                    backendsID.remove(be);
+                }
+            }
+        }
+        return needRebalanceDiskBackends;
+    }
+
     @Override
     protected void checkSupportedInCloudMode(ConnectContext ctx) throws DdlException {
         LOG.info("AdminRebalanceDiskCommand not supported in cloud mode");