This is an automated email from the ASF dual-hosted git repository.
wangbo pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new f7437380ca3 [Fix] some fix for workload group (#50543)
f7437380ca3 is described below
commit f7437380ca39ff02e42a57a368b4f0a09fcf3ed1
Author: wangbo <wan...@selectdb.com>
AuthorDate: Tue May 6 14:20:49 2025 +0800
[Fix] some fix for workload group (#50543)
---
.../org/apache/doris/mysql/privilege/Auth.java | 10 ++
.../resource/workloadgroup/WorkloadGroup.java | 4 +
.../resource/workloadgroup/WorkloadGroupKey.java | 78 +++++++++++++
.../resource/workloadgroup/WorkloadGroupMgr.java | 126 +++++++++++----------
.../org/apache/doris/mysql/privilege/AuthTest.java | 3 +
.../workloadgroup/WorkloadGroupMgrTest.java | 64 ++++++++---
.../resource/workloadgroup/WorkloadGroupTest.java | 43 +++++++
.../account_p0/test_grant_priv_workload.groovy | 13 +++
.../workload_manager_p0/test_curd_wlg.groovy | 14 +--
9 files changed, 276 insertions(+), 79 deletions(-)
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java
b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java
index aba0f9571f9..cb252bd983b 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java
@@ -774,6 +774,16 @@ public class Auth implements Writable {
private void grantInternal(UserIdentity userIdent, String role,
WorkloadGroupPattern workloadGroupPattern,
PrivBitSet privs, boolean errOnNonExist, boolean isReplay) throws
DdlException {
+ if (!isReplay) {
+ if (!FeConstants.runningUnitTest) {
+ if (!"%".equals(workloadGroupPattern.getworkloadGroupName())
&& !Env.getCurrentEnv()
+ .getWorkloadGroupMgr()
+
.isWorkloadGroupExists(workloadGroupPattern.getworkloadGroupName())) {
+ throw new DdlException(
+ "Can not find workload group " +
workloadGroupPattern.getworkloadGroupName());
+ }
+ }
+ }
writeLock();
try {
if (role == null) {
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroup.java
b/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroup.java
index f7a0fdce49c..5cf893fc004 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroup.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroup.java
@@ -539,6 +539,10 @@ public class WorkloadGroup implements Writable,
GsonPostProcessable {
return name;
}
+ public WorkloadGroupKey getWorkloadGroupKey() {
+ return WorkloadGroupKey.get(this.getComputeGroup(), this.getName());
+ }
+
public Map<String, String> getProperties() {
return properties;
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupKey.java
b/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupKey.java
new file mode 100644
index 00000000000..f894003fb2f
--- /dev/null
+++
b/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupKey.java
@@ -0,0 +1,78 @@
+// 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.workloadgroup;
+
+import com.google.common.base.Preconditions;
+import org.apache.commons.lang3.StringUtils;
+
+import java.util.Objects;
+
+public class WorkloadGroupKey {
+
+ private String computeGroup;
+
+ private String workloadGroupName;
+
+ private WorkloadGroupKey(String computeGroup, String workloadGroupName) {
+ this.computeGroup = computeGroup;
+ this.workloadGroupName = workloadGroupName;
+ }
+
+ public static WorkloadGroupKey get(String computeGroup, String
workloadGroupName) {
+ Preconditions.checkState(!StringUtils.isEmpty(workloadGroupName));
+ return new WorkloadGroupKey(
+ StringUtils.isEmpty(computeGroup) ?
WorkloadGroupMgr.EMPTY_COMPUTE_GROUP : computeGroup,
+ workloadGroupName);
+ }
+
+ public String getComputeGroup() {
+ return computeGroup;
+ }
+
+ public String getWorkloadGroupName() {
+ return workloadGroupName;
+ }
+
+ @Override
+ public String toString() {
+ return computeGroup + "." + workloadGroupName;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+
+ if (o instanceof WorkloadGroupKey) {
+ WorkloadGroupKey otherWgKey = (WorkloadGroupKey) o;
+
+ boolean cgEqual =
this.computeGroup.equals(otherWgKey.computeGroup);
+ boolean wgEqual =
this.workloadGroupName.equals(otherWgKey.workloadGroupName);
+ return cgEqual && wgEqual;
+ }
+
+ return false;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(computeGroup, workloadGroupName);
+ }
+
+}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgr.java
b/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgr.java
index dbe91a2898a..662a3714eb7 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgr.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgr.java
@@ -88,7 +88,7 @@ public class WorkloadGroupMgr extends MasterDaemon implements
Writable, GsonPost
private static final Logger LOG =
LogManager.getLogger(WorkloadGroupMgr.class);
@SerializedName(value = "idToWorkloadGroup")
private final Map<Long, WorkloadGroup> idToWorkloadGroup =
Maps.newHashMap();
- private final Map<Pair<String, String>, WorkloadGroup> nameToWorkloadGroup
= Maps.newHashMap();
+ private final Map<WorkloadGroupKey, WorkloadGroup> keyToWorkloadGroup =
Maps.newHashMap();
private final Map<Long, QueryQueue> idToQueryQueue = Maps.newHashMap();
private final ResourceProcNode procNode = new ResourceProcNode();
private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
@@ -163,12 +163,13 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
lock.writeLock().unlock();
}
- private WorkloadGroup getWorkloadGroupByComputeGroupUnlock(String cgName,
String wgName)
+ private WorkloadGroup
getWorkloadGroupByComputeGroupUnlock(WorkloadGroupKey wgKey)
throws DdlException {
- WorkloadGroup wg = nameToWorkloadGroup.get(Pair.of(cgName, wgName));
+ WorkloadGroup wg = keyToWorkloadGroup.get(wgKey);
if (wg == null) {
throw new DdlException(
- "Can not find workload group " + wgName + " in compute
group " + cgName + ".");
+ "Can not find workload group " +
wgKey.getWorkloadGroupName() + " in compute group "
+ + wgKey.getComputeGroup() + ".");
}
return wg;
}
@@ -181,7 +182,8 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
readLock();
try {
for (String cgName : cgNames) {
- WorkloadGroup workloadGroup =
getWorkloadGroupByComputeGroupUnlock(cgName, wgName);
+ WorkloadGroup workloadGroup =
getWorkloadGroupByComputeGroupUnlock(
+ WorkloadGroupKey.get(cgName, wgName));
workloadGroups.add(workloadGroup.toThrift());
}
context.setWorkloadGroupName(wgName);
@@ -271,7 +273,7 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
valueList.add(String.valueOf(qdtail.first));
valueList.add(String.valueOf(qdtail.second));
}
- ret.put(wg.getName(), valueList);
+ ret.put(wg.getWorkloadGroupKey().toString(), valueList);
}
} finally {
readUnlock();
@@ -300,20 +302,21 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
String workloadGroupName = workloadGroup.getName();
writeLock();
try {
- Pair<String, String> key = Pair.of(computeGroup,
workloadGroupName);
- if (nameToWorkloadGroup.containsKey(key)) {
+ WorkloadGroupKey wgKey = WorkloadGroupKey.get(computeGroup,
workloadGroupName);
+ if (keyToWorkloadGroup.containsKey(wgKey)) {
if (isIfNotExists) {
return;
}
throw new DdlException(
- "Compute group " + key.first + " already has workload
group " + key.second + ".");
+ "Compute group " + wgKey.getComputeGroup() + " already
has workload group "
+ + wgKey.getWorkloadGroupName() + ".");
}
if (idToWorkloadGroup.size() >= Config.workload_group_max_num) {
throw new DdlException(
"Workload group number can not be exceed " +
Config.workload_group_max_num);
}
checkGlobalUnlock(workloadGroup, null);
- nameToWorkloadGroup.put(key, workloadGroup);
+ keyToWorkloadGroup.put(wgKey, workloadGroup);
idToWorkloadGroup.put(workloadGroup.getId(), workloadGroup);
Env.getCurrentEnv().getEditLog().logCreateWorkloadGroup(workloadGroup);
} finally {
@@ -386,12 +389,13 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
}
WorkloadGroup newWorkloadGroup;
+ WorkloadGroupKey wgKey = WorkloadGroupKey.get(computeGroup,
workloadGroupName);
writeLock();
try {
- WorkloadGroup currentWorkloadGroup =
getWorkloadGroupByComputeGroupUnlock(computeGroup, workloadGroupName);
+ WorkloadGroup currentWorkloadGroup =
getWorkloadGroupByComputeGroupUnlock(wgKey);
newWorkloadGroup =
WorkloadGroup.copyAndUpdate(currentWorkloadGroup, properties);
checkGlobalUnlock(newWorkloadGroup, currentWorkloadGroup);
- nameToWorkloadGroup.put(Pair.of(computeGroup, workloadGroupName),
newWorkloadGroup);
+ keyToWorkloadGroup.put(wgKey, newWorkloadGroup);
idToWorkloadGroup.put(newWorkloadGroup.getId(), newWorkloadGroup);
Env.getCurrentEnv().getEditLog().logAlterWorkloadGroup(newWorkloadGroup);
} finally {
@@ -421,11 +425,13 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
+ "then you can drop the group.");
}
+ WorkloadGroupKey wgKey = WorkloadGroupKey.get(computeGroup,
workloadGroupName);
+
// A group with related policies should not be deleted.
Long wgId = null;
readLock();
try {
- WorkloadGroup wg = nameToWorkloadGroup.get(Pair.of(computeGroup,
workloadGroupName));
+ WorkloadGroup wg = keyToWorkloadGroup.get(wgKey);
if (wg != null) {
wgId = wg.getId();
}
@@ -441,52 +447,43 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
}
}
- Pair<String, String> key = Pair.of(computeGroup, workloadGroupName);
writeLock();
try {
- if (!nameToWorkloadGroup.containsKey(key)) {
+ if (!keyToWorkloadGroup.containsKey(wgKey)) {
if (ifExists) {
return;
}
throw new DdlException(
- "Can not find workload group " + workloadGroupName + "
in compute group " + computeGroup + ".");
+ "Can not find workload group " +
wgKey.getWorkloadGroupName() + " in compute group "
+ + wgKey.getComputeGroup() + ".");
}
- WorkloadGroup workloadGroup = nameToWorkloadGroup.get(key);
- long groupId = workloadGroup.getId();
- idToWorkloadGroup.remove(groupId);
- nameToWorkloadGroup.remove(key);
- idToQueryQueue.remove(groupId);
- Env.getCurrentEnv().getEditLog().logDropWorkloadGroup(new
DropWorkloadGroupOperationLog(groupId));
+ WorkloadGroup workloadGroup = keyToWorkloadGroup.get(wgKey);
+ keyToWorkloadGroup.remove(wgKey);
+ idToWorkloadGroup.remove(workloadGroup.getId());
+ idToQueryQueue.remove(workloadGroup.getId());
+ Env.getCurrentEnv().getEditLog()
+ .logDropWorkloadGroup(new
DropWorkloadGroupOperationLog(workloadGroup.getId()));
} finally {
writeUnlock();
}
- LOG.info("Drop workload group success: {} for compute group {}",
workloadGroupName, computeGroup);
+ LOG.info("Drop workload group success: {} for compute group {}",
wgKey.getWorkloadGroupName(),
+ wgKey.getComputeGroup());
}
private void insertWorkloadGroup(WorkloadGroup workloadGroup) {
writeLock();
try {
- LOG.info("[init_wg] before, name map, {}, id map: {}, name map:
{}", "replay",
- idToWorkloadGroup, nameToWorkloadGroup);
+ LOG.info("[init_wg] before, {}, id map: {}, name map: {}",
"replay",
+ idToWorkloadGroup, keyToWorkloadGroup);
idToWorkloadGroup.put(workloadGroup.getId(), workloadGroup);
- Pair<String, String> key =
Pair.of(workloadGroup.getComputeGroup(), workloadGroup.getName());
- nameToWorkloadGroup.put(key, workloadGroup);
- LOG.info("[init_wg] after, name map, {}, id map: {}, name map:
{}", "replay",
- idToWorkloadGroup, nameToWorkloadGroup);
+ keyToWorkloadGroup.put(workloadGroup.getWorkloadGroupKey(),
workloadGroup);
+ LOG.info("[init_wg] after, {}, id map: {}, name map: {}", "replay",
+ idToWorkloadGroup, keyToWorkloadGroup);
} finally {
writeUnlock();
}
}
- public boolean isWorkloadGroupExists(String workloadGroupName) {
- readLock();
- try {
- return nameToWorkloadGroup.containsKey(workloadGroupName);
- } finally {
- readUnlock();
- }
- }
-
public void replayCreateWorkloadGroup(WorkloadGroup workloadGroup) {
insertWorkloadGroup(workloadGroup);
}
@@ -503,7 +500,7 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
return;
}
WorkloadGroup workloadGroup = idToWorkloadGroup.get(id);
- nameToWorkloadGroup.remove(workloadGroup.getName());
+ keyToWorkloadGroup.remove(workloadGroup.getWorkloadGroupKey());
idToWorkloadGroup.remove(id);
} finally {
writeUnlock();
@@ -554,8 +551,8 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
}
// for ut
- public Map<Pair<String, String>, WorkloadGroup> getNameToWorkloadGroup() {
- return nameToWorkloadGroup;
+ public Map<WorkloadGroupKey, WorkloadGroup> getNameToWorkloadGroup() {
+ return keyToWorkloadGroup;
}
// for ut
@@ -576,14 +573,13 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
@Override
public void gsonPostProcess() throws IOException {
- LOG.info("[init_wg] before, name map, {}, id map: {}, name map: {}",
"gson",
- idToWorkloadGroup, nameToWorkloadGroup);
+ LOG.info("[init_wg] before, {}, id map: {}, name map: {}", "gson",
+ idToWorkloadGroup, keyToWorkloadGroup);
for (Map.Entry<Long, WorkloadGroup> entry :
idToWorkloadGroup.entrySet()) {
- String computeGroupName = entry.getValue().getComputeGroup();
- nameToWorkloadGroup.put(Pair.of(computeGroupName,
entry.getValue().getName()), entry.getValue());
+ keyToWorkloadGroup.put(entry.getValue().getWorkloadGroupKey(),
entry.getValue());
}
- LOG.info("[init_wg] after, name map, {}, id map: {}, name map: {}",
"gson",
- idToWorkloadGroup, nameToWorkloadGroup);
+ LOG.info("[init_wg] after, {}, id map: {}, name map: {}", "gson",
+ idToWorkloadGroup, keyToWorkloadGroup);
}
public class ResourceProcNode {
@@ -611,8 +607,8 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
List<WorkloadGroup> oldWgList = Lists.newArrayList();
readLock();
try {
- for (Map.Entry<Pair<String, String>, WorkloadGroup> entry :
nameToWorkloadGroup.entrySet()) {
- if (EMPTY_COMPUTE_GROUP.equals(entry.getKey().first)) {
+ for (Map.Entry<WorkloadGroupKey, WorkloadGroup> entry :
keyToWorkloadGroup.entrySet()) {
+ if
(EMPTY_COMPUTE_GROUP.equals(entry.getKey().getComputeGroup())) {
oldWgList.add(entry.getValue());
}
}
@@ -626,7 +622,7 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
writeLock();
try {
LOG.info("[init_wg] before create normal wg, id map: {}, name map:
{}", idToWorkloadGroup,
- nameToWorkloadGroup);
+ keyToWorkloadGroup);
if (idToWorkloadGroup.isEmpty()) {
String defaultCgName = Config.isCloudMode() ?
Tag.VALUE_DEFAULT_COMPUTE_GROUP_NAME
: Tag.VALUE_DEFAULT_TAG;
@@ -636,14 +632,14 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
WorkloadGroup defaultWg = new
WorkloadGroup(Env.getCurrentEnv().getNextId(), DEFAULT_GROUP_NAME,
properties);
idToWorkloadGroup.put(defaultWg.getId(), defaultWg);
- nameToWorkloadGroup.put(Pair.of(defaultWg.getComputeGroup(),
defaultWg.getName()), defaultWg);
+ keyToWorkloadGroup.put(defaultWg.getWorkloadGroupKey(),
defaultWg);
Env.getCurrentEnv().getEditLog().logCreateWorkloadGroup(defaultWg);
LOG.info("[init_wg]Create default workload group success: {}",
defaultWg);
} else {
LOG.info("[init_wg]This is not a new cluster, skip create
default wg");
}
LOG.info("[init_wg] after create normal wg, id map: {}, name map:
{}", idToWorkloadGroup,
- nameToWorkloadGroup);
+ keyToWorkloadGroup);
} finally {
writeUnlock();
}
@@ -652,16 +648,16 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
public void bindWorkloadGroupToComputeGroup(Set<String> cgSet,
WorkloadGroup oldWg) {
writeLock();
try {
- Pair<String, String> oldKey = Pair.of(EMPTY_COMPUTE_GROUP,
oldWg.getName());
+ WorkloadGroupKey oldKey =
WorkloadGroupKey.get(EMPTY_COMPUTE_GROUP, oldWg.getName());
// it means old compute group has been dropped, just return;
- if (!nameToWorkloadGroup.containsKey(oldKey)) {
+ if (!keyToWorkloadGroup.containsKey(oldKey)) {
LOG.info("[init_wg]Old workload group {} has been dropped,
skip it.", oldWg.getName());
return;
}
// create new workload group for all compute group.
for (String computeGroup : cgSet) {
- Pair<String, String> newKey = Pair.of(computeGroup,
oldWg.getName());
- if (nameToWorkloadGroup.containsKey(newKey)) {
+ WorkloadGroupKey newKey = WorkloadGroupKey.get(computeGroup,
oldWg.getName());
+ if (keyToWorkloadGroup.containsKey(newKey)) {
LOG.info("[init_wg]Workload group {} already exists in
compute group {}, skip it.",
oldWg.getName(), computeGroup);
continue;
@@ -673,7 +669,7 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
newProp.put(WorkloadGroup.COMPUTE_GROUP, computeGroup);
WorkloadGroup newWg = new
WorkloadGroup(Env.getCurrentEnv().getNextId(), oldWg.getName(),
newProp);
- nameToWorkloadGroup.put(newKey, newWg);
+ keyToWorkloadGroup.put(newKey, newWg);
idToWorkloadGroup.put(newWg.getId(), newWg);
Env.getCurrentEnv().getEditLog().logCreateWorkloadGroup(newWg);
LOG.info("[init_wg]Create workload group {} for compute group
{} success.", oldWg.getName(),
@@ -681,7 +677,7 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
}
// drop old workload group
- nameToWorkloadGroup.remove(oldKey);
+ keyToWorkloadGroup.remove(oldKey);
idToWorkloadGroup.remove(oldWg.getId());
idToQueryQueue.remove(oldWg.getId());
Env.getCurrentEnv().getEditLog().logDropWorkloadGroup(new
DropWorkloadGroupOperationLog(oldWg.getId()));
@@ -693,4 +689,18 @@ public class WorkloadGroupMgr extends MasterDaemon
implements Writable, GsonPost
}
}
+ public boolean isWorkloadGroupExists(String wgName) {
+ readLock();
+ try {
+ for (WorkloadGroup wg : idToWorkloadGroup.values()) {
+ if (wg.getName().equals(wgName)) {
+ return true;
+ }
+ }
+ return false;
+ } finally {
+ readUnlock();
+ }
+ }
+
}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/mysql/privilege/AuthTest.java
b/fe/fe-core/src/test/java/org/apache/doris/mysql/privilege/AuthTest.java
index e7f81c31a64..333904bf98d 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/mysql/privilege/AuthTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/mysql/privilege/AuthTest.java
@@ -38,6 +38,7 @@ import org.apache.doris.catalog.Env;
import org.apache.doris.common.AnalysisException;
import org.apache.doris.common.DdlException;
import org.apache.doris.common.ExceptionChecker;
+import org.apache.doris.common.FeConstants;
import org.apache.doris.common.UserException;
import org.apache.doris.datasource.InternalCatalog;
import org.apache.doris.mysql.MysqlPassword;
@@ -1987,6 +1988,7 @@ public class AuthTest {
@Test
public void testWorkloadGroupPriv() throws UserException {
+ FeConstants.runningUnitTest = true;
UserIdentity userIdentity = new UserIdentity("testUser", "%");
String role = "role0";
String workloadGroupName = "g1";
@@ -2447,6 +2449,7 @@ public class AuthTest {
@Test
public void testShowRoles() {
+ FeConstants.runningUnitTest = true;
String role = "test_wg_role";
CreateRoleStmt roleStmt = new CreateRoleStmt(role);
try {
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgrTest.java
b/fe/fe-core/src/test/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgrTest.java
index 323f95f1f15..8489466bcb2 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgrTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/resource/workloadgroup/WorkloadGroupMgrTest.java
@@ -25,6 +25,7 @@ import org.apache.doris.common.UserException;
import org.apache.doris.mysql.privilege.AccessControllerManager;
import org.apache.doris.mysql.privilege.Auth;
import org.apache.doris.mysql.privilege.PrivPredicate;
+import org.apache.doris.persist.DropWorkloadGroupOperationLog;
import org.apache.doris.persist.EditLog;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.resource.computegroup.ComputeGroup;
@@ -127,13 +128,13 @@ public class WorkloadGroupMgrTest {
WorkloadGroup wg1 = new WorkloadGroup(wgId1, wgName1, properties1);
workloadGroupMgr.createWorkloadGroup(cg1, wg1, false);
- Pair<String, String> key1 = Pair.of(cg1, wgName1);
- Map<Pair<String, String>, WorkloadGroup> nameToRG =
workloadGroupMgr.getNameToWorkloadGroup();
+ WorkloadGroupKey key1 = WorkloadGroupKey.get(cg1, wgName1);
+ Map<WorkloadGroupKey, WorkloadGroup> nameToRG =
workloadGroupMgr.getNameToWorkloadGroup();
Assert.assertEquals(1, nameToRG.size());
Assert.assertTrue(nameToRG.containsKey(key1));
WorkloadGroup group1 = nameToRG.get(key1);
- Assert.assertEquals(key1.second, group1.getName());
- Assert.assertEquals(key1.first, group1.getComputeGroup());
+ Assert.assertEquals(key1.getWorkloadGroupName(), group1.getName());
+ Assert.assertEquals(key1.getComputeGroup(), group1.getComputeGroup());
Map<Long, WorkloadGroup> idToRG =
workloadGroupMgr.getIdToWorkloadGroup();
Assert.assertEquals(1, idToRG.size());
@@ -150,7 +151,7 @@ public class WorkloadGroupMgrTest {
WorkloadGroup wg2 = new WorkloadGroup(wgId2, wgName2, properties2);
workloadGroupMgr.createWorkloadGroup(cg2, wg2, false);
- Pair<String, String> key2 = Pair.of(cg2, wgName2);
+ WorkloadGroupKey key2 = WorkloadGroupKey.get(cg2, wgName2);
nameToRG = workloadGroupMgr.getNameToWorkloadGroup();
Assert.assertEquals(2, nameToRG.size());
Assert.assertTrue(nameToRG.containsKey(key2));
@@ -158,7 +159,7 @@ public class WorkloadGroupMgrTest {
idToRG = workloadGroupMgr.getIdToWorkloadGroup();
Assert.assertEquals(2, idToRG.size());
Assert.assertTrue(idToRG.containsKey(group2.getId()));
- Assert.assertTrue(key2.first.equals(wg2.getComputeGroup()));
+
Assert.assertTrue(key2.getComputeGroup().equals(wg2.getComputeGroup()));
// 3 test memory limit exceeds
Map<String, String> properties3 = Maps.newHashMap();
@@ -294,7 +295,7 @@ public class WorkloadGroupMgrTest {
prop1.put(WorkloadGroup.CPU_SHARE, "10");
workloadGroupMgr.createWorkloadGroup(cgName1, new WorkloadGroup(wgId1,
wgName1, prop1), false);
Assert.assertTrue(Long.valueOf(
- workloadGroupMgr.getNameToWorkloadGroup().get(Pair.of(cgName1,
wgName1)).getProperties()
+
workloadGroupMgr.getNameToWorkloadGroup().get(WorkloadGroupKey.get(cgName1,
wgName1)).getProperties()
.get(WorkloadGroup.CPU_SHARE)) == 10);
// test alter failed
@@ -311,7 +312,7 @@ public class WorkloadGroupMgrTest {
// test alter success
workloadGroupMgr.alterWorkloadGroup(cgName1, wgName1, prop2);
- WorkloadGroup wg =
workloadGroupMgr.getNameToWorkloadGroup().get(Pair.of(cgName1, wgName1));
+ WorkloadGroup wg =
workloadGroupMgr.getNameToWorkloadGroup().get(WorkloadGroupKey.get(cgName1,
wgName1));
Assert.assertTrue(Long.valueOf(wg.getProperties().get(WorkloadGroup.CPU_SHARE))
== 20);
}
@@ -333,7 +334,7 @@ public class WorkloadGroupMgrTest {
prop1.put(WorkloadGroup.CPU_SHARE, "123");
WorkloadGroup wg1 = new WorkloadGroup(wgId1, wgName1, prop1);
wgMgr.getIdToWorkloadGroup().put(wgId1, wg1);
-
wgMgr.getNameToWorkloadGroup().put(Pair.of(WorkloadGroupMgr.EMPTY_COMPUTE_GROUP,
wgName1), wg1);
+
wgMgr.getNameToWorkloadGroup().put(WorkloadGroupKey.get(WorkloadGroupMgr.EMPTY_COMPUTE_GROUP,
wgName1), wg1);
wgMgr.getIdToQueryQueue().put(wgId1, new QueryQueue(0, 0, 0, 0, 0));
long wgId2 = 2;
@@ -342,7 +343,7 @@ public class WorkloadGroupMgrTest {
prop2.put(WorkloadGroup.CPU_SHARE, "123");
WorkloadGroup wg2 = new WorkloadGroup(wgId2, wgName2, prop2);
wgMgr.getIdToWorkloadGroup().put(wgId2, wg2);
-
wgMgr.getNameToWorkloadGroup().put(Pair.of(WorkloadGroupMgr.EMPTY_COMPUTE_GROUP,
wgName2), wg2);
+
wgMgr.getNameToWorkloadGroup().put(WorkloadGroupKey.get(WorkloadGroupMgr.EMPTY_COMPUTE_GROUP,
wgName2), wg2);
wgMgr.getIdToQueryQueue().put(wgId2, new QueryQueue(0, 0, 0, 0, 0));
long wgId3 = 3;
@@ -351,7 +352,7 @@ public class WorkloadGroupMgrTest {
prop3.put(WorkloadGroup.CPU_SHARE, "123");
WorkloadGroup wg3 = new WorkloadGroup(wgId3, wgName3, prop3);
wgMgr.getIdToWorkloadGroup().put(wgId3, wg3);
-
wgMgr.getNameToWorkloadGroup().put(Pair.of(WorkloadGroupMgr.EMPTY_COMPUTE_GROUP,
wgName3), wg3);
+
wgMgr.getNameToWorkloadGroup().put(WorkloadGroupKey.get(WorkloadGroupMgr.EMPTY_COMPUTE_GROUP,
wgName3), wg3);
wgMgr.getIdToQueryQueue().put(wgId3, new QueryQueue(0, 0, 0, 0, 0));
@@ -392,9 +393,9 @@ public class WorkloadGroupMgrTest {
Assert.assertTrue(wgMgr.getIdToWorkloadGroup().get(wgId4).equals(wg4));
for (String cgName : cgSet) {
- WorkloadGroup wg11 =
wgMgr.getNameToWorkloadGroup().get(Pair.of(cgName, wgName1));
- WorkloadGroup wg22 =
wgMgr.getNameToWorkloadGroup().get(Pair.of(cgName, wgName2));
- WorkloadGroup wg33 =
wgMgr.getNameToWorkloadGroup().get(Pair.of(cgName, wgName3));
+ WorkloadGroup wg11 =
wgMgr.getNameToWorkloadGroup().get(WorkloadGroupKey.get(cgName, wgName1));
+ WorkloadGroup wg22 =
wgMgr.getNameToWorkloadGroup().get(WorkloadGroupKey.get(cgName, wgName2));
+ WorkloadGroup wg33 =
wgMgr.getNameToWorkloadGroup().get(WorkloadGroupKey.get(cgName, wgName3));
Assert.assertTrue(wgMgr.getIdToWorkloadGroup().containsKey(wg11.getId()));
Assert.assertTrue(wgMgr.getIdToWorkloadGroup().containsKey(wg22.getId()));
@@ -560,4 +561,39 @@ public class WorkloadGroupMgrTest {
}
}
}
+
+ @Test
+ public void testReplayWorkloadGroup() {
+ WorkloadGroupMgr wgMgr = new WorkloadGroupMgr();
+ Assert.assertTrue(wgMgr.getNameToWorkloadGroup().size() == 0);
+ Assert.assertTrue(wgMgr.getIdToWorkloadGroup().size() == 0);
+
+
+ // 1 test replay create
+ WorkloadGroup wg1 = new WorkloadGroup(1, "wg1", Maps.newHashMap());
+ wgMgr.replayCreateWorkloadGroup(wg1);
+
+ Assert.assertTrue(wgMgr.getNameToWorkloadGroup().size() == 1);
+ Assert.assertTrue(wgMgr.getIdToWorkloadGroup().size() == 1);
+
Assert.assertTrue(wgMgr.getNameToWorkloadGroup().get(wg1.getWorkloadGroupKey())
+ .equals(wgMgr.getIdToWorkloadGroup().get(wg1.getId())));
+
+ // 2 test replay alter
+ Map<String, String> pop2 = Maps.newHashMap();
+ pop2.put("cpu_share", "2345");
+ WorkloadGroup wg2 = new WorkloadGroup(1, "wg1", pop2);
+ wgMgr.replayAlterWorkloadGroup(wg2);
+
Assert.assertTrue(wgMgr.getNameToWorkloadGroup().get(wg2.getWorkloadGroupKey())
+ .equals(wgMgr.getIdToWorkloadGroup().get(wg2.getId())));
+
Assert.assertTrue(wgMgr.getNameToWorkloadGroup().get(wg2.getWorkloadGroupKey()).getProperties().get("cpu_share")
+ .equals("2345"));
+ Assert.assertTrue(wgMgr.getNameToWorkloadGroup().size() == 1);
+ Assert.assertTrue(wgMgr.getIdToWorkloadGroup().size() == 1);
+
+ // 3 test replay drop
+ DropWorkloadGroupOperationLog dropLog = new
DropWorkloadGroupOperationLog(1);
+ wgMgr.replayDropWorkloadGroup(dropLog);
+ Assert.assertTrue(wgMgr.getNameToWorkloadGroup().size() == 0);
+ Assert.assertTrue(wgMgr.getIdToWorkloadGroup().size() == 0);
+ }
}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/resource/workloadgroup/WorkloadGroupTest.java
b/fe/fe-core/src/test/java/org/apache/doris/resource/workloadgroup/WorkloadGroupTest.java
index a7176f2d50a..e75f1ef904a 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/resource/workloadgroup/WorkloadGroupTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/resource/workloadgroup/WorkloadGroupTest.java
@@ -109,4 +109,47 @@ public class WorkloadGroupTest {
}
Assert.assertEquals(hasException, true);
}
+
+ @Test
+ public void testWorkloadGroupKey() {
+ // equal
+ WorkloadGroupKey eqKey1 = WorkloadGroupKey.get("cg1", "wg1");
+ WorkloadGroupKey eqKey2 = WorkloadGroupKey.get("cg1", "wg1");
+ WorkloadGroupKey eqKey3 = WorkloadGroupKey.get("cg1", "wg2");
+ Assert.assertTrue(eqKey1.equals(eqKey1));
+ Assert.assertTrue(eqKey1.equals(eqKey2));
+ Assert.assertTrue(eqKey2.equals(eqKey1));
+ Assert.assertTrue(eqKey1.hashCode() == eqKey2.hashCode());
+
+ Assert.assertFalse(eqKey3.equals(eqKey1));
+ Assert.assertFalse(eqKey1.equals(eqKey3));
+ Assert.assertTrue(eqKey1.hashCode() != eqKey3.hashCode());
+
+ WorkloadGroupKey eqKey4 = WorkloadGroupKey.get("cg2", "wg2");
+ Assert.assertFalse(eqKey4.equals(eqKey3));
+ Assert.assertFalse(eqKey3.equals(eqKey4));
+ Assert.assertFalse(eqKey4.hashCode() == eqKey3.hashCode());
+
+
+ // test wg name exception
+ try {
+ WorkloadGroupKey.get("cg1", "");
+ Assert.fail();
+ } catch (IllegalStateException e) {
+ Assert.assertTrue(true);
+ }
+
+
+ // test null equal
+ Assert.assertTrue(!eqKey1.equals(null));
+
+ WorkloadGroupKey nullkey2 = WorkloadGroupKey.get(null, "wg1");
+ WorkloadGroupKey nullkey3 = WorkloadGroupKey.get("", "wg1");
+ Assert.assertTrue(nullkey2.equals(nullkey3));
+ Assert.assertTrue(nullkey3.equals(nullkey2));
+
+ Assert.assertFalse(nullkey2.equals(eqKey1));
+ Assert.assertFalse(eqKey1.equals(nullkey2));
+
+ }
}
diff --git a/regression-test/suites/account_p0/test_grant_priv_workload.groovy
b/regression-test/suites/account_p0/test_grant_priv_workload.groovy
index 12149f0fa7a..c5ce48318c2 100644
--- a/regression-test/suites/account_p0/test_grant_priv_workload.groovy
+++ b/regression-test/suites/account_p0/test_grant_priv_workload.groovy
@@ -18,6 +18,16 @@
import org.junit.Assert;
suite("test_grant_priv_workload") {
+ def forComputeGroupStr = "";
+
+ //cloud-mode
+ if (isCloudMode()) {
+ def clusters = sql " SHOW CLUSTERS; "
+ assertTrue(!clusters.isEmpty())
+ def validCluster = clusters[0][0]
+ forComputeGroupStr = " for $validCluster "
+ }
+
def user1 = 'test_grant_priv_workload_user1'
def user2 = 'test_grant_priv_workload_user2'
def pwd = '123456'
@@ -31,6 +41,8 @@ suite("test_grant_priv_workload") {
sql """CREATE USER '${user1}' IDENTIFIED BY '${pwd}'"""
sql """CREATE USER '${user2}' IDENTIFIED BY '${pwd}'"""
+ sql "create workload group if not exists ${workload1}
${forComputeGroupStr} properties('cpu_share'='123')"
+
// test only have USAGE_PRIV, can not grant to other user
sql """grant USAGE_PRIV on WORKLOAD GROUP ${workload1} to ${user1}"""
connect(user1, "${pwd}", url) {
@@ -52,6 +64,7 @@ suite("test_grant_priv_workload") {
}
}
+ sql "drop workload group ${workload1} ${forComputeGroupStr}"
sql """drop user if exists ${user1}"""
sql """drop user if exists ${user2}"""
}
diff --git a/regression-test/suites/workload_manager_p0/test_curd_wlg.groovy
b/regression-test/suites/workload_manager_p0/test_curd_wlg.groovy
index 53d800bb4e8..e3a443afe3d 100644
--- a/regression-test/suites/workload_manager_p0/test_curd_wlg.groovy
+++ b/regression-test/suites/workload_manager_p0/test_curd_wlg.groovy
@@ -88,13 +88,6 @@ suite("test_crud_wlg") {
sql "ADMIN SET FRONTEND CONFIG ('enable_workload_group' = 'true');"
sql "ADMIN SET FRONTEND CONFIG ('query_queue_update_interval_ms' = '100');"
- sql "create workload group if not exists normal $forComputeGroupStr " +
- "properties ( " +
- " 'cpu_share'='1024', " +
- " 'memory_limit'='50%', " +
- " 'enable_memory_overcommit'='true' " +
- ");"
-
// reset normal group property
sql "alter workload group normal $forComputeGroupStr properties (
'cpu_share'='1024' );"
sql "alter workload group normal $forComputeGroupStr properties (
'memory_limit'='50%' );"
@@ -340,6 +333,13 @@ suite("test_crud_wlg") {
}
}
+
+ sql "drop workload group if exists grant_test_wg $forComputeGroupStr;"
+ test {
+ sql " GRANT USAGE_PRIV ON WORKLOAD GROUP grant_test_wg TO
'test_wlg_user'@'%';"
+ exception "Can not find workload group"
+ }
+
sql "GRANT USAGE_PRIV ON WORKLOAD GROUP 'test_group' TO
'test_wlg_user'@'%';"
connect('test_wlg_user', '12345', context.config.jdbcUrl) {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org
For additional commands, e-mail: commits-h...@doris.apache.org
