wangbo commented on code in PR #28443: URL: https://github.com/apache/doris/pull/28443#discussion_r1429753869
########## fe/fe-core/src/main/java/org/apache/doris/resource/workloadschedpolicy/WorkloadSchedPolicyMgr.java: ########## @@ -0,0 +1,509 @@ +// 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.resource.workloadschedpolicy; + +import org.apache.doris.analysis.AlterWorkloadSchedPolicyStmt; +import org.apache.doris.analysis.CreateWorkloadSchedPolicyStmt; +import org.apache.doris.analysis.DropWorkloadSchedPolicyStmt; +import org.apache.doris.analysis.UserIdentity; +import org.apache.doris.catalog.Env; +import org.apache.doris.common.Config; +import org.apache.doris.common.UserException; +import org.apache.doris.common.io.Text; +import org.apache.doris.common.io.Writable; +import org.apache.doris.common.proc.BaseProcResult; +import org.apache.doris.common.proc.ProcResult; +import org.apache.doris.common.util.DebugUtil; +import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.persist.gson.GsonPostProcessable; +import org.apache.doris.persist.gson.GsonUtils; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.service.ExecuteEnv; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Maps; +import com.google.gson.annotations.SerializedName; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import java.io.DataInput; +import java.io.DataOutput; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.PriorityQueue; +import java.util.Queue; +import java.util.Set; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +public class WorkloadSchedPolicyMgr implements Writable, GsonPostProcessable { + + private static final Logger LOG = LogManager.getLogger(WorkloadSchedPolicyMgr.class); + + @SerializedName(value = "idToPolicy") + private Map<Long, WorkloadSchedPolicy> idToPolicy = Maps.newConcurrentMap(); + private Map<String, WorkloadSchedPolicy> nameToPolicy = Maps.newHashMap(); + + private PolicyProcNode policyProcNode = new PolicyProcNode(); + + public static final ImmutableList<String> WORKLOAD_SCHED_POLICY_NODE_TITLE_NAMES + = new ImmutableList.Builder<String>() + .add("Id").add("Name").add("ItemName").add("ItemValue") + .build(); + + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + + public static Comparator<WorkloadSchedPolicy> policyComparator = new Comparator<WorkloadSchedPolicy>() { + @Override + public int compare(WorkloadSchedPolicy p1, WorkloadSchedPolicy p2) { + return p2.getPriority() - p1.getPriority(); + } + }; + + private Thread policyExecThread = new Thread() { + @Override + public void run() { + while (true) { + try { + // 1 get query info map + Map<Integer, ConnectContext> connectMap = ExecuteEnv.getInstance().getScheduler() + .getConnectionMap(); + List<WorkloadQueryInfo> queryInfoList = new ArrayList<>(); + + // a snapshot for connect context + Set<Integer> keySet = new HashSet<>(); + keySet.addAll(connectMap.keySet()); + + for (Integer connectId : keySet) { + ConnectContext cctx = connectMap.get(connectId); + if (cctx == null) { + continue; + } + + String username = cctx.getQualifiedUser(); + long queryTime = System.currentTimeMillis() - cctx.getStartTime(); + + WorkloadQueryInfo policyQueryInfo = new WorkloadQueryInfo(); + policyQueryInfo.queryId = DebugUtil.printId(cctx.queryId()); + policyQueryInfo.tUniqueId = cctx.queryId(); + policyQueryInfo.context = cctx; + policyQueryInfo.metricMap = new HashMap<>(); + policyQueryInfo.metricMap.put(WorkloadMetricType.username, username); + policyQueryInfo.metricMap.put(WorkloadMetricType.query_time, String.valueOf(queryTime)); + + queryInfoList.add(policyQueryInfo); + } + + // 2 exec policy + if (queryInfoList.size() > 0) { + execPolicy(queryInfoList); + } + } catch (Throwable t) { + LOG.error("[policy thread]error happens when exec policy"); + } + + // 3 sleep + try { + Thread.sleep(Config.workload_sched_policy_interval_ms); + } catch (InterruptedException e) { + LOG.error("error happends when policy exec thread sleep"); + } + } + } + }; + + public void start() { + policyExecThread.start(); + } + + public void createWorkloadSchedPolicy(CreateWorkloadSchedPolicyStmt createStmt) throws UserException { + String policyName = createStmt.getPolicyName(); + + // 1 create condition + List<WorkloadConditionMeta> originConditions = createStmt.getConditions(); + List<WorkloadCondition> policyConditionList = new ArrayList<>(); + for (WorkloadConditionMeta cm : originConditions) { + WorkloadCondition cond = WorkloadCondition.createWorkloadCondition(cm.metricName, cm.op, + cm.value); + policyConditionList.add(cond); + } + + // 2 create action + Map<String, String> originActions = createStmt.getActions(); + List<WorkloadAction> policyActionList = new ArrayList<>(); + for (Map.Entry<String, String> actionEntry : originActions.entrySet()) { + String actionName = actionEntry.getKey(); + String actionArgs = actionEntry.getValue(); + + // we need convert wgName to wgId, because wgName may change + if (WorkloadActionType.move_query_to_group.toString().equals(actionName)) { + Long wgId = Env.getCurrentEnv().getWorkloadGroupMgr().getWorkloadGroupIdByName(actionArgs); + if (wgId == null) { + throw new UserException( + "can not find workload group " + actionArgs + " when set workload sched policy"); + } + actionEntry.setValue(wgId.toString()); + } + + WorkloadAction ret = WorkloadAction.createWorkloadAction(actionEntry.getKey(), actionEntry.getValue()); + policyActionList.add(ret); + } + checkPolicyActionConflicts(policyActionList); + + // 3 create policy + Map<String, String> propMap = createStmt.getProperties(); + if (propMap == null) { + propMap = new HashMap<>(); + } + if (propMap.size() != 0) { + checkProperties(propMap); + } + writeLock(); + try { + if (nameToPolicy.containsKey(createStmt.getPolicyName())) { + if (createStmt.isIfNotExists()) { + return; + } else { + throw new UserException("workload schedule policy " + policyName + " already exists "); + } + } + long id = Env.getCurrentEnv().getNextId(); + WorkloadSchedPolicy policy = new WorkloadSchedPolicy(id, policyName, + policyConditionList, policyActionList, propMap); + policy.setConditionMeta(originConditions); + policy.setActionMeta(originActions); + // Env.getCurrentEnv().getEditLog().logCreateWorkloadSchedPolicy(policy); + idToPolicy.put(id, policy); + nameToPolicy.put(policyName, policy); + } finally { + writeUnlock(); + } + } + + private void checkPolicyActionConflicts(List<WorkloadAction> actionList) throws UserException { + Set<WorkloadActionType> actionTypeSet = new HashSet<>(); + for (WorkloadAction action : actionList) { + if (!actionTypeSet.add(action.getWorkloadActionType())) { + throw new UserException("duplicate action in one policy"); + } + } + if (actionTypeSet.contains(WorkloadActionType.cancel_query) && actionTypeSet.contains( + WorkloadActionType.move_query_to_group)) { + throw new UserException(String.format("%s and %s can not exist in one policy at same time", + WorkloadActionType.cancel_query, WorkloadActionType.move_query_to_group)); + } + } + + public void execPolicy(List<WorkloadQueryInfo> queryInfoList) { + // 1 get a snapshot of policy + Set<Long> policyIdSet = new HashSet<>(); + readLock(); + try { + policyIdSet.addAll(idToPolicy.keySet()); + } finally { + readUnlock(); + } + + for (WorkloadQueryInfo queryInfo : queryInfoList) { + try { + // 1 check policy is match + Map<WorkloadActionType, Queue<WorkloadSchedPolicy>> matchedPolicyMap = Maps.newHashMap(); + for (Long policyId : policyIdSet) { + WorkloadSchedPolicy policy = idToPolicy.get(policyId); + if (policy == null) { + continue; + } + if (policy.isEnabled() && policy.isMatch(queryInfo)) { + WorkloadActionType actionType = policy.getFirstActionType(); + // add to priority queue + Queue<WorkloadSchedPolicy> queue = matchedPolicyMap.get(actionType); + if (queue == null) { + queue = new PriorityQueue<>(policyComparator); + matchedPolicyMap.put(actionType, queue); + } + queue.offer(policy); Review Comment: 这会个是根据priority对policy排序的话，刚测过了，设置并发的两个规则，优先级更高的会生效 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For additional commands, e-mail: commits-h...@doris.apache.org
