[Refactor](backend) Refactor the logic of selecting Backend in FE.

fe/fe-core/src/main/java/org/apache/doris/backup/RestoreJob.java

@@ -74,9 +74,6 @@
 import org.apache.doris.thrift.TStorageType;
 import org.apache.doris.thrift.TTaskType;
 
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
-
 import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ArrayListMultimap;
@@ -87,6 +84,9 @@
 import com.google.common.collect.Multimap;
 import com.google.common.collect.Table.Cell;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -1018,7 +1018,8 @@ private Partition resetPartitionForRestore(OlapTable localTbl, OlapTable remoteT
 
                 // replicas
                 try {
-                    Map<Tag, List<Long>> beIds = Catalog.getCurrentSystemInfo().chooseBackendIdByFilters(replicaAlloc, clusterName, null);
+                    Map<Tag, List<Long>> beIds = Catalog.getCurrentSystemInfo()
+                            .selectBackendIdsForReplicaCreation(replicaAlloc, clusterName, null);
                     for (Map.Entry<Tag, List<Long>> entry : beIds.entrySet()) {
                         for (Long beId : entry.getValue()) {
                             long newReplicaId = catalog.getNextId();

fe/fe-core/src/main/java/org/apache/doris/catalog/Catalog.java

@@ -271,6 +271,7 @@
 import org.apache.logging.log4j.Logger;
 import org.codehaus.jackson.map.ObjectMapper;
 
+import javax.annotation.Nullable;
 import java.io.BufferedReader;
 import java.io.DataInputStream;
 import java.io.DataOutputStream;
@@ -296,7 +297,6 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
-import javax.annotation.Nullable;
 
 public class Catalog {
     private static final Logger LOG = LogManager.getLogger(Catalog.class);
@@ -4542,10 +4542,12 @@ private void createTablets(String clusterName, MaterializedIndex index, ReplicaS
                 // This is the first colocate table in the group, or just a normal table,
                 // randomly choose backends
                 if (!Config.disable_storage_medium_check) {
-                    chosenBackendIds = getCurrentSystemInfo().chooseBackendIdByFilters(replicaAlloc, clusterName,
+                    chosenBackendIds =
+                            getCurrentSystemInfo().selectBackendIdsForReplicaCreation(replicaAlloc, clusterName,
                             tabletMeta.getStorageMedium());
                 } else {
-                    chosenBackendIds = getCurrentSystemInfo().chooseBackendIdByFilters(replicaAlloc, clusterName, null);
+                    chosenBackendIds =
+                            getCurrentSystemInfo().selectBackendIdsForReplicaCreation(replicaAlloc, clusterName, null);
                 }
 
                 for (Map.Entry<Tag, List<Long>> entry : chosenBackendIds.entrySet()) {

fe/fe-core/src/main/java/org/apache/doris/catalog/OlapTable.java

@@ -500,7 +500,8 @@ public Status resetIdsForRestore(Catalog catalog, Database db, ReplicaAllocation
 
                     // replicas
                     try {
-                        Map<Tag, List<Long>> tag2beIds = Catalog.getCurrentSystemInfo().chooseBackendIdByFilters(
+                        Map<Tag, List<Long>> tag2beIds =
+                                Catalog.getCurrentSystemInfo().selectBackendIdsForReplicaCreation(
                                 replicaAlloc, db.getClusterName(), null);
                         for (Map.Entry<Tag, List<Long>> entry3 : tag2beIds.entrySet()) {
                             for (Long beId : entry3.getValue()) {

fe/fe-core/src/main/java/org/apache/doris/httpv2/rest/LoadAction.java

@@ -21,12 +21,14 @@
 import org.apache.doris.cluster.ClusterNamespace;
 import org.apache.doris.common.Config;
 import org.apache.doris.common.DdlException;
+import org.apache.doris.common.LoadException;
 import org.apache.doris.httpv2.entity.ResponseEntityBuilder;
 import org.apache.doris.httpv2.entity.RestBaseResult;
 import org.apache.doris.mysql.privilege.PrivPredicate;
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.service.ExecuteEnv;
 import org.apache.doris.system.Backend;
+import org.apache.doris.system.BeSelectionPolicy;
 import org.apache.doris.system.SystemInfoService;
 import org.apache.doris.thrift.TNetworkAddress;
 
@@ -41,12 +43,10 @@
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.servlet.view.RedirectView;
 
-import java.util.List;
-
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
-
 import io.netty.handler.codec.http.HttpHeaderNames;
+import java.util.List;
 
 @RestController
 public class LoadAction extends RestBaseController {
@@ -145,21 +145,7 @@ private Object executeWithoutPassword(HttpServletRequest request,
                     return new RestBaseResult(e.getMessage());
                 }
             } else {
-                // Choose a backend sequentially.
-                SystemInfoService.BeAvailablePredicate beAvailablePredicate =
-                        new SystemInfoService.BeAvailablePredicate(false, false, true);
-                List<Long> backendIds = Catalog.getCurrentSystemInfo().seqChooseBackendIdsByStorageMediumAndTag(
-                        1, beAvailablePredicate, false, clusterName, null, null);
-                if (backendIds == null) {
-                    return new RestBaseResult(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG);
-                }
-
-                Backend backend = Catalog.getCurrentSystemInfo().getBackend(backendIds.get(0));
-                if (backend == null) {
-                    return new RestBaseResult(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG);
-                }
-
-                redirectAddr = new TNetworkAddress(backend.getHost(), backend.getHttpPort());
+                redirectAddr = selectRedirectBackend(clusterName);
             }
 
             LOG.info("redirect load action to destination={}, stream: {}, db: {}, tbl: {}, label: {}",
@@ -194,22 +180,7 @@ private Object executeStreamLoad2PC(HttpServletRequest request, String db) {
                 return new RestBaseResult("No transaction operation(\'commit\' or \'abort\') selected.");
             }
 
-            // Choose a backend sequentially.
-            SystemInfoService.BeAvailablePredicate beAvailablePredicate =
-                    new SystemInfoService.BeAvailablePredicate(false, false, true);
-            List<Long> backendIds = Catalog.getCurrentSystemInfo().seqChooseBackendIdsByStorageMediumAndTag(
-                    1, beAvailablePredicate, false, clusterName, null, null);
-            if (backendIds == null) {
-                return new RestBaseResult(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG);
-            }
-
-            Backend backend = Catalog.getCurrentSystemInfo().getBackend(backendIds.get(0));
-            if (backend == null) {
-                return new RestBaseResult(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG);
-            }
-
-            TNetworkAddress redirectAddr = new TNetworkAddress(backend.getHost(), backend.getHttpPort());
-
+            TNetworkAddress redirectAddr = selectRedirectBackend(clusterName);
             LOG.info("redirect stream load 2PC action to destination={}, db: {}, txn: {}, operation: {}",
                     redirectAddr.toString(), dbName, request.getHeader(TXN_ID_KEY), txnOperation);
 
@@ -220,4 +191,18 @@ private Object executeStreamLoad2PC(HttpServletRequest request, String db) {
             return new RestBaseResult(e.getMessage());
         }
     }
+
+    private TNetworkAddress selectRedirectBackend(String clusterName) throws LoadException {
+        BeSelectionPolicy policy = new BeSelectionPolicy.Builder().setCluster(clusterName).needLoadAvailable().build();
+        List<Long> backendIds = Catalog.getCurrentSystemInfo().selectBackendIdsByPolicy(policy, 1);
+        if (backendIds.isEmpty()) {
+            throw new LoadException(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG + ", policy: " + policy);
+        }
+
+        Backend backend = Catalog.getCurrentSystemInfo().getBackend(backendIds.get(0));
+        if (backend == null) {
+            throw new LoadException(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG + ", policy: " + policy);
+        }
+        return new TNetworkAddress(backend.getHost(), backend.getHttpPort());
+    }
 }

fe/fe-core/src/main/java/org/apache/doris/httpv2/util/LoadSubmitter.java

@@ -19,10 +19,11 @@
 
 import org.apache.doris.catalog.Catalog;
 import org.apache.doris.cluster.ClusterNamespace;
-import org.apache.doris.common.DdlException;
+import org.apache.doris.common.LoadException;
 import org.apache.doris.common.ThreadPoolManager;
 import org.apache.doris.httpv2.rest.UploadAction;
 import org.apache.doris.system.Backend;
+import org.apache.doris.system.BeSelectionPolicy;
 import org.apache.doris.system.SystemInfoService;
 
 import com.google.common.base.Strings;
@@ -136,19 +137,15 @@ private File checkAndGetFile(TmpFileMgr.TmpFile tmpFile) {
             return file;
         }
 
-        private Backend selectOneBackend() throws DdlException {
-            SystemInfoService.BeAvailablePredicate beAvailablePredicate =
-                    new SystemInfoService.BeAvailablePredicate(false, false, true);
-            List<Long> backendIds = Catalog.getCurrentSystemInfo().seqChooseBackendIdsByStorageMediumAndTag(
-                    1, beAvailablePredicate, false,
-                    SystemInfoService.DEFAULT_CLUSTER, null, null);
-            if (backendIds == null) {
-                throw new DdlException("No alive backend");
+        private Backend selectOneBackend() throws LoadException {
+            BeSelectionPolicy policy = new BeSelectionPolicy.Builder().needLoadAvailable().build();
+            List<Long> backendIds = Catalog.getCurrentSystemInfo().selectBackendIdsByPolicy(policy, 1);
+            if (backendIds.isEmpty()) {
+                throw new LoadException(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG + ", policy: " + policy);
             }
-
             Backend backend = Catalog.getCurrentSystemInfo().getBackend(backendIds.get(0));
             if (backend == null) {
-                throw new DdlException("No alive backend");
+                throw new LoadException(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG + ", policy: " + policy);
             }
             return backend;
         }

fe/fe-core/src/main/java/org/apache/doris/qe/InsertStreamTxnExecutor.java

@@ -26,7 +26,7 @@
 import org.apache.doris.rpc.BackendServiceProxy;
 import org.apache.doris.rpc.RpcException;
 import org.apache.doris.system.Backend;
-import org.apache.doris.system.SystemInfoService;
+import org.apache.doris.system.BeSelectionPolicy;
 import org.apache.doris.task.StreamLoadTask;
 import org.apache.doris.thrift.TBrokerRangeDesc;
 import org.apache.doris.thrift.TExecPlanFragmentParams;
@@ -63,13 +63,11 @@ public void beginTransaction(TStreamLoadPutRequest request) throws UserException
         StreamLoadTask streamLoadTask = StreamLoadTask.fromTStreamLoadPutRequest(request);
         StreamLoadPlanner planner = new StreamLoadPlanner(txnEntry.getDb(), (OlapTable) txnEntry.getTable(), streamLoadTask);
         TExecPlanFragmentParams tRequest = planner.plan(streamLoadTask.getId());
-        SystemInfoService.BeAvailablePredicate beAvailablePredicate =
-                new SystemInfoService.BeAvailablePredicate(false, true, true);
-        List<Long> beIds = Catalog.getCurrentSystemInfo().seqChooseBackendIdsByStorageMediumAndTag(
-                1, beAvailablePredicate, false,
-                txnEntry.getDb().getClusterName(), null, null);
-        if (beIds == null || beIds.isEmpty()) {
-            throw new UserException("there is no backend load available or scanNode backend available.");
+        BeSelectionPolicy policy = new BeSelectionPolicy.Builder().setCluster(txnEntry.getDb().getClusterName())
+                .needLoadAvailable().needQueryAvailable().build();
+        List<Long> beIds = Catalog.getCurrentSystemInfo().selectBackendIdsByPolicy(policy, 1);
+        if (beIds.isEmpty()) {
+            throw new UserException("No available backend to match the policy: " + policy);
         }
 
         tRequest.setTxnConf(txnConf).setImportLabel(txnEntry.getLabel());

fe/fe-core/src/main/java/org/apache/doris/qe/MultiLoadMgr.java

@@ -41,6 +41,7 @@
 import org.apache.doris.load.loadv2.LoadJob;
 import org.apache.doris.load.loadv2.LoadTask;
 import org.apache.doris.system.Backend;
+import org.apache.doris.system.BeSelectionPolicy;
 import org.apache.doris.system.SystemInfoService;
 import org.apache.doris.thrift.TMiniLoadRequest;
 import org.apache.doris.thrift.TNetworkAddress;
@@ -91,14 +92,13 @@ public void startMulti(String fullDbName, String label, Map<String, String> prop
             if (infoMap.containsKey(multiLabel)) {
                 throw new LabelAlreadyUsedException(label);
             }
-            MultiLoadDesc multiLoadDesc = new MultiLoadDesc(multiLabel, properties);
-            SystemInfoService.BeAvailablePredicate beAvailablePredicate =
-                    new SystemInfoService.BeAvailablePredicate(false, false, true);
-            List<Long> backendIds = Catalog.getCurrentSystemInfo().seqChooseBackendIdsByStorageMediumAndTag(1,
-                    beAvailablePredicate, false, ConnectContext.get().getClusterName(), null, null);
-            if (backendIds == null) {
-                throw new DdlException(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG);
+            BeSelectionPolicy policy = new BeSelectionPolicy.Builder().setCluster(ConnectContext.get().getClusterName())
+                    .needLoadAvailable().build();
+            List<Long> backendIds = Catalog.getCurrentSystemInfo().selectBackendIdsByPolicy(policy, 1);
+            if (backendIds.isEmpty()) {
+                throw new DdlException(SystemInfoService.NO_BACKEND_LOAD_AVAILABLE_MSG + " policy: " + policy);
             }
+            MultiLoadDesc multiLoadDesc = new MultiLoadDesc(multiLabel, properties);
             multiLoadDesc.setBackendId(backendIds.get(0));
             infoMap.put(multiLabel, multiLoadDesc);
         } finally {

fe/fe-core/src/main/java/org/apache/doris/system/BeSelectionPolicy.java

@@ -0,0 +1,131 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.system;
+
+import org.apache.doris.resource.Tag;
+import org.apache.doris.thrift.TStorageMedium;
+
+import com.google.common.collect.Sets;
+
+import java.util.Set;
+
+/**
+ * Selection policy for building BE nodes
+ */
+public class BeSelectionPolicy {
+    public String cluster = SystemInfoService.DEFAULT_CLUSTER;
+    public boolean needScheduleAvailable = false;
+    public boolean needQueryAvailable = false;
+    public boolean needLoadAvailable = false;
+    // Resource tag. Empty means no need to consider resource tag.
+    public Set<Tag> resourceTags = Sets.newHashSet();
+    // storage medium. null means no need to consider storage medium.
+    public TStorageMedium storageMedium = null;
+    // Check if disk usage reaches limit. false means no need to check.
+    public boolean checkDiskUsage = false;
+    // If set to false, do not select backends on same host.
+    public boolean allowOnSameHost = false;
+
+    private BeSelectionPolicy() {
+
+    }
+
+    public static class Builder {
+        private BeSelectionPolicy policy;
+        public Builder() {
+            policy = new BeSelectionPolicy();
+        }
+
+        public Builder setCluster(String cluster) {
+            policy.cluster = cluster;
+            return this;
+        }
+
+        public Builder needScheduleAvailable() {
+            policy.needScheduleAvailable = true;
+            return this;
+        }
+
+        public Builder needQueryAvailable() {
+            policy.needQueryAvailable = true;
+            return this;
+        }
+
+        public Builder needLoadAvailable() {
+            policy.needLoadAvailable = true;
+            return this;
+        }
+
+        public Builder addTags(Set<Tag> tags) {
+            policy.resourceTags.addAll(tags);
+            return this;
+        }
+
+        public Builder setStorageMedium(TStorageMedium medium) {
+            policy.storageMedium = medium;
+            return this;
+        }
+
+        public Builder needCheckDiskUsage() {
+            policy.checkDiskUsage = true;
+            return this;
+        }
+
+        public Builder allowOnSameHost() {
+            policy.allowOnSameHost = true;
+            return this;
+        }
+
+        public BeSelectionPolicy build() {
+            return policy;
+        }
+    }
+
+    public boolean isMatch(Backend backend) {
+        if (needScheduleAvailable && !backend.isScheduleAvailable()
+                || needQueryAvailable && !backend.isQueryAvailable()
+                || needLoadAvailable && !backend.isLoadAvailable()
+                || !resourceTags.isEmpty() && !resourceTags.contains(backend.getTag())
+                || storageMedium != null && !backend.hasSpecifiedStorageMedium(storageMedium)) {
+            return false;
+        }
+
+        if (checkDiskUsage) {
+            if (storageMedium == null && backend.diskExceedLimit()) {
+                return false;
+            }
+            if (storageMedium != null && backend.diskExceedLimitByStorageMedium(storageMedium)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append("cluster|query|load|schedule|tags|medium: ");
+        sb.append(cluster).append("|");
+        sb.append(needQueryAvailable).append("|");
+        sb.append(needLoadAvailable).append("|");
+        sb.append(needScheduleAvailable).append("|");
+        sb.append(resourceTags).append("|");
+        sb.append(storageMedium);
+        return sb.toString();
+    }
+}