From 3bba33b1eb60d8092aac7d7cb9ae0461f7dfe36a Mon Sep 17 00:00:00 2001 From: 761417898 <761417898@qq.com> Date: Thu, 16 Jul 2026 12:28:30 +0800 Subject: [PATCH] Add unified user resource quota framework with SET/SHOW/DELETE USER QUOTA. Introduce per-user min/max quotas for read/write CPU, memory, temp disk and disk IO, with ConfigNode persistence, DataNode acquire/release, THROTTLE compatibility, and related tests. Usage is reported via a dedicated low-frequency reportUserResourceUsage RPC (not main heartbeat); disk_io/temp_disk use fixed long units (bytes/sec and bytes). --- .../iotdb/db/it/quotas/IoTDBSpaceQuotaIT.java | 7 + .../it/quotas/IoTDBUserResourceQuotaIT.java | 252 +++++++++ .../iotdb/db/qp/sql/IdentifierParser.g4 | 1 + .../apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 | 13 + .../org/apache/iotdb/db/qp/sql/SqlLexer.g4 | 4 + .../confignode/i18n/ManagerMessages.java | 4 + .../confignode/i18n/ManagerMessages.java | 4 + .../client/async/CnToDnAsyncRequestType.java | 1 + ...oDnInternalServiceAsyncRequestManager.java | 6 + .../consensus/request/ConfigPhysicalPlan.java | 8 + .../request/ConfigPhysicalPlanType.java | 2 + .../quota/DeleteUserResourceQuotaPlan.java | 82 +++ .../write/quota/SetUserResourceQuotaPlan.java | 158 ++++++ .../manager/ClusterQuotaManager.java | 165 ++++++ .../confignode/manager/ConfigManager.java | 32 ++ .../executor/ConfigPlanExecutor.java | 6 + .../persistence/quota/QuotaInfo.java | 219 ++++++++ .../thrift/ConfigNodeRPCServiceProcessor.java | 24 + .../request/ConfigPhysicalPlanSerDeTest.java | 32 ++ .../iotdb/db/i18n/DataNodeQueryMessages.java | 10 + .../iotdb/db/i18n/StorageEngineMessages.java | 17 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 10 + .../iotdb/db/i18n/StorageEngineMessages.java | 15 + .../org/apache/iotdb/db/conf/IoTDBConfig.java | 62 +++ .../apache/iotdb/db/conf/IoTDBDescriptor.java | 14 + .../db/protocol/client/ConfigNodeClient.java | 31 ++ .../impl/DataNodeInternalRPCServiceImpl.java | 8 + .../common/header/DatasetHeaderFactory.java | 4 + .../execution/schedule/DriverScheduler.java | 77 +-- .../execution/schedule/task/DriverTask.java | 10 + .../config/TreeConfigTaskVisitor.java | 24 + .../executor/ClusterConfigTaskExecutor.java | 86 ++++ .../config/executor/IConfigTaskExecutor.java | 13 + .../quota/DeleteUserResourceQuotaTask.java | 42 ++ .../sys/quota/SetUserResourceQuotaTask.java | 42 ++ .../sys/quota/ShowUserResourceQuotaTask.java | 247 +++++++++ .../queryengine/plan/parser/ASTVisitor.java | 145 ++++++ .../security/TreeAccessCheckVisitor.java | 21 + .../plan/statement/StatementType.java | 3 + .../plan/statement/StatementVisitor.java | 18 + .../DeleteUserResourceQuotaStatement.java | 61 +++ .../quota/SetUserResourceQuotaStatement.java | 126 +++++ .../quota/ShowUserResourceQuotaStatement.java | 79 +++ .../rescon/quotas/AcquireContext.java | 64 +++ .../rescon/quotas/AcquirePolicy.java | 49 ++ .../rescon/quotas/AcquireResult.java | 53 ++ .../quotas/CapacityResourceLimiter.java | 98 ++++ .../rescon/quotas/CpuSlotLimiter.java | 22 + .../quotas/DataNodeThrottleQuotaManager.java | 28 +- .../rescon/quotas/DiskIoLimiter.java | 72 +++ .../rescon/quotas/LimiterAcquireResult.java | 50 ++ .../rescon/quotas/MemoryBudgetLimiter.java | 22 + .../rescon/quotas/NodeQuotaState.java | 111 ++++ .../rescon/quotas/QuotaToken.java | 70 +++ .../rescon/quotas/QuotaTokenBundle.java | 50 ++ .../quotas/ResourceAwareOperationQuota.java | 55 ++ .../rescon/quotas/ResourceLimiter.java | 39 ++ .../rescon/quotas/TempDiskLimiter.java | 23 + .../rescon/quotas/ThrottleQuotaLimit.java | 9 +- .../UserResourceQuotaExceededException.java | 30 ++ .../quotas/UserResourceQuotaManager.java | 484 ++++++++++++++++++ .../rescon/quotas/WriteMemoryEstimator.java | 129 +++++ .../quota/ShowUserResourceQuotaTaskTest.java | 95 ++++ .../UserResourceQuotaParseSmokeTest.java | 168 ++++++ .../quotas/CapacityResourceLimiterTest.java | 310 +++++++++++ .../iotdb/commons/quota/OperationType.java | 25 + .../commons/quota/ResourceQuotaRange.java | 74 +++ .../iotdb/commons/quota/ResourceType.java | 27 + .../commons/quota/UserResourceQuota.java | 92 ++++ .../quota/UserResourceQuotaConverter.java | 155 ++++++ .../schema/column/ColumnHeaderConstant.java | 14 + .../quota/UserResourceQuotaConverterTest.java | 74 +++ .../src/main/thrift/common.thrift | 35 ++ .../src/main/thrift/confignode.thrift | 28 + .../src/main/thrift/datanode.thrift | 2 + 75 files changed, 4705 insertions(+), 37 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/db/it/quotas/IoTDBUserResourceQuotaIT.java create mode 100644 iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/quota/DeleteUserResourceQuotaPlan.java create mode 100644 iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/quota/SetUserResourceQuotaPlan.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/DeleteUserResourceQuotaTask.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/SetUserResourceQuotaTask.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/ShowUserResourceQuotaTask.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/DeleteUserResourceQuotaStatement.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/SetUserResourceQuotaStatement.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/ShowUserResourceQuotaStatement.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquireContext.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquirePolicy.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquireResult.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/CapacityResourceLimiter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/CpuSlotLimiter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/DiskIoLimiter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/LimiterAcquireResult.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/MemoryBudgetLimiter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/NodeQuotaState.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/QuotaToken.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/QuotaTokenBundle.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ResourceAwareOperationQuota.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ResourceLimiter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/TempDiskLimiter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/UserResourceQuotaExceededException.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/UserResourceQuotaManager.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/WriteMemoryEstimator.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/ShowUserResourceQuotaTaskTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/UserResourceQuotaParseSmokeTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/quotas/CapacityResourceLimiterTest.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/OperationType.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/ResourceQuotaRange.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/ResourceType.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/UserResourceQuota.java create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/UserResourceQuotaConverter.java create mode 100644 iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/quota/UserResourceQuotaConverterTest.java diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/quotas/IoTDBSpaceQuotaIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/quotas/IoTDBSpaceQuotaIT.java index 8930985717f5d..682ec050dad37 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/quotas/IoTDBSpaceQuotaIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/quotas/IoTDBSpaceQuotaIT.java @@ -25,6 +25,7 @@ import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.itbase.category.ClusterIT; import org.apache.iotdb.itbase.category.LocalStandaloneIT; +import org.apache.iotdb.itbase.exception.InconsistentDataException; import org.awaitility.Awaitility; import org.junit.After; @@ -510,11 +511,17 @@ public void showSpaceQuotaTest1() throws SQLException { } private void validateResultSetEventually(Statement statement, String sql, String ans) { + // ClusterIT fans SHOW out to every DataNode and compares cells. Space-quota usage is updated + // asynchronously via heartbeat, so coordinators may briefly disagree (e.g. [0, 2, 2]) until CN + // converges. That mismatch surfaces as InconsistentDataException from getString(), which + // untilAsserted() does not retry by default — ignore only that type so genuine SQL errors + // still fail fast. Awaitility.await() .pollInSameThread() .pollDelay(0, TimeUnit.MILLISECONDS) .pollInterval(1, TimeUnit.SECONDS) .atMost(30, TimeUnit.SECONDS) + .ignoreExceptionsMatching(e -> e instanceof InconsistentDataException) .untilAsserted( () -> { try { diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/quotas/IoTDBUserResourceQuotaIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/quotas/IoTDBUserResourceQuotaIT.java new file mode 100644 index 0000000000000..3df1043b4ac8e --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/quotas/IoTDBUserResourceQuotaIT.java @@ -0,0 +1,252 @@ +/* + * 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.iotdb.db.it.quotas; + +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.itbase.category.LocalStandaloneIT; + +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +@RunWith(IoTDBTestRunner.class) +@Category({LocalStandaloneIT.class, ClusterIT.class}) +public class IoTDBUserResourceQuotaIT { + + @Before + public void setUp() throws Exception { + EnvFactory.getEnv().getConfig().getDataNodeCommonConfig().setQuotaEnable(true); + EnvFactory.getEnv().initClusterEnvironment(); + } + + @After + public void tearDown() throws Exception { + EnvFactory.getEnv().getConfig().getDataNodeCommonConfig().setQuotaEnable(false); + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void throttleQuotaMapsToReadCpuMax() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE USER quota_user 'pass'"); + statement.execute("GRANT READ_DATA ON root.** TO USER quota_user"); + statement.execute("SET THROTTLE QUOTA cpu=1 ON quota_user"); + statement.execute("CREATE DATABASE root.quota"); + statement.execute("CREATE TIMESERIES root.quota.d1.s1 WITH DATATYPE=INT32, ENCODING=PLAIN"); + statement.execute("INSERT INTO root.quota.d1(time, s1) VALUES (1, 1)"); + try (ResultSet rs = statement.executeQuery("SHOW THROTTLE QUOTA quota_user")) { + Assert.assertTrue(rs.next()); + } + try (Connection userCon = EnvFactory.getEnv().getConnection("quota_user", "pass"); + Statement userStmt = userCon.createStatement()) { + userStmt.executeQuery("SELECT s1 FROM root.quota.d1"); + } + } + } + + @Test + public void setAndShowUserQuota() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE USER urq_user 'pass'"); + statement.execute( + "SET USER QUOTA ON urq_user WITH read_cpu_min=1, read_cpu_max=4, write_memory_max=10M, write_disk_io_max=10485760, write_temp_disk_max=8589934592"); + Map maxByType = new HashMap<>(); + Set types = new HashSet<>(); + try (ResultSet rs = statement.executeQuery("SHOW USER QUOTA urq_user")) { + ResultSetMetaData meta = rs.getMetaData(); + int userIdx = columnIndex(meta, "User"); + int nodeIdx = columnIndex(meta, "NodeID"); + int rwIdx = columnIndex(meta, "Read/Write"); + int typeIdx = columnIndex(meta, "QuotaType"); + int minIdx = columnIndex(meta, "Min"); + int maxIdx = columnIndex(meta, "Max"); + int usedIdx = columnIndex(meta, "Used"); + Assert.assertTrue(columnIndex(meta, "MinGap") > 0); + while (rs.next()) { + Assert.assertEquals("urq_user", rs.getString(userIdx)); + String nodeId = rs.getString(nodeIdx); + Assert.assertNotNull(nodeId); + Assert.assertFalse("-".equals(nodeId)); + Assert.assertTrue(Integer.parseInt(nodeId) > 0); + String key = rs.getString(rwIdx) + ":" + rs.getString(typeIdx); + maxByType.put(key, rs.getString(maxIdx)); + types.add(rs.getString(typeIdx)); + if ("read:cpu".equals(key)) { + Assert.assertEquals("1", rs.getString(minIdx)); + Assert.assertEquals("4", rs.getString(maxIdx)); + } + Assert.assertNotNull(rs.getString(usedIdx)); + } + } + Assert.assertEquals("4", maxByType.get("read:cpu")); + Assert.assertTrue(maxByType.containsKey("write:memory")); + Assert.assertTrue(types.contains("temp_disk")); + Assert.assertTrue(types.contains("disk_io")); + + // At least one alive DataNode id appears in SHOW (dedicated usage-report aggregation). + Set nodeIds = new HashSet<>(); + try (ResultSet rs = statement.executeQuery("SHOW USER QUOTA urq_user")) { + int nodeIdx = columnIndex(rs.getMetaData(), "NodeID"); + while (rs.next()) { + nodeIds.add(Integer.parseInt(rs.getString(nodeIdx))); + } + } + Assert.assertFalse(nodeIds.isEmpty()); + for (Integer id : nodeIds) { + Assert.assertTrue(id > 0); + } + + statement.execute("DELETE USER QUOTA ON urq_user"); + try (ResultSet rs = statement.executeQuery("SHOW USER QUOTA urq_user")) { + Assert.assertFalse(rs.next()); + } + } + } + + @Test + public void lowerMaxDoesNotBreakShowAndNewAcquirePath() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE USER urq_lower 'pass'"); + statement.execute("SET USER QUOTA ON urq_lower WITH read_cpu_max=8"); + statement.execute("SET USER QUOTA ON urq_lower WITH read_cpu_max=2"); + try (ResultSet rs = statement.executeQuery("SHOW USER QUOTA urq_lower")) { + boolean seen = false; + while (rs.next()) { + if ("cpu".equalsIgnoreCase(rs.getString(columnIndex(rs.getMetaData(), "QuotaType"))) + && "read" + .equalsIgnoreCase(rs.getString(columnIndex(rs.getMetaData(), "Read/Write")))) { + Assert.assertEquals("2", rs.getString(columnIndex(rs.getMetaData(), "Max"))); + seen = true; + } + } + Assert.assertTrue(seen); + } + } + } + + @Test + public void rejectRootAndPartialUpdateAndCrossThrottle() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(); + Statement statement = connection.createStatement()) { + try { + statement.execute("SET USER QUOTA ON `root` WITH read_cpu_max=1"); + Assert.fail("root SET should fail"); + } catch (Exception e) { + // expected + } + + statement.execute("CREATE USER urq_partial 'pass'"); + statement.execute("SET USER QUOTA ON urq_partial WITH read_cpu_max=4, write_memory_max=10M"); + statement.execute("SET USER QUOTA ON urq_partial WITH read_cpu_min=1"); + Map values = new HashMap<>(); + try (ResultSet rs = statement.executeQuery("SHOW USER QUOTA urq_partial")) { + ResultSetMetaData meta = rs.getMetaData(); + int rw = columnIndex(meta, "Read/Write"); + int type = columnIndex(meta, "QuotaType"); + int min = columnIndex(meta, "Min"); + int max = columnIndex(meta, "Max"); + while (rs.next()) { + values.put( + rs.getString(rw) + ":" + rs.getString(type), + rs.getString(min) + "/" + rs.getString(max)); + } + } + Assert.assertEquals("1/4", values.get("read:cpu")); + Assert.assertTrue(values.containsKey("write:memory")); + + statement.execute("CREATE USER urq_cross 'pass'"); + statement.execute("SET THROTTLE QUOTA cpu=2 ON urq_cross"); + statement.execute("SET USER QUOTA ON urq_cross WITH write_disk_io_max=1048576"); + Assert.assertTrue(hasQuotaType(statement, "SHOW USER QUOTA urq_cross", "cpu")); + Assert.assertTrue(hasQuotaType(statement, "SHOW USER QUOTA urq_cross", "disk_io")); + Assert.assertTrue(hasAnyRow(statement, "SHOW THROTTLE QUOTA urq_cross")); + } + } + + @Test + public void readCpuMaxEnforcementAndRelease() throws Exception { + try (Connection connection = EnvFactory.getEnv().getConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE USER urq_cpu 'pass'"); + statement.execute("GRANT READ_DATA ON root.** TO USER urq_cpu"); + statement.execute("SET USER QUOTA ON urq_cpu WITH read_cpu_max=1"); + statement.execute("CREATE DATABASE root.urq_cpu"); + statement.execute("CREATE TIMESERIES root.urq_cpu.d1.s1 WITH DATATYPE=INT32,ENCODING=PLAIN"); + statement.execute("INSERT INTO root.urq_cpu.d1(time,s1) VALUES(1,1)"); + + // Single query should succeed (and release). + try (Connection userCon = EnvFactory.getEnv().getConnection("urq_cpu", "pass"); + Statement userStmt = userCon.createStatement()) { + try (ResultSet rs = userStmt.executeQuery("SELECT s1 FROM root.urq_cpu.d1")) { + Assert.assertTrue(rs.next()); + } + try (ResultSet rs = userStmt.executeQuery("SELECT s1 FROM root.urq_cpu.d1")) { + Assert.assertTrue(rs.next()); + } + } + } + } + + private static boolean hasQuotaType(Statement statement, String sql, String quotaType) + throws Exception { + try (ResultSet rs = statement.executeQuery(sql)) { + int typeIdx = columnIndex(rs.getMetaData(), "QuotaType"); + while (rs.next()) { + if (quotaType.equalsIgnoreCase(rs.getString(typeIdx))) { + return true; + } + } + } + return false; + } + + private static boolean hasAnyRow(Statement statement, String sql) throws Exception { + try (ResultSet rs = statement.executeQuery(sql)) { + return rs.next(); + } + } + + private static int columnIndex(ResultSetMetaData meta, String name) throws Exception { + for (int i = 1; i <= meta.getColumnCount(); i++) { + if (name.equalsIgnoreCase(meta.getColumnLabel(i)) + || name.equalsIgnoreCase(meta.getColumnName(i))) { + return i; + } + } + throw new AssertionError("column not found: " + name); + } +} diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 index 135b674554b52..82dc7d68d03a8 100644 --- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 +++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IdentifierParser.g4 @@ -244,6 +244,7 @@ keyWords | SUBSCRIPTION | SUBSCRIPTIONS | SUBSTRING + | SUMMARY | SYSTEM | TABLE | TAG diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 index 5353535ad344e..b61175913972b 100644 --- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 +++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/IoTDBSqlParser.g4 @@ -72,6 +72,7 @@ ddlStatement | callInference | loadModel | unloadModel // Quota | setSpaceQuota | showSpaceQuota | setThrottleQuota | showThrottleQuota + | setUserResourceQuota | showUserResourceQuota | deleteUserResourceQuota // View | createLogicalView | dropLogicalView | showLogicalView | renameLogicalView | alterLogicalView // Table View @@ -400,6 +401,18 @@ showThrottleQuota : SHOW THROTTLE QUOTA (userName=identifier)? ; +setUserResourceQuota + : SET USER QUOTA ON userName=identifier WITH attributePair (COMMA attributePair)* + ; + +showUserResourceQuota + : SHOW USER QUOTA (userName=identifier)? (SUMMARY)? (ON DATANODE dataNodeId=INTEGER_LITERAL)? + ; + +deleteUserResourceQuota + : DELETE USER QUOTA ON userName=identifier + ; + // Trigger ========================================================================================= // ---- Create Trigger createTrigger diff --git a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 index 4d915459c39a8..a252d75e2cace 100644 --- a/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 +++ b/iotdb-core/antlr/src/main/antlr4/org/apache/iotdb/db/qp/sql/SqlLexer.g4 @@ -907,6 +907,10 @@ SUBSTRING : S U B S T R I N G ; +SUMMARY + : S U M M A R Y + ; + SYSTEM : S Y S T E M ; diff --git a/iotdb-core/confignode/src/main/i18n/en/org/apache/iotdb/confignode/i18n/ManagerMessages.java b/iotdb-core/confignode/src/main/i18n/en/org/apache/iotdb/confignode/i18n/ManagerMessages.java index 812c38c043775..47c06c1b8eaf2 100644 --- a/iotdb-core/confignode/src/main/i18n/en/org/apache/iotdb/confignode/i18n/ManagerMessages.java +++ b/iotdb-core/confignode/src/main/i18n/en/org/apache/iotdb/confignode/i18n/ManagerMessages.java @@ -572,6 +572,10 @@ private ManagerMessages() {} public static final String MESSAGE_DATABASE_LIMIT_THRESHOLD_45C23274 = "database_limit_threshold"; public static final String LOG_UNEXPECTED_ERROR_HAPPENED_SETTING_SPACE_QUOTA_DATABASE_ARG_F6ED7586 = "Unexpected error happened while setting space quota on database: %s "; public static final String LOG_UNEXPECTED_ERROR_HAPPENED_SETTING_THROTTLE_QUOTA_USER_ARG_C111BE81 = "Unexpected error happened while setting throttle quota on user: %s "; + public static final String LOG_UNEXPECTED_ERROR_SETTING_USER_RESOURCE_QUOTA_91A4C2E8 = + "Unexpected error happened when setting user resource quota"; + public static final String LOG_FAILED_TO_AGGREGATE_RUNNING_DATANODE_USAGE_FOR_SHOW_USER_QUOTA_00017902 = + "Failed to aggregate Running DataNode usage for SHOW USER QUOTA"; public static final String LOG_SCHEMA_TEMPLATE_NEED_TWO_FILES_1E57542A = "schema_template need two files"; public static final String LOG_GOT_IOEXCEPTION_DESERIALIZE_USE_ROLE_FILE_TYPE_ARG_1B548759 = "Got IOException when deserialize use&role file, type:{}"; public static final String LOG_GOT_IOEXCEPTION_DESERIALIZE_ROLELIST_1354F29E = "Got IOException when deserialize roleList"; diff --git a/iotdb-core/confignode/src/main/i18n/zh/org/apache/iotdb/confignode/i18n/ManagerMessages.java b/iotdb-core/confignode/src/main/i18n/zh/org/apache/iotdb/confignode/i18n/ManagerMessages.java index 9c508beacf438..4f7affa79696a 100644 --- a/iotdb-core/confignode/src/main/i18n/zh/org/apache/iotdb/confignode/i18n/ManagerMessages.java +++ b/iotdb-core/confignode/src/main/i18n/zh/org/apache/iotdb/confignode/i18n/ManagerMessages.java @@ -564,6 +564,10 @@ private ManagerMessages() {} public static final String MESSAGE_DATABASE_LIMIT_THRESHOLD_45C23274 = "database_limit_threshold"; public static final String LOG_UNEXPECTED_ERROR_HAPPENED_SETTING_SPACE_QUOTA_DATABASE_ARG_F6ED7586 = "设置数据库 %s 的空间配额时发生意外错误 "; public static final String LOG_UNEXPECTED_ERROR_HAPPENED_SETTING_THROTTLE_QUOTA_USER_ARG_C111BE81 = "设置用户 %s 的限流配额时发生意外错误 "; + public static final String LOG_UNEXPECTED_ERROR_SETTING_USER_RESOURCE_QUOTA_91A4C2E8 = + "设置用户资源配额时发生意外错误"; + public static final String LOG_FAILED_TO_AGGREGATE_RUNNING_DATANODE_USAGE_FOR_SHOW_USER_QUOTA_00017902 = + "汇总 Running DataNode 用量失败，无法完整展示 SHOW USER QUOTA"; public static final String LOG_SCHEMA_TEMPLATE_NEED_TWO_FILES_1E57542A = "schema_template 需要两个文件"; public static final String LOG_GOT_IOEXCEPTION_DESERIALIZE_USE_ROLE_FILE_TYPE_ARG_1B548759 = "反序列化 use&role 文件时发生 IOException，类型：{}"; public static final String LOG_GOT_IOEXCEPTION_DESERIALIZE_ROLELIST_1354F29E = "反序列化 roleList 时发生 IOException"; diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnAsyncRequestType.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnAsyncRequestType.java index 2d44c214967f6..c9ae3aa5947e4 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnAsyncRequestType.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnAsyncRequestType.java @@ -118,6 +118,7 @@ public enum CnToDnAsyncRequestType { // Quota SET_SPACE_QUOTA, SET_THROTTLE_QUOTA, + SET_USER_RESOURCE_QUOTA, // Table UPDATE_TABLE, diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java index 4048016548a16..a69558d5930fd 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/client/async/CnToDnInternalServiceAsyncRequestManager.java @@ -26,6 +26,7 @@ import org.apache.iotdb.common.rpc.thrift.TSetSpaceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSetTTLReq; import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TSetUserResourceQuotaReq; import org.apache.iotdb.commons.client.async.AsyncDataNodeInternalServiceClient; import org.apache.iotdb.commons.client.request.AsyncRequestContext; import org.apache.iotdb.commons.client.request.AsyncRequestRPCHandler; @@ -410,6 +411,11 @@ protected void initActionMapBuilder() { (req, client, handler) -> client.setThrottleQuota( (TSetThrottleQuotaReq) req, (DataNodeTSStatusRPCHandler) handler)); + actionMapBuilder.put( + CnToDnAsyncRequestType.SET_USER_RESOURCE_QUOTA, + (req, client, handler) -> + client.setUserResourceQuota( + (TSetUserResourceQuotaReq) req, (DataNodeTSStatusRPCHandler) handler)); actionMapBuilder.put( CnToDnAsyncRequestType.RESET_PEER_LIST, (req, client, handler) -> diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java index b7453fb987665..22bb4e8ca36e3 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlan.java @@ -83,8 +83,10 @@ import org.apache.iotdb.confignode.consensus.request.write.pipe.task.SetPipeStatusWithStoppedByRuntimeExceptionPlanV2; import org.apache.iotdb.confignode.consensus.request.write.procedure.DeleteProcedurePlan; import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan; +import org.apache.iotdb.confignode.consensus.request.write.quota.DeleteUserResourceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetSpaceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetThrottleQuotaPlan; +import org.apache.iotdb.confignode.consensus.request.write.quota.SetUserResourceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.region.CreateRegionGroupsPlan; import org.apache.iotdb.confignode.consensus.request.write.region.OfferRegionMaintainTasksPlan; import org.apache.iotdb.confignode.consensus.request.write.region.PollRegionMaintainTaskPlan; @@ -612,6 +614,12 @@ public static ConfigPhysicalPlan create(final ByteBuffer buffer) throws IOExcept case setThrottleQuota: plan = new SetThrottleQuotaPlan(); break; + case setUserResourceQuota: + plan = new SetUserResourceQuotaPlan(); + break; + case deleteUserResourceQuota: + plan = new DeleteUserResourceQuotaPlan(); + break; case CreateExternalService: plan = new CreateExternalServicePlan(); break; diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java index dce1db12cd032..4a81bae6ec0f0 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanType.java @@ -287,6 +287,8 @@ public enum ConfigPhysicalPlanType { /** Quota. */ setSpaceQuota((short) 1400), setThrottleQuota((short) 1401), + setUserResourceQuota((short) 1402), + deleteUserResourceQuota((short) 1403), /** Pipe Task. */ CreatePipeV2((short) 1500), diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/quota/DeleteUserResourceQuotaPlan.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/quota/DeleteUserResourceQuotaPlan.java new file mode 100644 index 0000000000000..97d22e15e2980 --- /dev/null +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/quota/DeleteUserResourceQuotaPlan.java @@ -0,0 +1,82 @@ +/* + * 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.iotdb.confignode.consensus.request.write.quota; + +import org.apache.iotdb.commons.utils.BasicStructureSerDeUtil; +import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlan; +import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Objects; + +public class DeleteUserResourceQuotaPlan extends ConfigPhysicalPlan { + + private String userName; + + public DeleteUserResourceQuotaPlan() { + super(ConfigPhysicalPlanType.deleteUserResourceQuota); + } + + public DeleteUserResourceQuotaPlan(String userName) { + super(ConfigPhysicalPlanType.deleteUserResourceQuota); + this.userName = userName; + } + + public String getUserName() { + return userName; + } + + public void setUserName(String userName) { + this.userName = userName; + } + + @Override + protected void serializeImpl(DataOutputStream stream) throws IOException { + stream.writeShort(getType().getPlanType()); + BasicStructureSerDeUtil.write(userName, stream); + } + + @Override + protected void deserializeImpl(ByteBuffer buffer) throws IOException { + userName = BasicStructureSerDeUtil.readString(buffer); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + if (!super.equals(o)) { + return false; + } + DeleteUserResourceQuotaPlan that = (DeleteUserResourceQuotaPlan) o; + return Objects.equals(userName, that.userName); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), userName); + } +} diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/quota/SetUserResourceQuotaPlan.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/quota/SetUserResourceQuotaPlan.java new file mode 100644 index 0000000000000..0fdd189cddc97 --- /dev/null +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/consensus/request/write/quota/SetUserResourceQuotaPlan.java @@ -0,0 +1,158 @@ +/* + * 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.iotdb.confignode.consensus.request.write.quota; + +import org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange; +import org.apache.iotdb.common.rpc.thrift.TResourceType; +import org.apache.iotdb.common.rpc.thrift.TTimedQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.common.rpc.thrift.ThrottleType; +import org.apache.iotdb.commons.utils.BasicStructureSerDeUtil; +import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlan; +import org.apache.iotdb.confignode.consensus.request.ConfigPhysicalPlanType; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.EnumMap; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +public class SetUserResourceQuotaPlan extends ConfigPhysicalPlan { + + private String userName; + private TUserResourceQuota userResourceQuota; + + public SetUserResourceQuotaPlan() { + super(ConfigPhysicalPlanType.setUserResourceQuota); + } + + public SetUserResourceQuotaPlan(String userName, TUserResourceQuota userResourceQuota) { + super(ConfigPhysicalPlanType.setUserResourceQuota); + this.userName = userName; + this.userResourceQuota = userResourceQuota; + } + + public String getUserName() { + return userName; + } + + public void setUserName(String userName) { + this.userName = userName; + } + + public TUserResourceQuota getUserResourceQuota() { + return userResourceQuota; + } + + public void setUserResourceQuota(TUserResourceQuota userResourceQuota) { + this.userResourceQuota = userResourceQuota; + } + + @Override + protected void serializeImpl(DataOutputStream stream) throws IOException { + stream.writeShort(getType().getPlanType()); + BasicStructureSerDeUtil.write(userName, stream); + serializeQuota(userResourceQuota, stream); + } + + @Override + protected void deserializeImpl(ByteBuffer buffer) throws IOException { + userName = BasicStructureSerDeUtil.readString(buffer); + userResourceQuota = deserializeQuota(buffer); + } + + static void serializeQuota(TUserResourceQuota quota, DataOutputStream stream) throws IOException { + writeRangeMap(quota.getReadQuota(), stream); + writeRangeMap(quota.getWriteQuota(), stream); + Map throttleLimit = + quota.isSetThrottleLimit() ? quota.getThrottleLimit() : new HashMap<>(); + BasicStructureSerDeUtil.write(throttleLimit.size(), stream); + for (Map.Entry entry : throttleLimit.entrySet()) { + BasicStructureSerDeUtil.write(entry.getKey().name(), stream); + BasicStructureSerDeUtil.write(entry.getValue().getTimeUnit(), stream); + BasicStructureSerDeUtil.write(entry.getValue().getSoftLimit(), stream); + } + } + + static TUserResourceQuota deserializeQuota(ByteBuffer buffer) throws IOException { + TUserResourceQuota quota = new TUserResourceQuota(); + quota.setReadQuota(readRangeMap(buffer)); + quota.setWriteQuota(readRangeMap(buffer)); + int throttleSize = BasicStructureSerDeUtil.readInt(buffer); + Map throttleLimit = new HashMap<>(); + for (int i = 0; i < throttleSize; i++) { + ThrottleType type = ThrottleType.valueOf(BasicStructureSerDeUtil.readString(buffer)); + long timeUnit = BasicStructureSerDeUtil.readLong(buffer); + long softLimit = BasicStructureSerDeUtil.readLong(buffer); + throttleLimit.put(type, new TTimedQuota(timeUnit, softLimit)); + } + quota.setThrottleLimit(throttleLimit); + return quota; + } + + private static void writeRangeMap( + Map rangeMap, DataOutputStream stream) + throws IOException { + Map map = + rangeMap == null ? new EnumMap<>(TResourceType.class) : rangeMap; + BasicStructureSerDeUtil.write(map.size(), stream); + for (Map.Entry entry : map.entrySet()) { + BasicStructureSerDeUtil.write(entry.getKey().name(), stream); + BasicStructureSerDeUtil.write(entry.getValue().getMinValue(), stream); + BasicStructureSerDeUtil.write(entry.getValue().getMaxValue(), stream); + } + } + + private static Map readRangeMap(ByteBuffer buffer) + throws IOException { + int size = BasicStructureSerDeUtil.readInt(buffer); + Map map = new EnumMap<>(TResourceType.class); + for (int i = 0; i < size; i++) { + TResourceType type = TResourceType.valueOf(BasicStructureSerDeUtil.readString(buffer)); + long min = BasicStructureSerDeUtil.readLong(buffer); + long max = BasicStructureSerDeUtil.readLong(buffer); + map.put(type, new TResourceQuotaRange(min, max)); + } + return map; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + if (!super.equals(o)) { + return false; + } + SetUserResourceQuotaPlan that = (SetUserResourceQuotaPlan) o; + return Objects.equals(userName, that.userName) + && Objects.equals(userResourceQuota, that.userResourceQuota); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), userName, userResourceQuota); + } +} diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ClusterQuotaManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ClusterQuotaManager.java index d53890f75c341..9719e285662cc 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ClusterQuotaManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ClusterQuotaManager.java @@ -23,20 +23,25 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.common.rpc.thrift.TSetSpaceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TSetUserResourceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSpaceQuota; import org.apache.iotdb.common.rpc.thrift.TThrottleQuota; import org.apache.iotdb.commons.conf.IoTDBConstant; import org.apache.iotdb.confignode.client.async.CnToDnAsyncRequestType; import org.apache.iotdb.confignode.client.async.CnToDnInternalServiceAsyncRequestManager; import org.apache.iotdb.confignode.client.async.handlers.DataNodeAsyncRequestContext; +import org.apache.iotdb.confignode.consensus.request.write.quota.DeleteUserResourceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetSpaceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetThrottleQuotaPlan; +import org.apache.iotdb.confignode.consensus.request.write.quota.SetUserResourceQuotaPlan; import org.apache.iotdb.confignode.i18n.ManagerMessages; import org.apache.iotdb.confignode.manager.partition.PartitionManager; import org.apache.iotdb.confignode.persistence.quota.QuotaInfo; import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq; +import org.apache.iotdb.confignode.rpc.thrift.TShowUserResourceQuotaReq; import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; import org.apache.iotdb.consensus.exception.ConsensusException; import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; @@ -64,6 +69,10 @@ public class ClusterQuotaManager { private final Map> dataRegionIdMap; private final Map regionDisk; + /** dataNodeId -> latest user-resource in-use snapshot from heartbeat. */ + private final Map + userResourceUsageByNode; + public ClusterQuotaManager(IManager configManager, QuotaInfo quotaInfo) { this.configManager = configManager; this.quotaInfo = quotaInfo; @@ -72,6 +81,7 @@ public ClusterQuotaManager(IManager configManager, QuotaInfo quotaInfo) { schemaRegionIdMap = new HashMap<>(); dataRegionIdMap = new HashMap<>(); regionDisk = new ConcurrentHashMap<>(); + userResourceUsageByNode = new ConcurrentHashMap<>(); } public TSStatus setSpaceQuota(final TSetSpaceQuotaReq req) { @@ -245,6 +255,144 @@ public TThrottleQuotaResp getThrottleQuota() { return throttleQuotaResp; } + public TSStatus setUserResourceQuota(TSetUserResourceQuotaReq req) { + if (isClearRequest(req.getUserResourceQuota())) { + return deleteUserResourceQuota(req.getUserName()); + } + try { + TSStatus response = + configManager + .getConsensusManager() + .write(new SetUserResourceQuotaPlan(req.getUserName(), req.getUserResourceQuota())); + if (response.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + // Broadcast the merged quota persisted on ConfigNode, not the partial request. + org.apache.iotdb.common.rpc.thrift.TUserResourceQuota merged = + quotaInfo.getUserResourceQuotaLimit().get(req.getUserName()); + TSetUserResourceQuotaReq broadcastReq = new TSetUserResourceQuotaReq(); + broadcastReq.setUserName(req.getUserName()); + broadcastReq.setUserResourceQuota(merged != null ? merged : req.getUserResourceQuota()); + Map dataNodeLocationMap = + configManager.getNodeManager().getRegisteredDataNodeLocations(); + DataNodeAsyncRequestContext clientHandler = + new DataNodeAsyncRequestContext<>( + CnToDnAsyncRequestType.SET_USER_RESOURCE_QUOTA, broadcastReq, dataNodeLocationMap); + CnToDnInternalServiceAsyncRequestManager.getInstance() + .sendAsyncRequestWithRetry(clientHandler); + return RpcUtils.squashResponseStatusList(clientHandler.getResponseList()); + } + return response; + } catch (ConsensusException e) { + LOGGER.warn(ManagerMessages.LOG_UNEXPECTED_ERROR_SETTING_USER_RESOURCE_QUOTA_91A4C2E8, e); + TSStatus res = new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode()); + res.setMessage(e.getMessage()); + return res; + } + } + + private static boolean isClearRequest( + org.apache.iotdb.common.rpc.thrift.TUserResourceQuota quota) { + if (quota == null) { + return true; + } + boolean readEmpty = + !quota.isSetReadQuota() || quota.getReadQuota() == null || quota.getReadQuota().isEmpty(); + boolean writeEmpty = + !quota.isSetWriteQuota() + || quota.getWriteQuota() == null + || quota.getWriteQuota().isEmpty(); + boolean throttleEmpty = + !quota.isSetThrottleLimit() + || quota.getThrottleLimit() == null + || quota.getThrottleLimit().isEmpty(); + return readEmpty && writeEmpty && throttleEmpty; + } + + public TSStatus deleteUserResourceQuota(String userName) { + try { + TSStatus response = + configManager.getConsensusManager().write(new DeleteUserResourceQuotaPlan(userName)); + if (response.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + TSetUserResourceQuotaReq broadcastReq = new TSetUserResourceQuotaReq(); + broadcastReq.setUserName(userName); + // Empty quota means clear on DataNode (see UserResourceQuotaManager.isClearRequest). + org.apache.iotdb.common.rpc.thrift.TUserResourceQuota empty = + new org.apache.iotdb.common.rpc.thrift.TUserResourceQuota(); + empty.setReadQuota( + new java.util.EnumMap<>(org.apache.iotdb.common.rpc.thrift.TResourceType.class)); + empty.setWriteQuota( + new java.util.EnumMap<>(org.apache.iotdb.common.rpc.thrift.TResourceType.class)); + empty.setThrottleLimit(new java.util.HashMap<>()); + broadcastReq.setUserResourceQuota(empty); + Map dataNodeLocationMap = + configManager.getNodeManager().getRegisteredDataNodeLocations(); + DataNodeAsyncRequestContext clientHandler = + new DataNodeAsyncRequestContext<>( + CnToDnAsyncRequestType.SET_USER_RESOURCE_QUOTA, broadcastReq, dataNodeLocationMap); + CnToDnInternalServiceAsyncRequestManager.getInstance() + .sendAsyncRequestWithRetry(clientHandler); + return RpcUtils.squashResponseStatusList(clientHandler.getResponseList()); + } + return response; + } catch (ConsensusException e) { + LOGGER.warn(ManagerMessages.LOG_UNEXPECTED_ERROR_SETTING_USER_RESOURCE_QUOTA_91A4C2E8, e); + TSStatus res = new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode()); + res.setMessage(e.getMessage()); + return res; + } + } + + public TUserResourceQuotaResp showUserResourceQuota(TShowUserResourceQuotaReq req) { + TUserResourceQuotaResp resp = new TUserResourceQuotaResp(); + if (req.getUserName() == null) { + resp.setUserResourceQuota(quotaInfo.getUserResourceQuotaLimit()); + } else { + Map map = new HashMap<>(); + org.apache.iotdb.common.rpc.thrift.TUserResourceQuota quota = + quotaInfo.getUserResourceQuotaLimit().get(req.getUserName()); + if (quota != null) { + map.put(req.getUserName(), quota); + } + resp.setUserResourceQuota(map); + } + // Aggregate usage for all Running DataNodes (heartbeat cache; empty if not yet reported). + Map usage = + new HashMap<>(); + try { + for (org.apache.iotdb.common.rpc.thrift.TDataNodeConfiguration dn : + configManager + .getNodeManager() + .filterDataNodeThroughStatus(org.apache.iotdb.commons.cluster.NodeStatus.Running)) { + int nodeId = dn.getLocation().getDataNodeId(); + if (req.isSetDataNodeId() && req.getDataNodeId() != nodeId) { + continue; + } + usage.put( + nodeId, + userResourceUsageByNode.getOrDefault( + nodeId, new org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot())); + } + } catch (Exception e) { + LOGGER.warn( + ManagerMessages + .LOG_FAILED_TO_AGGREGATE_RUNNING_DATANODE_USAGE_FOR_SHOW_USER_QUOTA_00017902, + e); + } + if (!usage.isEmpty()) { + resp.setUsageByDataNode(usage); + } + resp.setStatus(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS)); + return resp; + } + + public TUserResourceQuotaResp getUserResourceQuota() { + TUserResourceQuotaResp resp = new TUserResourceQuotaResp(); + if (!quotaInfo.getUserResourceQuotaLimit().isEmpty()) { + resp.setUserResourceQuota(quotaInfo.getUserResourceQuotaLimit()); + } + resp.setStatus(RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS)); + return resp; + } + public Map getSpaceQuotaUsage() { return quotaInfo.getSpaceQuotaUsage(); } @@ -261,6 +409,23 @@ public Map getRegionDisk() { return regionDisk; } + public Map + getUserResourceUsageByNode() { + return userResourceUsageByNode; + } + + /** Accept DN-side usage snapshot from dedicated report RPC (not main heartbeat). */ + public TSStatus reportUserResourceUsage( + int dataNodeId, org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot usage) { + if (usage != null) { + userResourceUsageByNode.put(dataNodeId, usage); + } else { + userResourceUsageByNode.put( + dataNodeId, new org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot()); + } + return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); + } + public void updateSpaceQuotaUsage() { AtomicLong deviceCount = new AtomicLong(); AtomicLong timeSeriesCount = new AtomicLong(); diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java index 21c6e6aedaa18..58a7d098f00ca 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ConfigManager.java @@ -35,6 +35,7 @@ import org.apache.iotdb.common.rpc.thrift.TSetConfigurationReq; import org.apache.iotdb.common.rpc.thrift.TSetSpaceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TSetUserResourceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TShowAppliedConfigurationsResp; import org.apache.iotdb.common.rpc.thrift.TShowConfigurationResp; import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; @@ -247,6 +248,7 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq; import org.apache.iotdb.confignode.rpc.thrift.TShowTopicReq; import org.apache.iotdb.confignode.rpc.thrift.TShowTopicResp; +import org.apache.iotdb.confignode.rpc.thrift.TShowUserResourceQuotaReq; import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp; import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TStartPipeReq; @@ -256,6 +258,7 @@ import org.apache.iotdb.confignode.rpc.thrift.TTimeSlotList; import org.apache.iotdb.confignode.rpc.thrift.TUnsetSchemaTemplateReq; import org.apache.iotdb.confignode.rpc.thrift.TUnsubscribeReq; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; import org.apache.iotdb.consensus.common.DataSet; import org.apache.iotdb.consensus.exception.ConsensusException; import org.apache.iotdb.db.schemaengine.template.TemplateAlterOperationType; @@ -3081,6 +3084,35 @@ public TThrottleQuotaResp getThrottleQuota() { : new TThrottleQuotaResp(status); } + public TSStatus setUserResourceQuota(TSetUserResourceQuotaReq req) { + TSStatus status = confirmLeader(); + return status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() + ? clusterQuotaManager.setUserResourceQuota(req) + : status; + } + + public TUserResourceQuotaResp showUserResourceQuota(TShowUserResourceQuotaReq req) { + TSStatus status = confirmLeader(); + return status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() + ? clusterQuotaManager.showUserResourceQuota(req) + : new TUserResourceQuotaResp(status); + } + + public TUserResourceQuotaResp getUserResourceQuota() { + TSStatus status = confirmLeader(); + return status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() + ? clusterQuotaManager.getUserResourceQuota() + : new TUserResourceQuotaResp(status); + } + + public TSStatus reportUserResourceUsage( + int dataNodeId, org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot usage) { + TSStatus status = confirmLeader(); + return status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() + ? clusterQuotaManager.reportUserResourceUsage(dataNodeId, usage) + : status; + } + @Override public TSStatus createTable(final ByteBuffer tableInfo) { final TSStatus status = confirmLeader(); diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java index 4a70ace8ca19e..c49af00644ce2 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/executor/ConfigPlanExecutor.java @@ -108,8 +108,10 @@ import org.apache.iotdb.confignode.consensus.request.write.pipe.task.SetPipeStatusWithStoppedByRuntimeExceptionPlanV2; import org.apache.iotdb.confignode.consensus.request.write.procedure.DeleteProcedurePlan; import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan; +import org.apache.iotdb.confignode.consensus.request.write.quota.DeleteUserResourceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetSpaceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetThrottleQuotaPlan; +import org.apache.iotdb.confignode.consensus.request.write.quota.SetUserResourceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.region.CreateRegionGroupsPlan; import org.apache.iotdb.confignode.consensus.request.write.region.OfferRegionMaintainTasksPlan; import org.apache.iotdb.confignode.consensus.request.write.region.PollSpecificRegionMaintainTaskPlan; @@ -681,6 +683,10 @@ public TSStatus executeNonQueryPlan(ConfigPhysicalPlan physicalPlan) return quotaInfo.setSpaceQuota((SetSpaceQuotaPlan) physicalPlan); case setThrottleQuota: return quotaInfo.setThrottleQuota((SetThrottleQuotaPlan) physicalPlan); + case setUserResourceQuota: + return quotaInfo.setUserResourceQuota((SetUserResourceQuotaPlan) physicalPlan); + case deleteUserResourceQuota: + return quotaInfo.deleteUserResourceQuota((DeleteUserResourceQuotaPlan) physicalPlan); case CreatePipeSinkV1: case DropPipeV1: case DropPipeSinkV1: diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java index 783ad6c336891..bffa527c7db3b 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/persistence/quota/QuotaInfo.java @@ -19,15 +19,21 @@ package org.apache.iotdb.confignode.persistence.quota; +import org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange; +import org.apache.iotdb.common.rpc.thrift.TResourceType; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.common.rpc.thrift.TSpaceQuota; import org.apache.iotdb.common.rpc.thrift.TThrottleQuota; import org.apache.iotdb.common.rpc.thrift.TTimedQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; import org.apache.iotdb.common.rpc.thrift.ThrottleType; import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.commons.quota.UserResourceQuotaConverter; import org.apache.iotdb.commons.snapshot.SnapshotProcessor; +import org.apache.iotdb.confignode.consensus.request.write.quota.DeleteUserResourceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetSpaceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetThrottleQuotaPlan; +import org.apache.iotdb.confignode.consensus.request.write.quota.SetUserResourceQuotaPlan; import org.apache.iotdb.confignode.i18n.ManagerMessages; import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; @@ -53,6 +59,7 @@ public class QuotaInfo implements SnapshotProcessor { private final Map spaceQuotaLimit; private final Map spaceQuotaUsage; private final Map throttleQuotaLimit; + private final Map userResourceQuotaLimit; private final String snapshotFileName = "quota_info.bin"; @@ -61,6 +68,7 @@ public QuotaInfo() { spaceQuotaLimit = new HashMap<>(); spaceQuotaUsage = new HashMap<>(); throttleQuotaLimit = new HashMap<>(); + userResourceQuotaLimit = new HashMap<>(); } public TSStatus setSpaceQuota(SetSpaceQuotaPlan setSpaceQuotaPlan) { @@ -138,9 +146,125 @@ public TSStatus setThrottleQuota(SetThrottleQuotaPlan setThrottleQuotaPlan) { throttleQuotaLimit.put( setThrottleQuotaPlan.getUserName(), setThrottleQuotaPlan.getThrottleQuota()); } + syncUserResourceFromThrottle(userName); return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); } + public TSStatus setUserResourceQuota(SetUserResourceQuotaPlan plan) { + String userName = plan.getUserName(); + TUserResourceQuota incoming = plan.getUserResourceQuota(); + // Seed from legacy throttle so partial USER QUOTA updates do not wipe cpu/mem limits. + if (!userResourceQuotaLimit.containsKey(userName) && throttleQuotaLimit.containsKey(userName)) { + syncUserResourceFromThrottle(userName); + } + TUserResourceQuota merged = + userResourceQuotaLimit.containsKey(userName) + ? mergeUserResourceQuota(userResourceQuotaLimit.get(userName), incoming) + : incoming; + userResourceQuotaLimit.put(userName, merged); + mergeThrottleFromUserResource(userName, merged); + return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); + } + + public TSStatus deleteUserResourceQuota(DeleteUserResourceQuotaPlan plan) { + String userName = plan.getUserName(); + userResourceQuotaLimit.remove(userName); + throttleQuotaLimit.remove(userName); + return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); + } + + /** + * Incrementally sync derived throttle fields from the merged user resource quota. Never replace + * the whole throttle entry, otherwise a write-only USER QUOTA update would clear prior cpu/mem. + */ + private void mergeThrottleFromUserResource(String userName, TUserResourceQuota merged) { + TThrottleQuota derived = + UserResourceQuotaConverter.toThrottleQuota(UserResourceQuotaConverter.fromThrift(merged)); + TThrottleQuota existing = throttleQuotaLimit.get(userName); + if (existing == null) { + throttleQuotaLimit.put(userName, derived); + return; + } + if (derived.isSetCpuLimit() && derived.getCpuLimit() > 0) { + existing.setCpuLimit(derived.getCpuLimit()); + } + if (derived.isSetMemLimit() && derived.getMemLimit() > 0) { + existing.setMemLimit(derived.getMemLimit()); + } + if (derived.isSetThrottleLimit() && !derived.getThrottleLimit().isEmpty()) { + if (!existing.isSetThrottleLimit()) { + existing.setThrottleLimit(new HashMap<>()); + } + existing.getThrottleLimit().putAll(derived.getThrottleLimit()); + } + } + + private TUserResourceQuota mergeUserResourceQuota( + TUserResourceQuota existing, TUserResourceQuota incoming) { + if (incoming.isSetReadQuota()) { + for (Map.Entry entry : + incoming.getReadQuota().entrySet()) { + mergeRange(existing, entry.getKey(), entry.getValue(), true); + } + } + if (incoming.isSetWriteQuota()) { + for (Map.Entry entry : + incoming.getWriteQuota().entrySet()) { + mergeRange(existing, entry.getKey(), entry.getValue(), false); + } + } + if (incoming.isSetThrottleLimit()) { + if (!existing.isSetThrottleLimit()) { + existing.setThrottleLimit(new HashMap<>()); + } + for (Map.Entry entry : incoming.getThrottleLimit().entrySet()) { + existing.getThrottleLimit().put(entry.getKey(), entry.getValue()); + } + } + return existing; + } + + private void mergeRange( + TUserResourceQuota existing, TResourceType type, TResourceQuotaRange incoming, boolean read) { + TResourceQuotaRange current = + read + ? (existing.isSetReadQuota() ? existing.getReadQuota().get(type) : null) + : (existing.isSetWriteQuota() ? existing.getWriteQuota().get(type) : null); + if (current == null) { + current = + new TResourceQuotaRange(IoTDBConstant.UNLIMITED_VALUE, IoTDBConstant.UNLIMITED_VALUE); + } + if (incoming.getMinValue() != IoTDBConstant.UNLIMITED_VALUE) { + current.setMinValue(incoming.getMinValue()); + } + if (incoming.getMaxValue() != IoTDBConstant.UNLIMITED_VALUE) { + current.setMaxValue(incoming.getMaxValue()); + } + if (read) { + if (!existing.isSetReadQuota()) { + existing.setReadQuota(new HashMap<>()); + } + existing.getReadQuota().put(type, current); + } else { + if (!existing.isSetWriteQuota()) { + existing.setWriteQuota(new HashMap<>()); + } + existing.getWriteQuota().put(type, current); + } + } + + private void syncUserResourceFromThrottle(String userName) { + TThrottleQuota throttle = throttleQuotaLimit.get(userName); + if (throttle == null) { + return; + } + TUserResourceQuota existing = + userResourceQuotaLimit.getOrDefault(userName, new TUserResourceQuota()); + TUserResourceQuota migrated = + UserResourceQuotaConverter.toThrift(UserResourceQuotaConverter.fromThrottleQuota(throttle)); + userResourceQuotaLimit.put(userName, mergeUserResourceQuota(existing, migrated)); + } + public Map getSpaceQuotaLimit() { return spaceQuotaLimit; } @@ -159,6 +283,7 @@ public boolean processTakeSnapshot(File snapshotDir) throws TException, IOExcept try (FileOutputStream fileOutputStream = new FileOutputStream(snapshotFile)) { serializeSpaceQuotaLimit(fileOutputStream); serializeThrottleQuotaLimit(fileOutputStream); + serializeUserResourceQuotaLimit(fileOutputStream); fileOutputStream.getFD().sync(); } finally { spaceQuotaReadWriteLock.writeLock().unlock(); @@ -207,6 +332,11 @@ public void processLoadSnapshot(File snapshotDir) throws TException, IOException clear(); deserializeSpaceQuotaLimit(fileInputStream); deserializeThrottleQuotaLimit(fileInputStream); + if (fileInputStream.available() > 0) { + deserializeUserResourceQuotaLimit(fileInputStream); + } else { + migrateThrottleToUserResourceQuota(); + } } finally { spaceQuotaReadWriteLock.writeLock().unlock(); } @@ -257,8 +387,97 @@ public Map getThrottleQuotaLimit() { return throttleQuotaLimit; } + public Map getUserResourceQuotaLimit() { + return userResourceQuotaLimit; + } + + private void serializeUserResourceQuotaLimit(FileOutputStream fileOutputStream) + throws IOException { + ReadWriteIOUtils.write(userResourceQuotaLimit.size(), fileOutputStream); + for (Map.Entry entry : userResourceQuotaLimit.entrySet()) { + ReadWriteIOUtils.write(entry.getKey(), fileOutputStream); + writeUserResourceQuota(entry.getValue(), fileOutputStream); + } + } + + private void deserializeUserResourceQuotaLimit(FileInputStream fileInputStream) + throws IOException { + int size = ReadWriteIOUtils.readInt(fileInputStream); + while (size > 0) { + String userName = ReadWriteIOUtils.readString(fileInputStream); + userResourceQuotaLimit.put(userName, readUserResourceQuota(fileInputStream)); + size--; + } + } + + private void writeUserResourceQuota(TUserResourceQuota quota, FileOutputStream stream) + throws IOException { + writeRangeMap(quota.isSetReadQuota() ? quota.getReadQuota() : new HashMap<>(), stream); + writeRangeMap(quota.isSetWriteQuota() ? quota.getWriteQuota() : new HashMap<>(), stream); + Map throttleLimit = + quota.isSetThrottleLimit() ? quota.getThrottleLimit() : new HashMap<>(); + ReadWriteIOUtils.write(throttleLimit.size(), stream); + for (Map.Entry entry : throttleLimit.entrySet()) { + ReadWriteIOUtils.write(entry.getKey().name(), stream); + ReadWriteIOUtils.write(entry.getValue().getTimeUnit(), stream); + ReadWriteIOUtils.write(entry.getValue().getSoftLimit(), stream); + } + } + + private TUserResourceQuota readUserResourceQuota(FileInputStream stream) throws IOException { + TUserResourceQuota quota = new TUserResourceQuota(); + quota.setReadQuota(readRangeMap(stream)); + quota.setWriteQuota(readRangeMap(stream)); + int throttleSize = ReadWriteIOUtils.readInt(stream); + Map throttleLimit = new HashMap<>(); + while (throttleSize > 0) { + ThrottleType type = ThrottleType.valueOf(ReadWriteIOUtils.readString(stream)); + long timeUnit = ReadWriteIOUtils.readLong(stream); + long softLimit = ReadWriteIOUtils.readLong(stream); + throttleLimit.put(type, new TTimedQuota(timeUnit, softLimit)); + throttleSize--; + } + quota.setThrottleLimit(throttleLimit); + return quota; + } + + private void writeRangeMap( + Map rangeMap, FileOutputStream stream) + throws IOException { + ReadWriteIOUtils.write(rangeMap.size(), stream); + for (Map.Entry entry : rangeMap.entrySet()) { + ReadWriteIOUtils.write(entry.getKey().name(), stream); + ReadWriteIOUtils.write(entry.getValue().getMinValue(), stream); + ReadWriteIOUtils.write(entry.getValue().getMaxValue(), stream); + } + } + + private Map readRangeMap(FileInputStream stream) + throws IOException { + int size = ReadWriteIOUtils.readInt(stream); + Map map = new HashMap<>(); + while (size > 0) { + TResourceType type = TResourceType.valueOf(ReadWriteIOUtils.readString(stream)); + long min = ReadWriteIOUtils.readLong(stream); + long max = ReadWriteIOUtils.readLong(stream); + map.put(type, new TResourceQuotaRange(min, max)); + size--; + } + return map; + } + + private void migrateThrottleToUserResourceQuota() { + for (Map.Entry entry : throttleQuotaLimit.entrySet()) { + userResourceQuotaLimit.put( + entry.getKey(), + UserResourceQuotaConverter.toThrift( + UserResourceQuotaConverter.fromThrottleQuota(entry.getValue()))); + } + } + public void clear() { spaceQuotaLimit.clear(); throttleQuotaLimit.clear(); + userResourceQuotaLimit.clear(); } } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java index 553b371668265..dd737c18f3f4c 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/service/thrift/ConfigNodeRPCServiceProcessor.java @@ -33,6 +33,7 @@ import org.apache.iotdb.common.rpc.thrift.TSetSpaceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSetTTLReq; import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TSetUserResourceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TShowAppliedConfigurationsResp; import org.apache.iotdb.common.rpc.thrift.TShowConfigurationResp; import org.apache.iotdb.common.rpc.thrift.TShowTTLReq; @@ -222,6 +223,7 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq; import org.apache.iotdb.confignode.rpc.thrift.TShowTopicReq; import org.apache.iotdb.confignode.rpc.thrift.TShowTopicResp; +import org.apache.iotdb.confignode.rpc.thrift.TShowUserResourceQuotaReq; import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp; import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TStartPipeReq; @@ -232,6 +234,7 @@ import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TUnsetSchemaTemplateReq; import org.apache.iotdb.confignode.rpc.thrift.TUnsubscribeReq; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; import org.apache.iotdb.confignode.service.ConfigNode; import org.apache.iotdb.consensus.exception.ConsensusException; import org.apache.iotdb.consensus.exception.ConsensusGroupNotExistException; @@ -1469,6 +1472,27 @@ public TThrottleQuotaResp getThrottleQuota() { return configManager.getThrottleQuota(); } + @Override + public TSStatus setUserResourceQuota(TSetUserResourceQuotaReq req) throws TException { + return configManager.setUserResourceQuota(req); + } + + @Override + public TUserResourceQuotaResp showUserResourceQuota(TShowUserResourceQuotaReq req) { + return configManager.showUserResourceQuota(req); + } + + @Override + public TUserResourceQuotaResp getUserResourceQuota() { + return configManager.getUserResourceQuota(); + } + + @Override + public TSStatus reportUserResourceUsage( + int dataNodeId, org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot usage) { + return configManager.reportUserResourceUsage(dataNodeId, usage); + } + @Override public TSStatus createTable(final ByteBuffer tableInfo) { return configManager.createTable(tableInfo); diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanSerDeTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanSerDeTest.java index e844dcf6910e9..3353f6fcf53bf 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanSerDeTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/consensus/request/ConfigPhysicalPlanSerDeTest.java @@ -29,11 +29,14 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.common.rpc.thrift.TNodeResource; import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange; +import org.apache.iotdb.common.rpc.thrift.TResourceType; import org.apache.iotdb.common.rpc.thrift.TSeriesPartitionSlot; import org.apache.iotdb.common.rpc.thrift.TSpaceQuota; import org.apache.iotdb.common.rpc.thrift.TThrottleQuota; import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.common.rpc.thrift.TTimedQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; import org.apache.iotdb.common.rpc.thrift.ThrottleType; import org.apache.iotdb.commons.consensus.index.impl.IoTProgressIndex; import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex; @@ -116,8 +119,10 @@ import org.apache.iotdb.confignode.consensus.request.write.pipe.task.SetPipeStatusWithStoppedByRuntimeExceptionPlanV2; import org.apache.iotdb.confignode.consensus.request.write.procedure.DeleteProcedurePlan; import org.apache.iotdb.confignode.consensus.request.write.procedure.UpdateProcedurePlan; +import org.apache.iotdb.confignode.consensus.request.write.quota.DeleteUserResourceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetSpaceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.quota.SetThrottleQuotaPlan; +import org.apache.iotdb.confignode.consensus.request.write.quota.SetUserResourceQuotaPlan; import org.apache.iotdb.confignode.consensus.request.write.region.CreateRegionGroupsPlan; import org.apache.iotdb.confignode.consensus.request.write.region.OfferRegionMaintainTasksPlan; import org.apache.iotdb.confignode.consensus.request.write.region.PollRegionMaintainTaskPlan; @@ -198,6 +203,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.EnumMap; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -2067,6 +2073,32 @@ public void setThrottleQuotaPlanTest() throws IOException { Assert.assertEquals(plan.getThrottleQuota(), deserializedPlan.getThrottleQuota()); } + @Test + public void setUserResourceQuotaPlanTest() throws IOException { + TUserResourceQuota quota = new TUserResourceQuota(); + Map readQuota = new EnumMap<>(TResourceType.class); + readQuota.put(TResourceType.CPU, new TResourceQuotaRange(1, 8)); + readQuota.put(TResourceType.MEMORY, new TResourceQuotaRange(-1, 1024)); + quota.setReadQuota(readQuota); + SetUserResourceQuotaPlan plan = new SetUserResourceQuotaPlan("u1", quota); + SetUserResourceQuotaPlan deserialized = + (SetUserResourceQuotaPlan) ConfigPhysicalPlan.Factory.create(plan.serializeToByteBuffer()); + Assert.assertEquals(plan.getUserName(), deserialized.getUserName()); + Assert.assertEquals( + plan.getUserResourceQuota().getReadQuota().get(TResourceType.CPU).getMaxValue(), + deserialized.getUserResourceQuota().getReadQuota().get(TResourceType.CPU).getMaxValue()); + } + + @Test + public void deleteUserResourceQuotaPlanTest() throws IOException { + DeleteUserResourceQuotaPlan plan = new DeleteUserResourceQuotaPlan("u_delete"); + DeleteUserResourceQuotaPlan deserialized = + (DeleteUserResourceQuotaPlan) + ConfigPhysicalPlan.Factory.create(plan.serializeToByteBuffer()); + Assert.assertEquals(plan.getUserName(), deserialized.getUserName()); + Assert.assertEquals(plan.getType(), deserialized.getType()); + } + @Test public void updateClusterIdPlanTest() throws IOException { final String clusterId = String.valueOf(UUID.randomUUID()); diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 7978fce51136e..1b6b4a9b34af3 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -640,6 +640,16 @@ public final class DataNodeQueryMessages { "Please set the number of timeseries greater than 0"; public static final String CANNOT_SET_THROTTLE_QUOTA_FOR_USER_ROOT = "Cannot set throttle quota for user root."; + public static final String EXCEPTION_CANNOT_SET_USER_QUOTA_FOR_USER_ROOT_3AAFE275 = + "Cannot set USER QUOTA for user root"; + public static final String EXCEPTION_INVALID_USER_QUOTA_ATTRIBUTE_ARG_D6CC7292 = + "Invalid USER QUOTA attribute: %s"; + public static final String EXCEPTION_INVALID_USER_QUOTA_DISK_IO_VALUE_ARG_EXPECTED_POSITIVE_LONG_BYTES_SEC_C140D430 = + "Invalid USER QUOTA disk_io value: %s, expected positive long (bytes/sec)"; + public static final String EXCEPTION_INVALID_USER_QUOTA_TEMP_DISK_VALUE_ARG_EXPECTED_POSITIVE_LONG_BYTES_B306C6BF = + "Invalid USER QUOTA temp_disk value: %s, expected positive long (bytes)"; + public static final String EXCEPTION_SHOW_USER_QUOTA_SUMMARY_OR_DATANODE_NOT_SUPPORTED_YET_B7A3E292 = + "SHOW USER QUOTA SUMMARY / ON DATANODE is not supported yet"; public static final String PLEASE_SET_THE_NUMBER_OF_REQUESTS_GREATER_THAN = "Please set the number of requests greater than 0"; public static final String PLEASE_SET_THE_NUMBER_OF_CPU_GREATER_THAN = diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java index f8b6e50ed5372..32be45772d4b5 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -93,6 +93,23 @@ private StorageEngineMessages() {} public static final String THROTTLE_QUOTA_RESTORED_SUCCESSFULLY = "Throttle quota limit restored successfully. "; public static final String THROTTLE_QUOTA_RESTORED_FAILED = "Throttle quota limit restored failed. "; public static final String INVALID_STATEMENT_TYPE = "Invalid statement type: "; + public static final String EXCEPTION_USER_RESOURCE_QUOTA_EXCEEDED_8B2D4E1A = + "User resource quota exceeded: %s"; + public static final String EXCEPTION_USER_RESOURCE_QUOTA_WAIT_TIMEOUT_6F3A1C2D = + "User resource quota wait timeout for %s %s"; + public static final String EXCEPTION_USER_RESOURCE_QUOTA_ACQUIRE_INTERRUPTED_31D4116D = + "User resource quota acquire interrupted for %s"; + public static final String EXCEPTION_USER_MAX_EXCEEDED_3D400C08 = "user max exceeded"; + public static final String EXCEPTION_NODE_CAPACITY_EXCEEDED_89601D9A = "node capacity exceeded"; + public static final String EXCEPTION_MIN_GAP_RESERVATION_B84C4EE4 = "min gap reservation"; + public static final String LOG_USER_RESOURCE_QUOTA_UPDATED_3C8F5A7B = + "User resource quota updated for user {}"; + public static final String LOG_USER_RESOURCE_QUOTA_ACQUIRE_REJECTED_F5079ADB = + "User resource quota acquire rejected for user {} op {} resource {}: {}"; + public static final String LOG_FAILED_TO_REPORT_USER_RESOURCE_USAGE_TO_CONFIGNODE_D08BB930 = + "Failed to report user resource usage to ConfigNode"; + public static final String LOG_USER_RESOURCE_USAGE_REPORT_STARTED_WITH_INTERVAL_ARG_SECONDS_C3CC4CC2 = + "User resource usage report started with interval %s seconds"; // ======================== DataRegion ======================== diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 34a7fa084bf30..7534c0ba42873 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -624,6 +624,16 @@ public final class DataNodeQueryMessages { "请将时间序列数设置为大于 0"; public static final String CANNOT_SET_THROTTLE_QUOTA_FOR_USER_ROOT = "不能为 root 用户设置限流配额。"; + public static final String EXCEPTION_CANNOT_SET_USER_QUOTA_FOR_USER_ROOT_3AAFE275 = + "不能为 root 用户设置 USER QUOTA"; + public static final String EXCEPTION_INVALID_USER_QUOTA_ATTRIBUTE_ARG_D6CC7292 = + "无效的 USER QUOTA 属性：%s"; + public static final String EXCEPTION_INVALID_USER_QUOTA_DISK_IO_VALUE_ARG_EXPECTED_POSITIVE_LONG_BYTES_SEC_C140D430 = + "无效的 USER QUOTA disk_io 值：%s，期望为正整数 long（单位 bytes/sec）"; + public static final String EXCEPTION_INVALID_USER_QUOTA_TEMP_DISK_VALUE_ARG_EXPECTED_POSITIVE_LONG_BYTES_B306C6BF = + "无效的 USER QUOTA temp_disk 值：%s，期望为正整数 long（单位 Byte）"; + public static final String EXCEPTION_SHOW_USER_QUOTA_SUMMARY_OR_DATANODE_NOT_SUPPORTED_YET_B7A3E292 = + "暂不支持 SHOW USER QUOTA SUMMARY / ON DATANODE"; public static final String PLEASE_SET_THE_NUMBER_OF_REQUESTS_GREATER_THAN = "请将请求数设置为大于 0"; public static final String PLEASE_SET_THE_NUMBER_OF_CPU_GREATER_THAN = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java index bffda3d52879e..c373c83d92de8 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -93,6 +93,21 @@ private StorageEngineMessages() {} public static final String THROTTLE_QUOTA_RESTORED_SUCCESSFULLY = "流量配额限制恢复成功。 "; public static final String THROTTLE_QUOTA_RESTORED_FAILED = "流量配额限制恢复失败。 "; public static final String INVALID_STATEMENT_TYPE = "无效的语句类型: "; + public static final String EXCEPTION_USER_RESOURCE_QUOTA_EXCEEDED_8B2D4E1A = "用户资源配额超限：%s"; + public static final String EXCEPTION_USER_RESOURCE_QUOTA_WAIT_TIMEOUT_6F3A1C2D = + "用户资源配额等待超时：%s %s"; + public static final String EXCEPTION_USER_RESOURCE_QUOTA_ACQUIRE_INTERRUPTED_31D4116D = + "用户资源配额申请被中断：%s"; + public static final String EXCEPTION_USER_MAX_EXCEEDED_3D400C08 = "超过用户 max"; + public static final String EXCEPTION_NODE_CAPACITY_EXCEEDED_89601D9A = "超过节点容量"; + public static final String EXCEPTION_MIN_GAP_RESERVATION_B84C4EE4 = "需为其他用户保留 minGap"; + public static final String LOG_USER_RESOURCE_QUOTA_UPDATED_3C8F5A7B = "用户 {} 的资源配额已更新"; + public static final String LOG_USER_RESOURCE_QUOTA_ACQUIRE_REJECTED_F5079ADB = + "用户资源配额申请被拒绝：用户 {}，操作 {}，资源 {}，原因：{}"; + public static final String LOG_FAILED_TO_REPORT_USER_RESOURCE_USAGE_TO_CONFIGNODE_D08BB930 = + "向 ConfigNode 上报用户资源用量失败"; + public static final String LOG_USER_RESOURCE_USAGE_REPORT_STARTED_WITH_INTERVAL_ARG_SECONDS_C3CC4CC2 = + "用户资源用量上报已启动，间隔 %s 秒"; // ======================== DataRegion ======================== diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java index b0c78e496e706..3ae640acad640 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java @@ -1210,6 +1210,15 @@ public class IoTDBConfig { /** Resource control */ private boolean quotaEnable = false; + private int dnQuotaCpuSlots = Math.max(4, Runtime.getRuntime().availableProcessors()); + private long dnQuotaMemoryBytes = Runtime.getRuntime().maxMemory(); + + /** + * Node capacity for user TEMP_DISK quota (temporary disk occupancy), in bytes. Default is + * computed at startup as min(100GiB, totalDataDirSpace/10) unless overridden. + */ + private long dnQuotaTempDiskBytes = -1L; + /** * 1. FixedIntervalRateLimiter : With this limiter resources will be refilled only after a fixed * interval of time. 2. AverageIntervalRateLimiter : This limiter will refill resources at every @@ -4337,6 +4346,59 @@ public void setQuotaEnable(boolean quotaEnable) { this.quotaEnable = quotaEnable; } + public int getDnQuotaCpuSlots() { + return dnQuotaCpuSlots; + } + + public void setDnQuotaCpuSlots(int dnQuotaCpuSlots) { + this.dnQuotaCpuSlots = dnQuotaCpuSlots; + } + + public long getDnQuotaMemoryBytes() { + return dnQuotaMemoryBytes; + } + + public void setDnQuotaMemoryBytes(long dnQuotaMemoryBytes) { + this.dnQuotaMemoryBytes = dnQuotaMemoryBytes; + } + + public long getDnQuotaTempDiskBytes() { + if (dnQuotaTempDiskBytes < 0) { + dnQuotaTempDiskBytes = computeDefaultTempDiskBytes(); + } + return dnQuotaTempDiskBytes; + } + + public void setDnQuotaTempDiskBytes(long dnQuotaTempDiskBytes) { + this.dnQuotaTempDiskBytes = dnQuotaTempDiskBytes; + } + + /** Default: min(100GiB, sum(dataDirs totalSpace) / 10). */ + public long computeDefaultTempDiskBytes() { + final long hundredGiB = 100L * 1024 * 1024 * 1024; + long totalSpace = 0L; + try { + String[] dirs = getDataDirs(); + if (dirs != null) { + for (String dir : dirs) { + if (dir == null) { + continue; + } + long space = new java.io.File(dir).getTotalSpace(); + if (space > 0) { + totalSpace += space; + } + } + } + } catch (Exception ignore) { + // fall through to hundredGiB + } + if (totalSpace <= 0) { + return hundredGiB; + } + return Math.min(hundredGiB, totalSpace / 10); + } + public String getRateLimiterType() { return RateLimiterType; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java index 3544f68662447..f3bb04ee7c2dd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java @@ -1050,6 +1050,20 @@ public void loadProperties(TrimProperties properties) throws BadNodeUrlException conf.setQuotaEnable( Boolean.parseBoolean( properties.getProperty("quota_enable", String.valueOf(conf.isQuotaEnable())))); + conf.setDnQuotaCpuSlots( + Integer.parseInt( + properties.getProperty( + "dn_quota_cpu_slots", String.valueOf(conf.getDnQuotaCpuSlots())))); + conf.setDnQuotaMemoryBytes( + Long.parseLong( + properties.getProperty( + "dn_quota_memory_bytes", String.valueOf(conf.getDnQuotaMemoryBytes())))); + if (properties.containsKey("dn_quota_temp_disk_bytes")) { + conf.setDnQuotaTempDiskBytes( + Long.parseLong(properties.getProperty("dn_quota_temp_disk_bytes"))); + } else { + conf.setDnQuotaTempDiskBytes(conf.computeDefaultTempDiskBytes()); + } // The buffer for sort operator to calculate loadFixedSizeLimitForQuery( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java index 950ac514ce12a..eacf173f9b3e4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/client/ConfigNodeClient.java @@ -32,6 +32,7 @@ import org.apache.iotdb.common.rpc.thrift.TSetSpaceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSetTTLReq; import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TSetUserResourceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TShowAppliedConfigurationsResp; import org.apache.iotdb.common.rpc.thrift.TShowConfigurationResp; import org.apache.iotdb.common.rpc.thrift.TShowTTLReq; @@ -179,6 +180,7 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq; import org.apache.iotdb.confignode.rpc.thrift.TShowTopicReq; import org.apache.iotdb.confignode.rpc.thrift.TShowTopicResp; +import org.apache.iotdb.confignode.rpc.thrift.TShowUserResourceQuotaReq; import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp; import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TStartPipeReq; @@ -189,6 +191,7 @@ import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TUnsetSchemaTemplateReq; import org.apache.iotdb.confignode.rpc.thrift.TUnsubscribeReq; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeMiscMessages; @@ -1486,6 +1489,34 @@ public TThrottleQuotaResp getThrottleQuota() throws TException { () -> client.getThrottleQuota(), resp -> !updateConfigNodeLeader(resp.status)); } + @Override + public TSStatus setUserResourceQuota(TSetUserResourceQuotaReq req) throws TException { + return executeRemoteCallWithRetry( + () -> client.setUserResourceQuota(req), status -> !updateConfigNodeLeader(status)); + } + + @Override + public TUserResourceQuotaResp showUserResourceQuota(TShowUserResourceQuotaReq req) + throws TException { + return executeRemoteCallWithRetry( + () -> client.showUserResourceQuota(req), resp -> !updateConfigNodeLeader(resp.status)); + } + + @Override + public TUserResourceQuotaResp getUserResourceQuota() throws TException { + return executeRemoteCallWithRetry( + () -> client.getUserResourceQuota(), resp -> !updateConfigNodeLeader(resp.status)); + } + + @Override + public TSStatus reportUserResourceUsage( + int dataNodeId, org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot usage) + throws TException { + return executeRemoteCallWithRetry( + () -> client.reportUserResourceUsage(dataNodeId, usage), + status -> !updateConfigNodeLeader(status)); + } + @Override public TSStatus createTable(final ByteBuffer tableInfo) throws TException { return executeRemoteCallWithRetry( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index 93cff6b8ff9a0..df0a4307770b9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java @@ -37,6 +37,7 @@ import org.apache.iotdb.common.rpc.thrift.TSetSpaceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSetTTLReq; import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TSetUserResourceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSettleReq; import org.apache.iotdb.common.rpc.thrift.TShowAppliedConfigurationsResp; import org.apache.iotdb.common.rpc.thrift.TShowConfigurationResp; @@ -218,6 +219,7 @@ import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; import org.apache.iotdb.db.storageengine.rescon.quotas.DataNodeSpaceQuotaManager; import org.apache.iotdb.db.storageengine.rescon.quotas.DataNodeThrottleQuotaManager; +import org.apache.iotdb.db.storageengine.rescon.quotas.UserResourceQuotaManager; import org.apache.iotdb.db.subscription.agent.SubscriptionAgent; import org.apache.iotdb.db.subscription.broker.consensus.ConsensusRegionRuntimeState; import org.apache.iotdb.db.subscription.broker.consensus.ConsensusSubscriptionSetupHandler; @@ -1849,6 +1851,12 @@ public TSStatus setThrottleQuota(TSetThrottleQuotaReq req) throws TException { return throttleQuotaManager.setThrottleQuota(req); } + @Override + public TSStatus setUserResourceQuota(TSetUserResourceQuotaReq req) throws TException { + UserResourceQuotaManager.getInstance().setUserResourceQuota(req); + return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); + } + @Override public TFetchFragmentInstanceStatisticsResp fetchFragmentInstanceStatistics( TFetchFragmentInstanceStatisticsReq req) throws TException { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java index e674199df52fb..a4331b4c6c914 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/common/header/DatasetHeaderFactory.java @@ -224,6 +224,10 @@ public static DatasetHeader getShowThrottleQuotaHeader() { return new DatasetHeader(ColumnHeaderConstant.showThrottleQuotaColumnHeaders, true); } + public static DatasetHeader getShowUserResourceQuotaHeader() { + return new DatasetHeader(ColumnHeaderConstant.showUserResourceQuotaColumnHeaders, true); + } + public static DatasetHeader getShowModelsHeader() { return new DatasetHeader(ColumnHeaderConstant.showModelsColumnHeaders, true); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java index 0485de570fe5b..bcb6e21592802 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverScheduler.java @@ -36,7 +36,6 @@ import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.common.QueryId; import org.apache.iotdb.db.queryengine.exception.CpuNotEnoughException; -import org.apache.iotdb.db.queryengine.execution.driver.DataDriver; import org.apache.iotdb.db.queryengine.execution.driver.IDriver; import org.apache.iotdb.db.queryengine.execution.exchange.IMPPDataExchangeManager; import org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeService; @@ -45,7 +44,10 @@ import org.apache.iotdb.db.queryengine.execution.schedule.queue.multilevelqueue.MultilevelPriorityQueue; import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTask; import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTaskStatus; -import org.apache.iotdb.db.storageengine.rescon.quotas.DataNodeThrottleQuotaManager; +import org.apache.iotdb.db.storageengine.rescon.quotas.AcquireContext; +import org.apache.iotdb.db.storageengine.rescon.quotas.AcquirePolicy; +import org.apache.iotdb.db.storageengine.rescon.quotas.UserResourceQuotaExceededException; +import org.apache.iotdb.db.storageengine.rescon.quotas.UserResourceQuotaManager; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.mpp.rpc.thrift.TFragmentInstanceId; @@ -65,7 +67,6 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; /** The manager of fragment instances scheduling. */ public class DriverScheduler implements IDriverScheduler, IService { @@ -199,10 +200,12 @@ public void submitDrivers( driver.isHighestPriority()))); List submittedTasks = new ArrayList<>(); + List deferredTasks = new ArrayList<>(); for (DriverTask task : tasks) { IDriver driver = task.getDriver(); int dependencyDriverIndex = driver.getDriverContext().getDependencyDriverIndex(); if (dependencyDriverIndex != -1) { + deferredTasks.add(task); SettableFuture blockedDependencyFuture = tasks.get(dependencyDriverIndex).getBlockedDependencyDriver(); blockedDependencyFuture.addListener( @@ -230,43 +233,45 @@ public void submitDrivers( if (IoTDBDescriptor.getInstance().getConfig().isQuotaEnable() && sessionInfo != null && !sessionInfo.getUserName().equals(IoTDBConstant.PATH_ROOT)) { - AtomicInteger usedCpu = new AtomicInteger(); - AtomicLong estimatedMemory = new AtomicLong(); - queryMap - .get(queryId) - .values() - .forEach( - driverTasks -> - driverTasks.forEach( - driverTask -> { - if (driverTask.getStatus().equals(DriverTaskStatus.RUNNING) - && driverTask.getDriver() instanceof DataDriver) { - usedCpu.addAndGet(1); - estimatedMemory.addAndGet(driverTask.getEstimatedMemorySize()); - } - })); - if (!DataNodeThrottleQuotaManager.getInstance() - .getThrottleQuotaLimit() - .checkCpu(sessionInfo.getUserName(), usedCpu.get())) { - throw new CpuNotEnoughException( - DataNodeQueryMessages - .QUERY_EXCEPTION_THERE_IS_NOT_ENOUGH_CPU_TO_EXECUTE_CURRENT_FRAGMENT_INSTANCE_E7719FB8); - } - if (!DataNodeThrottleQuotaManager.getInstance() - .getThrottleQuotaLimit() - .checkMemory(sessionInfo.getUserName(), estimatedMemory.get())) { - throw new MemoryNotEnoughException( - DataNodeQueryMessages - .QUERY_EXCEPTION_THERE_IS_NO_ENOUGH_MEMORY_TO_EXECUTE_CURRENT_FRAGMENT_INSTANCE_CB632843); + List allTasks = new ArrayList<>(submittedTasks.size() + deferredTasks.size()); + allTasks.addAll(submittedTasks); + allTasks.addAll(deferredTasks); + for (DriverTask task : allTasks) { + AcquireContext ctx = + new AcquireContext() + .setQueryId(queryId.getId()) + .setFragmentId(task.getDriverTaskId().getFragmentInstanceId().getFullId()) + .setStatementType("QUERY"); + try { + task.setQuotaTokenBundle( + UserResourceQuotaManager.getInstance() + .acquireReadResources( + sessionInfo.getUserName(), + task.getEstimatedMemorySize(), + ctx, + AcquirePolicy.defaults())); + } catch (UserResourceQuotaExceededException e) { + for (DriverTask acquired : allTasks) { + if (acquired.getQuotaTokenBundle() != null) { + acquired.getQuotaTokenBundle().close(); + acquired.setQuotaTokenBundle(null); + } + } + throw new CpuNotEnoughException(e.getMessage()); + } } } for (DriverTask task : submittedTasks) { registerTaskToQueryMap(queryId, task); } - scheduler.enforceTimeLimit(submittedTasks.get(submittedTasks.size() - 1)); - for (DriverTask task : submittedTasks) { - submitTaskToReadyQueue(task); + // Deferred (dependency) drivers are registered when their listener fires; still need a + // timeout sentinel if this query only has independent drivers. + if (!submittedTasks.isEmpty()) { + scheduler.enforceTimeLimit(submittedTasks.get(submittedTasks.size() - 1)); + for (DriverTask task : submittedTasks) { + submitTaskToReadyQueue(task); + } } } @@ -373,6 +378,10 @@ private void clearDriverTask(DriverTask task) { } timeoutQueue.remove(task.getDriverTaskId()); + if (task.getQuotaTokenBundle() != null) { + task.getQuotaTokenBundle().close(); + task.setQuotaTokenBundle(null); + } Map> queryRelatedTasks = queryMap.get(task.getDriverTaskId().getQueryId()); if (queryRelatedTasks != null) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java index 285c67dd7babc..604d8d3c2e698 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/task/DriverTask.java @@ -31,6 +31,7 @@ import org.apache.iotdb.db.queryengine.execution.schedule.ExecutionContext; import org.apache.iotdb.db.queryengine.execution.schedule.queue.multilevelqueue.DriverTaskHandle; import org.apache.iotdb.db.queryengine.execution.schedule.queue.multilevelqueue.Priority; +import org.apache.iotdb.db.storageengine.rescon.quotas.QuotaTokenBundle; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; @@ -63,6 +64,7 @@ public class DriverTask implements IDIndexedAccessible { private boolean reservedInReadyQueue; private long estimatedMemorySize; + private QuotaTokenBundle quotaTokenBundle; /** Initialize a dummy instance for queryHolder. */ public DriverTask() { @@ -108,6 +110,14 @@ public long getEstimatedMemorySize() { return driver.getEstimatedMemorySize(); } + public QuotaTokenBundle getQuotaTokenBundle() { + return quotaTokenBundle; + } + + public void setQuotaTokenBundle(QuotaTokenBundle quotaTokenBundle) { + this.quotaTokenBundle = quotaTokenBundle; + } + @Override public void setId(ID id) { driver.setDriverTaskId((DriverTaskId) id); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TreeConfigTaskVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TreeConfigTaskVisitor.java index 360116edccede..2a38e87e0504d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TreeConfigTaskVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/TreeConfigTaskVisitor.java @@ -120,10 +120,13 @@ import org.apache.iotdb.db.queryengine.plan.execution.config.sys.pipe.ShowPipeTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.pipe.StartPipeTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.pipe.StopPipeTask; +import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.DeleteUserResourceQuotaTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.SetSpaceQuotaTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.SetThrottleQuotaTask; +import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.SetUserResourceQuotaTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.ShowSpaceQuotaTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.ShowThrottleQuotaTask; +import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.ShowUserResourceQuotaTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.subscription.AlterTopicTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.subscription.CreateTopicTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.subscription.DropSubscriptionTask; @@ -229,10 +232,13 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.StartRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StopRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.TestConnectionStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.DeleteUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowUserResourceQuotaStatement; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.tsfile.exception.NotImplementedException; @@ -932,6 +938,24 @@ public IConfigTask visitShowThrottleQuota( return new ShowThrottleQuotaTask(showThrottleQuotaStatement); } + @Override + public IConfigTask visitSetUserResourceQuota( + SetUserResourceQuotaStatement statement, MPPQueryContext context) { + return new SetUserResourceQuotaTask(statement); + } + + @Override + public IConfigTask visitShowUserResourceQuota( + ShowUserResourceQuotaStatement statement, MPPQueryContext context) { + return new ShowUserResourceQuotaTask(statement); + } + + @Override + public IConfigTask visitDeleteUserResourceQuota( + DeleteUserResourceQuotaStatement statement, MPPQueryContext context) { + return new DeleteUserResourceQuotaTask(statement); + } + @Override public IConfigTask visitSetSqlDialect( SetSqlDialectStatement setSqlDialectStatement, MPPQueryContext context) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java index 92b9b80f1c6e7..b2376aab2dcd9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/ClusterConfigTaskExecutor.java @@ -37,16 +37,19 @@ import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; import org.apache.iotdb.common.rpc.thrift.TExternalServiceEntry; import org.apache.iotdb.common.rpc.thrift.TFlushReq; +import org.apache.iotdb.common.rpc.thrift.TResourceType; import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.common.rpc.thrift.TSetConfigurationReq; import org.apache.iotdb.common.rpc.thrift.TSetSpaceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TSetTTLReq; import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TSetUserResourceQuotaReq; import org.apache.iotdb.common.rpc.thrift.TShowAppliedConfigurationsResp; import org.apache.iotdb.common.rpc.thrift.TShowTTLReq; import org.apache.iotdb.common.rpc.thrift.TSpaceQuota; import org.apache.iotdb.common.rpc.thrift.TTestConnectionResp; import org.apache.iotdb.common.rpc.thrift.TThrottleQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.commons.client.exception.ClientManagerException; import org.apache.iotdb.commons.cluster.NodeStatus; @@ -179,12 +182,14 @@ import org.apache.iotdb.confignode.rpc.thrift.TShowThrottleReq; import org.apache.iotdb.confignode.rpc.thrift.TShowTopicReq; import org.apache.iotdb.confignode.rpc.thrift.TShowTopicResp; +import org.apache.iotdb.confignode.rpc.thrift.TShowUserResourceQuotaReq; import org.apache.iotdb.confignode.rpc.thrift.TShowVariablesResp; import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TStartPipeReq; import org.apache.iotdb.confignode.rpc.thrift.TStopPipeReq; import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TUnsetSchemaTemplateReq; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.BatchProcessException; import org.apache.iotdb.db.exception.StorageEngineException; @@ -264,6 +269,7 @@ import org.apache.iotdb.db.queryengine.plan.execution.config.sys.pipe.ShowPipeTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.ShowSpaceQuotaTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.ShowThrottleQuotaTask; +import org.apache.iotdb.db.queryengine.plan.execution.config.sys.quota.ShowUserResourceQuotaTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.subscription.ShowCreateTopicTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.subscription.ShowSubscriptionsTask; import org.apache.iotdb.db.queryengine.plan.execution.config.sys.subscription.ShowTopicsTask; @@ -326,10 +332,13 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.view.RenameLogicalViewStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.KillQueryStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowConfigurationStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.DeleteUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowUserResourceQuotaStatement; import org.apache.iotdb.db.schemaengine.SchemaEngine; import org.apache.iotdb.db.schemaengine.rescon.DataNodeSchemaQuotaManager; import org.apache.iotdb.db.schemaengine.table.InformationSchemaUtils; @@ -4474,6 +4483,83 @@ public TThrottleQuotaResp getThrottleQuota() { return throttleQuotaResp; } + @Override + public TUserResourceQuotaResp getUserResourceQuota() { + TUserResourceQuotaResp resp = new TUserResourceQuotaResp(); + try (ConfigNodeClient configNodeClient = + CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) { + resp = configNodeClient.getUserResourceQuota(); + } catch (final Exception e) { + LOGGER.error(e.getMessage()); + } + return resp; + } + + @Override + public SettableFuture setUserResourceQuota( + SetUserResourceQuotaStatement statement) { + SettableFuture future = SettableFuture.create(); + TSetUserResourceQuotaReq req = new TSetUserResourceQuotaReq(); + req.setUserName(statement.getUserName()); + req.setUserResourceQuota(statement.getUserResourceQuota()); + try (ConfigNodeClient client = + CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) { + TSStatus tsStatus = client.setUserResourceQuota(req); + if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS)); + } else { + future.setException(new IoTDBException(tsStatus)); + } + } catch (Exception e) { + future.setException(e); + } + return future; + } + + @Override + public SettableFuture deleteUserResourceQuota( + DeleteUserResourceQuotaStatement statement) { + SettableFuture future = SettableFuture.create(); + TSetUserResourceQuotaReq req = new TSetUserResourceQuotaReq(); + req.setUserName(statement.getUserName()); + TUserResourceQuota empty = new TUserResourceQuota(); + empty.setReadQuota(new java.util.EnumMap<>(TResourceType.class)); + empty.setWriteQuota(new java.util.EnumMap<>(TResourceType.class)); + empty.setThrottleLimit(new java.util.HashMap<>()); + req.setUserResourceQuota(empty); + try (ConfigNodeClient client = + CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) { + TSStatus tsStatus = client.setUserResourceQuota(req); + if (tsStatus.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS)); + } else { + future.setException(new IoTDBException(tsStatus)); + } + } catch (Exception e) { + future.setException(e); + } + return future; + } + + @Override + public SettableFuture showUserResourceQuota( + ShowUserResourceQuotaStatement statement) { + SettableFuture future = SettableFuture.create(); + try (ConfigNodeClient client = + CONFIG_NODE_CLIENT_MANAGER.borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) { + TShowUserResourceQuotaReq req = new TShowUserResourceQuotaReq(); + req.setUserName(statement.getUserName()); + if (statement.getDataNodeId() != null) { + req.setDataNodeId(statement.getDataNodeId()); + } + req.setSummary(statement.isSummary()); + ShowUserResourceQuotaTask.buildTSBlock(client.showUserResourceQuota(req), statement, future); + } catch (Exception e) { + future.setException(e); + } + return future; + } + @Override public TSpaceQuotaResp getSpaceQuota() { TSpaceQuotaResp spaceQuotaResp = new TSpaceQuotaResp(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java index 46f82cefe28cf..b966541f426e8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/executor/IConfigTaskExecutor.java @@ -33,6 +33,7 @@ import org.apache.iotdb.confignode.rpc.thrift.TFetchTableResp; import org.apache.iotdb.confignode.rpc.thrift.TSpaceQuotaResp; import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; import org.apache.iotdb.db.protocol.session.IClientSession; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult; @@ -96,10 +97,13 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.view.RenameLogicalViewStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.KillQueryStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.ShowConfigurationStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.DeleteUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowUserResourceQuotaStatement; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; @@ -324,6 +328,15 @@ SettableFuture showThrottleQuota( TThrottleQuotaResp getThrottleQuota(); + TUserResourceQuotaResp getUserResourceQuota(); + + SettableFuture setUserResourceQuota(SetUserResourceQuotaStatement statement); + + SettableFuture deleteUserResourceQuota( + DeleteUserResourceQuotaStatement statement); + + SettableFuture showUserResourceQuota(ShowUserResourceQuotaStatement statement); + TPipeTransferResp handleTransferConfigPlan(String clientId, TPipeTransferReq req); void handlePipeConfigClientExit(String clientId); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/DeleteUserResourceQuotaTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/DeleteUserResourceQuotaTask.java new file mode 100644 index 0000000000000..63447d239b113 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/DeleteUserResourceQuotaTask.java @@ -0,0 +1,42 @@ +/* + * 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.iotdb.db.queryengine.plan.execution.config.sys.quota; + +import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult; +import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask; +import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.DeleteUserResourceQuotaStatement; + +import com.google.common.util.concurrent.ListenableFuture; + +public class DeleteUserResourceQuotaTask implements IConfigTask { + + private final DeleteUserResourceQuotaStatement statement; + + public DeleteUserResourceQuotaTask(DeleteUserResourceQuotaStatement statement) { + this.statement = statement; + } + + @Override + public ListenableFuture execute(IConfigTaskExecutor configTaskExecutor) + throws InterruptedException { + return configTaskExecutor.deleteUserResourceQuota(statement); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/SetUserResourceQuotaTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/SetUserResourceQuotaTask.java new file mode 100644 index 0000000000000..5d136f09ff9da --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/SetUserResourceQuotaTask.java @@ -0,0 +1,42 @@ +/* + * 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.iotdb.db.queryengine.plan.execution.config.sys.quota; + +import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult; +import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask; +import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetUserResourceQuotaStatement; + +import com.google.common.util.concurrent.ListenableFuture; + +public class SetUserResourceQuotaTask implements IConfigTask { + + private final SetUserResourceQuotaStatement statement; + + public SetUserResourceQuotaTask(SetUserResourceQuotaStatement statement) { + this.statement = statement; + } + + @Override + public ListenableFuture execute(IConfigTaskExecutor configTaskExecutor) + throws InterruptedException { + return configTaskExecutor.setUserResourceQuota(statement); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/ShowUserResourceQuotaTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/ShowUserResourceQuotaTask.java new file mode 100644 index 0000000000000..8579a7e079bc3 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/ShowUserResourceQuotaTask.java @@ -0,0 +1,247 @@ +/* + * 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.iotdb.db.queryengine.plan.execution.config.sys.quota; + +import org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange; +import org.apache.iotdb.common.rpc.thrift.TResourceType; +import org.apache.iotdb.common.rpc.thrift.TTimedQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot; +import org.apache.iotdb.common.rpc.thrift.ThrottleType; +import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.commons.schema.column.ColumnHeader; +import org.apache.iotdb.commons.schema.column.ColumnHeaderConstant; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; +import org.apache.iotdb.db.queryengine.common.header.DatasetHeaderFactory; +import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult; +import org.apache.iotdb.db.queryengine.plan.execution.config.IConfigTask; +import org.apache.iotdb.db.queryengine.plan.execution.config.executor.IConfigTaskExecutor; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowUserResourceQuotaStatement; +import org.apache.iotdb.rpc.TSStatusCode; + +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.apache.tsfile.utils.BytesUtils; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.TreeSet; +import java.util.stream.Collectors; + +/** + * SHOW USER QUOTA: one row set per alive DataNode (usage from CN heartbeat cache), NodeID = + * DataNode id. + */ +public class ShowUserResourceQuotaTask implements IConfigTask { + + private final ShowUserResourceQuotaStatement statement; + + public ShowUserResourceQuotaTask(ShowUserResourceQuotaStatement statement) { + this.statement = statement; + } + + @Override + public ListenableFuture execute(IConfigTaskExecutor configTaskExecutor) + throws InterruptedException { + return configTaskExecutor.showUserResourceQuota(statement); + } + + public static void buildTSBlock( + TUserResourceQuotaResp resp, SettableFuture future) { + buildTSBlock(resp, null, future); + } + + public static void buildTSBlock( + TUserResourceQuotaResp resp, + ShowUserResourceQuotaStatement statement, + SettableFuture future) { + List types = + ColumnHeaderConstant.showUserResourceQuotaColumnHeaders.stream() + .map(ColumnHeader::getColumnType) + .collect(Collectors.toList()); + TsBlockBuilder builder = new TsBlockBuilder(types); + + TreeSet nodeIds = new TreeSet<>(); + Map usageByNode = + resp.isSetUsageByDataNode() && resp.getUsageByDataNode() != null + ? resp.getUsageByDataNode() + : Collections.emptyMap(); + nodeIds.addAll(usageByNode.keySet()); + Integer requestedNodeId = statement == null ? null : statement.getDataNodeId(); + if (requestedNodeId != null) { + nodeIds.clear(); + nodeIds.add(requestedNodeId); + } else if (nodeIds.isEmpty()) { + // CN did not return usageByDataNode (e.g. no Running DN visible yet): fall back so SHOW still + // returns quota rows. Normal path always lists all Running DNs from CN. + nodeIds.add(IoTDBDescriptor.getInstance().getConfig().getDataNodeId()); + } + + if (resp.getUserResourceQuota() != null) { + for (Integer nodeId : nodeIds) { + String nodeIdText = String.valueOf(nodeId); + TUserResourceUsageSnapshot snap = usageByNode.get(nodeId); + for (Map.Entry entry : resp.getUserResourceQuota().entrySet()) { + appendRanges( + builder, + entry.getKey(), + nodeIdText, + IoTDBConstant.REQUEST_TYPE_READ, + entry.getValue().getReadQuota(), + snap == null ? null : snap.getReadInUse()); + appendRanges( + builder, + entry.getKey(), + nodeIdText, + IoTDBConstant.REQUEST_TYPE_WRITE, + entry.getValue().getWriteQuota(), + snap == null ? null : snap.getWriteInUse()); + appendDiskIo(builder, entry.getKey(), nodeIdText, entry.getValue().getThrottleLimit()); + } + } + } + DatasetHeader header = DatasetHeaderFactory.getShowUserResourceQuotaHeader(); + future.set(new ConfigTaskResult(TSStatusCode.SUCCESS_STATUS, builder.build(), header)); + } + + private static void appendRanges( + TsBlockBuilder builder, + String user, + String nodeIdText, + String opLabel, + Map ranges, + Map> inUseByUser) { + if (ranges == null) { + return; + } + for (Map.Entry entry : ranges.entrySet()) { + long used = 0L; + if (inUseByUser != null && inUseByUser.containsKey(user)) { + Long v = inUseByUser.get(user).get(entry.getKey()); + if (v != null) { + used = v; + } + } + long min = entry.getValue().getMinValue(); + long minGap = + (min == IoTDBConstant.UNLIMITED_VALUE || min < 0) ? 0L : Math.max(0L, min - used); + appendRow( + builder, + user, + nodeIdText, + opLabel, + entry.getKey().name().toLowerCase(), + format(min), + format(entry.getValue().getMaxValue()), + String.valueOf(used), + formatGap(min, minGap)); + } + } + + private static void appendDiskIo( + TsBlockBuilder builder, + String user, + String nodeIdText, + Map throttleLimit) { + if (throttleLimit == null) { + return; + } + appendDiskIoSide( + builder, + user, + nodeIdText, + IoTDBConstant.REQUEST_TYPE_READ, + throttleLimit.get(ThrottleType.READ_SIZE)); + appendDiskIoSide( + builder, + user, + nodeIdText, + IoTDBConstant.REQUEST_TYPE_WRITE, + throttleLimit.get(ThrottleType.WRITE_SIZE)); + } + + private static void appendDiskIoSide( + TsBlockBuilder builder, + String user, + String nodeIdText, + String opLabel, + TTimedQuota timedQuota) { + if (timedQuota == null) { + return; + } + appendRow( + builder, + user, + nodeIdText, + opLabel, + "disk_io", + "-", + formatDiskIoLimit(timedQuota), + "-", + "-"); + } + + private static void appendRow( + TsBlockBuilder builder, + String user, + String nodeIdText, + String opLabel, + String quotaType, + String min, + String max, + String used, + String minGap) { + builder.getTimeColumnBuilder().writeLong(0L); + builder.getColumnBuilder(0).writeBinary(BytesUtils.valueOf(user)); + builder.getColumnBuilder(1).writeBinary(BytesUtils.valueOf(nodeIdText)); + builder.getColumnBuilder(2).writeBinary(BytesUtils.valueOf(opLabel)); + builder.getColumnBuilder(3).writeBinary(BytesUtils.valueOf(quotaType)); + builder.getColumnBuilder(4).writeBinary(BytesUtils.valueOf(min)); + builder.getColumnBuilder(5).writeBinary(BytesUtils.valueOf(max)); + builder.getColumnBuilder(6).writeBinary(BytesUtils.valueOf(used)); + builder.getColumnBuilder(7).writeBinary(BytesUtils.valueOf(minGap)); + builder.declarePosition(); + } + + private static String format(long value) { + return value == IoTDBConstant.UNLIMITED_VALUE ? "-" : String.valueOf(value); + } + + private static String formatGap(long minValue, long minGap) { + if (minValue == IoTDBConstant.UNLIMITED_VALUE || minValue < 0) { + return "-"; + } + return String.valueOf(minGap); + } + + private static String formatDiskIoLimit(TTimedQuota timedQuota) { + // USER QUOTA disk_io uses fixed unit bytes/sec stored as softLimit with timeUnit=SEC. + long bytesPerSec = + timedQuota.getTimeUnit() <= 0 + ? timedQuota.getSoftLimit() + : timedQuota.getSoftLimit() * IoTDBConstant.SEC / timedQuota.getTimeUnit(); + return String.valueOf(bytesPerSec); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java index 6ddd95593a1ec..e1af166981a78 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/parser/ASTVisitor.java @@ -256,10 +256,13 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.StartRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StopRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.TestConnectionStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.DeleteUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowUserResourceQuotaStatement; import org.apache.iotdb.db.schemaengine.template.TemplateAlterOperationType; import org.apache.iotdb.db.storageengine.load.config.LoadTsFileConfigurator; import org.apache.iotdb.db.utils.DataNodeDateTimeUtils; @@ -4912,6 +4915,148 @@ public Statement visitShowThrottleQuota(IoTDBSqlParser.ShowThrottleQuotaContext return showThrottleQuotaStatement; } + @Override + public Statement visitSetUserResourceQuota(IoTDBSqlParser.SetUserResourceQuotaContext ctx) { + if (!IoTDBDescriptor.getInstance().getConfig().isQuotaEnable()) { + throw new SemanticException(LIMIT_CONFIGURATION_ENABLED_ERROR_MSG); + } + if (parseIdentifier(ctx.userName.getText()).equals(IoTDBConstant.PATH_ROOT)) { + throw new SemanticException( + DataNodeQueryMessages.EXCEPTION_CANNOT_SET_USER_QUOTA_FOR_USER_ROOT_3AAFE275); + } + SetUserResourceQuotaStatement statement = new SetUserResourceQuotaStatement(); + statement.setUserName(parseIdentifier(ctx.userName.getText())); + for (IoTDBSqlParser.AttributePairContext pair : ctx.attributePair()) { + applyUserResourceQuotaAttribute( + statement, + parseAttributeKey(pair.attributeKey()), + parseAttributeValue(pair.attributeValue())); + } + return statement; + } + + @Override + public Statement visitShowUserResourceQuota(IoTDBSqlParser.ShowUserResourceQuotaContext ctx) { + if (!IoTDBDescriptor.getInstance().getConfig().isQuotaEnable()) { + throw new SemanticException(LIMIT_CONFIGURATION_ENABLED_ERROR_MSG); + } + ShowUserResourceQuotaStatement statement = new ShowUserResourceQuotaStatement(); + if (ctx.userName != null) { + statement.setUserName(parseIdentifier(ctx.userName.getText())); + } + if (ctx.SUMMARY() != null) { + statement.setSummary(true); + } + if (ctx.dataNodeId != null) { + statement.setDataNodeId(Integer.parseInt(ctx.dataNodeId.getText())); + } + return statement; + } + + @Override + public Statement visitDeleteUserResourceQuota(IoTDBSqlParser.DeleteUserResourceQuotaContext ctx) { + if (!IoTDBDescriptor.getInstance().getConfig().isQuotaEnable()) { + throw new SemanticException(LIMIT_CONFIGURATION_ENABLED_ERROR_MSG); + } + if (parseIdentifier(ctx.userName.getText()).equals(IoTDBConstant.PATH_ROOT)) { + throw new SemanticException( + DataNodeQueryMessages.EXCEPTION_CANNOT_SET_USER_QUOTA_FOR_USER_ROOT_3AAFE275); + } + DeleteUserResourceQuotaStatement statement = new DeleteUserResourceQuotaStatement(); + statement.setUserName(parseIdentifier(ctx.userName.getText())); + return statement; + } + + private void applyUserResourceQuotaAttribute( + SetUserResourceQuotaStatement statement, String key, String value) { + String[] parts = key.toLowerCase().split("_"); + if (parts.length < 3) { + throw new SemanticException( + String.format( + DataNodeQueryMessages.EXCEPTION_INVALID_USER_QUOTA_ATTRIBUTE_ARG_D6CC7292, key)); + } + SetUserResourceQuotaStatement.OperationSide side; + try { + side = SetUserResourceQuotaStatement.OperationSide.valueOf(parts[0].toUpperCase()); + } catch (IllegalArgumentException e) { + throw new SemanticException( + String.format( + DataNodeQueryMessages.EXCEPTION_INVALID_USER_QUOTA_ATTRIBUTE_ARG_D6CC7292, key)); + } + String bound = parts[parts.length - 1]; + if ("disk".equals(parts[1]) && "io".equals(parts[2])) { + // Fixed unit: bytes/sec as positive long (avoid complex size/time strings). + long bytesPerSec; + try { + bytesPerSec = Long.parseLong(value); + } catch (NumberFormatException e) { + throw new SemanticException( + String.format( + DataNodeQueryMessages + .EXCEPTION_INVALID_USER_QUOTA_DISK_IO_VALUE_ARG_EXPECTED_POSITIVE_LONG_BYTES_SEC_C140D430, + value)); + } + if (bytesPerSec <= 0) { + throw new SemanticException( + String.format( + DataNodeQueryMessages + .EXCEPTION_INVALID_USER_QUOTA_DISK_IO_VALUE_ARG_EXPECTED_POSITIVE_LONG_BYTES_SEC_C140D430, + value)); + } + statement.putDiskIo(side, new TTimedQuota(IoTDBConstant.SEC, bytesPerSec)); + return; + } + SetUserResourceQuotaStatement.ResourceSide resource; + if ("temp".equals(parts[1]) && "disk".equals(parts[2]) && parts.length >= 4) { + resource = SetUserResourceQuotaStatement.ResourceSide.TEMP_DISK; + } else { + try { + resource = SetUserResourceQuotaStatement.ResourceSide.valueOf(parts[1].toUpperCase()); + } catch (IllegalArgumentException e) { + throw new SemanticException( + String.format( + DataNodeQueryMessages.EXCEPTION_INVALID_USER_QUOTA_ATTRIBUTE_ARG_D6CC7292, key)); + } + } + if (!"min".equals(bound) && !"max".equals(bound)) { + throw new SemanticException( + String.format( + DataNodeQueryMessages.EXCEPTION_INVALID_USER_QUOTA_ATTRIBUTE_ARG_D6CC7292, key)); + } + long parsed; + if (resource == SetUserResourceQuotaStatement.ResourceSide.CPU) { + parsed = Long.parseLong(value); + } else if (resource == SetUserResourceQuotaStatement.ResourceSide.TEMP_DISK) { + // Fixed unit: bytes as positive long. + try { + parsed = Long.parseLong(value); + } catch (NumberFormatException e) { + throw new SemanticException( + String.format( + DataNodeQueryMessages + .EXCEPTION_INVALID_USER_QUOTA_TEMP_DISK_VALUE_ARG_EXPECTED_POSITIVE_LONG_BYTES_B306C6BF, + value)); + } + if (parsed <= 0) { + throw new SemanticException( + String.format( + DataNodeQueryMessages + .EXCEPTION_INVALID_USER_QUOTA_TEMP_DISK_VALUE_ARG_EXPECTED_POSITIVE_LONG_BYTES_B306C6BF, + value)); + } + } else { + parsed = parseThrottleQuotaSizeUnit(value.toLowerCase()); + } + long min = -1; + long max = -1; + if ("min".equals(bound)) { + min = parsed; + } else { + max = parsed; + } + statement.putRange(side, resource, min, max); + } + private long parseThrottleQuotaTimeUnit(String timeUnit) { switch (timeUnit.toLowerCase()) { case IoTDBConstant.SEC_UNIT: diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java index a4d20b9b7c78f..457362333d780 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java @@ -165,10 +165,13 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.StartRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StopRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.TestConnectionStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.DeleteUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowUserResourceQuotaStatement; import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; @@ -1964,6 +1967,24 @@ public TSStatus visitShowThrottleQuota( return checkGlobalAuth(context, PrivilegeType.MAINTAIN, () -> ""); } + @Override + public TSStatus visitSetUserResourceQuota( + SetUserResourceQuotaStatement statement, TreeAccessCheckContext context) { + return checkGlobalAuth(context, PrivilegeType.MAINTAIN, () -> ""); + } + + @Override + public TSStatus visitShowUserResourceQuota( + ShowUserResourceQuotaStatement statement, TreeAccessCheckContext context) { + return checkGlobalAuth(context, PrivilegeType.MAINTAIN, () -> ""); + } + + @Override + public TSStatus visitDeleteUserResourceQuota( + DeleteUserResourceQuotaStatement statement, TreeAccessCheckContext context) { + return checkGlobalAuth(context, PrivilegeType.MAINTAIN, () -> ""); + } + @Override public TSStatus visitShowSpaceQuota( ShowSpaceQuotaStatement showSpaceQuotaStatement, TreeAccessCheckContext context) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java index 00d36a03e236f..d2d9af37c976b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementType.java @@ -172,6 +172,9 @@ public enum StatementType { SHOW_SPACE_QUOTA, SET_THROTTLE_QUOTA, SHOW_THROTTLE_QUOTA, + SET_USER_RESOURCE_QUOTA, + SHOW_USER_RESOURCE_QUOTA, + DELETE_USER_RESOURCE_QUOTA, CREATE_LOGICAL_VIEW, DELETE_LOGICAL_VIEW, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java index 5f2f72f431397..63ff5014f05f6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/StatementVisitor.java @@ -153,10 +153,13 @@ import org.apache.iotdb.db.queryengine.plan.statement.sys.StartRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.StopRepairDataStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.TestConnectionStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.DeleteUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetUserResourceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowSpaceQuotaStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowThrottleQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowUserResourceQuotaStatement; /** * This class provides a visitor of {@link StatementNode}, which can be extended to create a visitor @@ -777,6 +780,21 @@ public R visitShowThrottleQuota( return visitStatement(showThrottleQuotaStatement, context); } + public R visitSetUserResourceQuota( + SetUserResourceQuotaStatement setUserResourceQuotaStatement, C context) { + return visitStatement(setUserResourceQuotaStatement, context); + } + + public R visitShowUserResourceQuota( + ShowUserResourceQuotaStatement showUserResourceQuotaStatement, C context) { + return visitStatement(showUserResourceQuotaStatement, context); + } + + public R visitDeleteUserResourceQuota( + DeleteUserResourceQuotaStatement deleteUserResourceQuotaStatement, C context) { + return visitStatement(deleteUserResourceQuotaStatement, context); + } + public R visitShowCurrentTimestamp( ShowCurrentTimestampStatement showCurrentTimestampStatement, C context) { return visitStatement(showCurrentTimestampStatement, context); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/DeleteUserResourceQuotaStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/DeleteUserResourceQuotaStatement.java new file mode 100644 index 0000000000000..d5c1b8c344e78 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/DeleteUserResourceQuotaStatement.java @@ -0,0 +1,61 @@ +/* + * 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.iotdb.db.queryengine.plan.statement.sys.quota; + +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement; +import org.apache.iotdb.db.queryengine.plan.statement.Statement; +import org.apache.iotdb.db.queryengine.plan.statement.StatementType; +import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor; + +import java.util.Collections; +import java.util.List; + +public class DeleteUserResourceQuotaStatement extends Statement implements IConfigStatement { + + private String userName; + + public DeleteUserResourceQuotaStatement() { + statementType = StatementType.DELETE_USER_RESOURCE_QUOTA; + } + + public String getUserName() { + return userName; + } + + public void setUserName(String userName) { + this.userName = userName; + } + + @Override + public QueryType getQueryType() { + return QueryType.WRITE; + } + + @Override + public List getPaths() { + return Collections.emptyList(); + } + + @Override + public R accept(StatementVisitor visitor, C context) { + return visitor.visitDeleteUserResourceQuota(this, context); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/SetUserResourceQuotaStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/SetUserResourceQuotaStatement.java new file mode 100644 index 0000000000000..9ea02e67f987c --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/SetUserResourceQuotaStatement.java @@ -0,0 +1,126 @@ +/* + * 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.iotdb.db.queryengine.plan.statement.sys.quota; + +import org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange; +import org.apache.iotdb.common.rpc.thrift.TResourceType; +import org.apache.iotdb.common.rpc.thrift.TTimedQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.common.rpc.thrift.ThrottleType; +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement; +import org.apache.iotdb.db.queryengine.plan.statement.Statement; +import org.apache.iotdb.db.queryengine.plan.statement.StatementType; +import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor; + +import java.util.Collections; +import java.util.List; + +public class SetUserResourceQuotaStatement extends Statement implements IConfigStatement { + + private String userName; + private TUserResourceQuota userResourceQuota = new TUserResourceQuota(); + + public SetUserResourceQuotaStatement() { + statementType = StatementType.SET_USER_RESOURCE_QUOTA; + } + + public String getUserName() { + return userName; + } + + public void setUserName(String userName) { + this.userName = userName; + } + + public TUserResourceQuota getUserResourceQuota() { + return userResourceQuota; + } + + public void setUserResourceQuota(TUserResourceQuota userResourceQuota) { + this.userResourceQuota = userResourceQuota; + } + + public void putRange(OperationSide side, ResourceSide resource, long min, long max) { + TResourceType type = TResourceType.valueOf(resource.name()); + java.util.Map target; + if (side == OperationSide.READ) { + if (!userResourceQuota.isSetReadQuota()) { + userResourceQuota.setReadQuota(new java.util.EnumMap<>(TResourceType.class)); + } + target = userResourceQuota.getReadQuota(); + } else { + if (!userResourceQuota.isSetWriteQuota()) { + userResourceQuota.setWriteQuota(new java.util.EnumMap<>(TResourceType.class)); + } + target = userResourceQuota.getWriteQuota(); + } + // Merge min/max from separate attributes (e.g. read_cpu_min + read_cpu_max) in one SET. + TResourceQuotaRange existing = target.get(type); + long mergedMin = min; + long mergedMax = max; + if (existing != null) { + if (min < 0) { + mergedMin = existing.getMinValue(); + } + if (max < 0) { + mergedMax = existing.getMaxValue(); + } + } + target.put(type, new TResourceQuotaRange(mergedMin, mergedMax)); + } + + public void putDiskIo(OperationSide side, TTimedQuota timedQuota) { + if (!userResourceQuota.isSetThrottleLimit()) { + userResourceQuota.setThrottleLimit(new java.util.HashMap<>()); + } + userResourceQuota + .getThrottleLimit() + .put( + side == OperationSide.READ ? ThrottleType.READ_SIZE : ThrottleType.WRITE_SIZE, + timedQuota); + } + + public enum OperationSide { + READ, + WRITE + } + + public enum ResourceSide { + CPU, + MEMORY, + TEMP_DISK + } + + @Override + public QueryType getQueryType() { + return QueryType.OTHER; + } + + @Override + public List getPaths() { + return Collections.emptyList(); + } + + @Override + public R accept(StatementVisitor visitor, C context) { + return visitor.visitSetUserResourceQuota(this, context); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/ShowUserResourceQuotaStatement.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/ShowUserResourceQuotaStatement.java new file mode 100644 index 0000000000000..2b2bcecd77b30 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/statement/sys/quota/ShowUserResourceQuotaStatement.java @@ -0,0 +1,79 @@ +/* + * 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.iotdb.db.queryengine.plan.statement.sys.quota; + +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.statement.IConfigStatement; +import org.apache.iotdb.db.queryengine.plan.statement.Statement; +import org.apache.iotdb.db.queryengine.plan.statement.StatementType; +import org.apache.iotdb.db.queryengine.plan.statement.StatementVisitor; + +import java.util.Collections; +import java.util.List; + +public class ShowUserResourceQuotaStatement extends Statement implements IConfigStatement { + + private String userName; + private Integer dataNodeId; + private boolean summary; + + public ShowUserResourceQuotaStatement() { + statementType = StatementType.SHOW_USER_RESOURCE_QUOTA; + } + + public String getUserName() { + return userName; + } + + public void setUserName(String userName) { + this.userName = userName; + } + + public Integer getDataNodeId() { + return dataNodeId; + } + + public void setDataNodeId(Integer dataNodeId) { + this.dataNodeId = dataNodeId; + } + + public boolean isSummary() { + return summary; + } + + public void setSummary(boolean summary) { + this.summary = summary; + } + + @Override + public QueryType getQueryType() { + return QueryType.READ; + } + + @Override + public List getPaths() { + return Collections.emptyList(); + } + + @Override + public R accept(StatementVisitor visitor, C context) { + return visitor.visitShowUserResourceQuota(this, context); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquireContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquireContext.java new file mode 100644 index 0000000000000..6972cb07447cd --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquireContext.java @@ -0,0 +1,64 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +public class AcquireContext { + + private String queryId; + private String fragmentId; + private String requestId; + private String statementType; + + public String getQueryId() { + return queryId; + } + + public AcquireContext setQueryId(String queryId) { + this.queryId = queryId; + return this; + } + + public String getFragmentId() { + return fragmentId; + } + + public AcquireContext setFragmentId(String fragmentId) { + this.fragmentId = fragmentId; + return this; + } + + public String getRequestId() { + return requestId; + } + + public AcquireContext setRequestId(String requestId) { + this.requestId = requestId; + return this; + } + + public String getStatementType() { + return statementType; + } + + public AcquireContext setStatementType(String statementType) { + this.statementType = statementType; + return this; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquirePolicy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquirePolicy.java new file mode 100644 index 0000000000000..088bb8da248fd --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquirePolicy.java @@ -0,0 +1,49 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +public class AcquirePolicy { + + public static final long DEFAULT_MAX_WAIT_MS = 100L; + public static final long DEFAULT_RETRY_INTERVAL_MS = 10L; + + private long maxWaitMs = DEFAULT_MAX_WAIT_MS; + private long retryIntervalMs = DEFAULT_RETRY_INTERVAL_MS; + + public long getMaxWaitMs() { + return maxWaitMs; + } + + public void setMaxWaitMs(long maxWaitMs) { + this.maxWaitMs = maxWaitMs; + } + + public long getRetryIntervalMs() { + return retryIntervalMs; + } + + public void setRetryIntervalMs(long retryIntervalMs) { + this.retryIntervalMs = retryIntervalMs; + } + + public static AcquirePolicy defaults() { + return new AcquirePolicy(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquireResult.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquireResult.java new file mode 100644 index 0000000000000..6e4f0787b22db --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/AcquireResult.java @@ -0,0 +1,53 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +public class AcquireResult { + + private final boolean success; + private final QuotaToken token; + private final String rejectReason; + + private AcquireResult(boolean success, QuotaToken token, String rejectReason) { + this.success = success; + this.token = token; + this.rejectReason = rejectReason; + } + + public static AcquireResult success(QuotaToken token) { + return new AcquireResult(true, token, null); + } + + public static AcquireResult reject(String reason) { + return new AcquireResult(false, null, reason); + } + + public boolean isSuccess() { + return success; + } + + public QuotaToken getToken() { + return token; + } + + public String getRejectReason() { + return rejectReason; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/CapacityResourceLimiter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/CapacityResourceLimiter.java new file mode 100644 index 0000000000000..29c4ae5b9673f --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/CapacityResourceLimiter.java @@ -0,0 +1,98 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.commons.quota.ResourceQuotaRange; +import org.apache.iotdb.db.i18n.StorageEngineMessages; + +/** + * Capacity-type limiter (CPU / MEMORY / TEMP_DISK). + * + *

Unlimited or unconfigured users still participate in node capacity and minGap scheduling so + * that users with a configured min are not starved. + */ +public class CapacityResourceLimiter implements ResourceLimiter { + + private static final ResourceQuotaRange UNLIMITED_RANGE = + new ResourceQuotaRange(IoTDBConstant.UNLIMITED_VALUE, IoTDBConstant.UNLIMITED_VALUE); + + @Override + public LimiterAcquireResult tryAcquire( + String user, long amount, ResourceQuotaRange range, NodeQuotaState node) { + ResourceQuotaRange effective = range == null ? UNLIMITED_RANGE : range; + long cur = node.inUse(user); + + if (effective.getMaxValue() != IoTDBConstant.UNLIMITED_VALUE + && cur + amount > effective.getMaxValue()) { + return LimiterAcquireResult.reject( + StorageEngineMessages.EXCEPTION_USER_MAX_EXCEEDED_3D400C08); + } + + if (node.totalInUse() + amount > node.getNodeCapacity()) { + return LimiterAcquireResult.reject( + StorageEngineMessages.EXCEPTION_NODE_CAPACITY_EXCEEDED_89601D9A); + } + + long freeAfter = node.getNodeCapacity() - node.totalInUse() - amount; + long minGapTotal = calcMinGapTotal(node, user, amount); + if (minGapTotal <= freeAfter) { + node.addInUse(user, amount); + return LimiterAcquireResult.ok(); + } + + if (effective.getMinValue() != IoTDBConstant.UNLIMITED_VALUE && cur < effective.getMinValue()) { + node.addInUse(user, amount); + return LimiterAcquireResult.ok(); + } + return LimiterAcquireResult.reject( + StorageEngineMessages.EXCEPTION_MIN_GAP_RESERVATION_B84C4EE4); + } + + @Override + public void release(String user, long amount, NodeQuotaState node) { + node.removeInUse(user, amount); + } + + @Override + public long getInUse(String user, NodeQuotaState node) { + return node.inUse(user); + } + + @Override + public boolean isEnforced() { + return true; + } + + static long calcMinGapTotal(NodeQuotaState node, String requestUser, long amount) { + long totalGap = 0; + for (String u : node.allUsers()) { + long inUse = node.inUse(u); + if (u.equals(requestUser)) { + inUse += amount; + } + long min = node.min(u); + if (min != IoTDBConstant.UNLIMITED_VALUE && inUse < min) { + totalGap += (min - inUse); + } + } + return totalGap; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/CpuSlotLimiter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/CpuSlotLimiter.java new file mode 100644 index 0000000000000..473b331185f24 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/CpuSlotLimiter.java @@ -0,0 +1,22 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +public class CpuSlotLimiter extends CapacityResourceLimiter {} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/DataNodeThrottleQuotaManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/DataNodeThrottleQuotaManager.java index db7ad470930ea..4ee71cecc66d2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/DataNodeThrottleQuotaManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/DataNodeThrottleQuotaManager.java @@ -21,7 +21,9 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.commons.conf.IoTDBConstant; import org.apache.iotdb.commons.exception.RpcThrottlingException; +import org.apache.iotdb.commons.quota.UserResourceQuotaConverter; import org.apache.iotdb.confignode.rpc.thrift.TThrottleQuotaResp; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.StorageEngineMessages; @@ -57,6 +59,14 @@ public static DataNodeThrottleQuotaManager getInstance() { public TSStatus setThrottleQuota(TSetThrottleQuotaReq req) { throttleQuotaLimit.setQuotas(req); + // Avoid re-entering UserResourceQuotaManager while either singleton is still constructing. + // Live RPC updates sync; recover() only restores local throttle maps. + if (UserResourceQuotaManager.isInitialized()) { + UserResourceQuotaManager.getInstance() + .updateQuotaWithoutThrottleSync( + req.getUserName(), + UserResourceQuotaConverter.fromThrottleQuota(req.getThrottleQuota())); + } return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS); } @@ -76,7 +86,8 @@ public void setThrottleQuotaLimit(ThrottleQuotaLimit throttleQuotaLimit) { * @return the {@link OperationQuota} * @throws RpcThrottlingException if the operation cannot be executed due to quota exceeded. */ - public OperationQuota checkQuota(String userName, Statement s) throws RpcThrottlingException { + public OperationQuota checkQuota(String userName, Statement s) + throws RpcThrottlingException, UserResourceQuotaExceededException { if (!IoTDBDescriptor.getInstance().getConfig().isQuotaEnable()) { return NoopOperationQuota.get(); } @@ -115,9 +126,22 @@ public OperationQuota checkQuota(String userName, Statement s) throws RpcThrottl * @throws RpcThrottlingException if the operation cannot be executed due to quota exceeded. */ private OperationQuota checkQuota(String userName, int numWrites, int numReads, Statement s) - throws RpcThrottlingException { + throws RpcThrottlingException, UserResourceQuotaExceededException { OperationQuota quota = getQuota(userName); quota.checkQuota(numWrites, numReads, s); + if (numWrites > 0 + && IoTDBDescriptor.getInstance().getConfig().isQuotaEnable() + && !IoTDBConstant.PATH_ROOT.equals(userName)) { + AcquireContext ctx = + new AcquireContext() + .setRequestId(String.valueOf(System.nanoTime())) + .setStatementType(s.getType().name()); + QuotaTokenBundle bundle = + UserResourceQuotaManager.getInstance() + .acquireWriteResources( + userName, WriteMemoryEstimator.estimate(s), ctx, AcquirePolicy.defaults()); + return new ResourceAwareOperationQuota(quota, bundle); + } return quota; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/DiskIoLimiter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/DiskIoLimiter.java new file mode 100644 index 0000000000000..048f4e5ab294e --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/DiskIoLimiter.java @@ -0,0 +1,72 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.common.rpc.thrift.ThrottleType; +import org.apache.iotdb.commons.exception.RpcThrottlingException; +import org.apache.iotdb.commons.quota.OperationType; +import org.apache.iotdb.commons.quota.ResourceQuotaRange; + +public class DiskIoLimiter implements ResourceLimiter { + + @Override + public LimiterAcquireResult tryAcquire( + String user, long amount, ResourceQuotaRange range, NodeQuotaState node) { + // DISK_IO uses throttle rate limiter; amount is checked via OperationQuota at RPC layer. + return LimiterAcquireResult.ok(); + } + + @Override + public void release(String user, long amount, NodeQuotaState node) {} + + @Override + public long getInUse(String user, NodeQuotaState node) { + QuotaLimiter limiter = + DataNodeThrottleQuotaManager.getInstance().getThrottleQuotaLimit().getUserLimiter(user); + if (limiter == null) { + return 0; + } + return limiter.getReadAvailable(); + } + + @Override + public boolean isEnforced() { + // DISK_IO is enforced by OperationQuota / throttle rate limiter at the RPC layer. + return false; + } + + public void checkDiskIo(String user, OperationType op, long amount) + throws RpcThrottlingException { + QuotaLimiter limiter = + DataNodeThrottleQuotaManager.getInstance().getThrottleQuotaLimit().getUserLimiter(user); + if (limiter == null) { + return; + } + if (op == OperationType.READ) { + limiter.checkQuota(0, 0, 0, amount); + } else { + limiter.checkQuota(0, amount, 0, 0); + } + } + + public static ThrottleType throttleTypeFor(OperationType op) { + return op == OperationType.READ ? ThrottleType.READ_SIZE : ThrottleType.WRITE_SIZE; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/LimiterAcquireResult.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/LimiterAcquireResult.java new file mode 100644 index 0000000000000..62a5be1f1a14e --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/LimiterAcquireResult.java @@ -0,0 +1,50 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +/** Result of a capacity limiter try-acquire attempt. */ +public final class LimiterAcquireResult { + + private static final LimiterAcquireResult OK = new LimiterAcquireResult(true, null); + + private final boolean success; + private final String rejectReason; + + private LimiterAcquireResult(boolean success, String rejectReason) { + this.success = success; + this.rejectReason = rejectReason; + } + + public static LimiterAcquireResult ok() { + return OK; + } + + public static LimiterAcquireResult reject(String reason) { + return new LimiterAcquireResult(false, reason); + } + + public boolean isSuccess() { + return success; + } + + public String getRejectReason() { + return rejectReason; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/MemoryBudgetLimiter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/MemoryBudgetLimiter.java new file mode 100644 index 0000000000000..81a4ac12aea60 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/MemoryBudgetLimiter.java @@ -0,0 +1,22 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +public class MemoryBudgetLimiter extends CapacityResourceLimiter {} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/NodeQuotaState.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/NodeQuotaState.java new file mode 100644 index 0000000000000..4695afa929ea6 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/NodeQuotaState.java @@ -0,0 +1,111 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.commons.quota.ResourceQuotaRange; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; + +public class NodeQuotaState { + + private final int nodeId; + private final long nodeCapacity; + private final Map inUseByUser = new HashMap<>(); + private final Map rangeByUser = new HashMap<>(); + + public NodeQuotaState(int nodeId, long nodeCapacity) { + this.nodeId = nodeId; + this.nodeCapacity = nodeCapacity; + } + + public int getNodeId() { + return nodeId; + } + + public long getNodeCapacity() { + return nodeCapacity; + } + + public long inUse(String user) { + return inUseByUser.getOrDefault(user, 0L); + } + + public void addInUse(String user, long amount) { + inUseByUser.put(user, inUse(user) + amount); + } + + public void removeInUse(String user, long amount) { + long newValue = Math.max(0, inUse(user) - amount); + if (newValue == 0) { + inUseByUser.remove(user); + } else { + inUseByUser.put(user, newValue); + } + } + + public long min(String user) { + ResourceQuotaRange range = rangeByUser.get(user); + return range == null ? IoTDBConstant.UNLIMITED_VALUE : range.getMinValue(); + } + + public void updateRange(String user, ResourceQuotaRange range) { + if (range == null || range.isUnlimited()) { + rangeByUser.remove(user); + } else { + rangeByUser.put(user, range); + } + } + + /** Remove configured ranges only; keep inUse so in-flight tokens can still release safely. */ + public void clearUserRange(String user) { + rangeByUser.remove(user); + } + + public void clearUser(String user) { + clearUserRange(user); + inUseByUser.remove(user); + } + + public Set allUsers() { + Set users = new java.util.HashSet<>(inUseByUser.keySet()); + users.addAll(rangeByUser.keySet()); + return users; + } + + public long totalInUse() { + return inUseByUser.values().stream().mapToLong(Long::longValue).sum(); + } + + public long remaining() { + return nodeCapacity - totalInUse(); + } + + public Map getInUseByUser() { + return Collections.unmodifiableMap(inUseByUser); + } + + public Map getRangeByUser() { + return Collections.unmodifiableMap(rangeByUser); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/QuotaToken.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/QuotaToken.java new file mode 100644 index 0000000000000..6690f65b032d8 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/QuotaToken.java @@ -0,0 +1,70 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.commons.quota.OperationType; +import org.apache.iotdb.commons.quota.ResourceType; + +public class QuotaToken implements AutoCloseable { + + private final UserResourceQuotaManager manager; + private final String user; + private final OperationType op; + private final ResourceType resource; + private final long amount; + private boolean released; + + public QuotaToken( + UserResourceQuotaManager manager, + String user, + OperationType op, + ResourceType resource, + long amount) { + this.manager = manager; + this.user = user; + this.op = op; + this.resource = resource; + this.amount = amount; + } + + public String getUser() { + return user; + } + + public OperationType getOp() { + return op; + } + + public ResourceType getResource() { + return resource; + } + + public long getAmount() { + return amount; + } + + @Override + public void close() { + if (!released) { + released = true; + manager.release(this); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/QuotaTokenBundle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/QuotaTokenBundle.java new file mode 100644 index 0000000000000..b8ceca3789cb3 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/QuotaTokenBundle.java @@ -0,0 +1,50 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +public class QuotaTokenBundle implements AutoCloseable { + + private final QuotaToken cpuToken; + private final QuotaToken memoryToken; + private final QuotaToken tempDiskToken; + + public QuotaTokenBundle(QuotaToken cpuToken, QuotaToken memoryToken) { + this(cpuToken, memoryToken, null); + } + + public QuotaTokenBundle(QuotaToken cpuToken, QuotaToken memoryToken, QuotaToken tempDiskToken) { + this.cpuToken = cpuToken; + this.memoryToken = memoryToken; + this.tempDiskToken = tempDiskToken; + } + + @Override + public void close() { + if (tempDiskToken != null) { + tempDiskToken.close(); + } + if (memoryToken != null) { + memoryToken.close(); + } + if (cpuToken != null) { + cpuToken.close(); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ResourceAwareOperationQuota.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ResourceAwareOperationQuota.java new file mode 100644 index 0000000000000..ceae9defc1b2d --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ResourceAwareOperationQuota.java @@ -0,0 +1,55 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.commons.exception.RpcThrottlingException; +import org.apache.iotdb.db.queryengine.plan.statement.Statement; + +import java.nio.ByteBuffer; +import java.util.List; + +public class ResourceAwareOperationQuota implements OperationQuota { + + private final OperationQuota delegate; + private final QuotaTokenBundle resourceBundle; + + public ResourceAwareOperationQuota(OperationQuota delegate, QuotaTokenBundle resourceBundle) { + this.delegate = delegate; + this.resourceBundle = resourceBundle; + } + + @Override + public void checkQuota(int numWrites, int numReads, Statement s) throws RpcThrottlingException { + delegate.checkQuota(numWrites, numReads, s); + } + + @Override + public void addReadResult(List queryResult) { + delegate.addReadResult(queryResult); + } + + @Override + public void close() { + delegate.close(); + if (resourceBundle != null) { + resourceBundle.close(); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ResourceLimiter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ResourceLimiter.java new file mode 100644 index 0000000000000..ad2f99e80ad12 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ResourceLimiter.java @@ -0,0 +1,39 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.commons.quota.ResourceQuotaRange; + +public interface ResourceLimiter { + + /** + * Try to acquire {@code amount} of the resource for {@code user}. + * + * @return success with no reason, or reject with a concrete reason string + */ + LimiterAcquireResult tryAcquire( + String user, long amount, ResourceQuotaRange range, NodeQuotaState node); + + void release(String user, long amount, NodeQuotaState node); + + long getInUse(String user, NodeQuotaState node); + + boolean isEnforced(); +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/TempDiskLimiter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/TempDiskLimiter.java new file mode 100644 index 0000000000000..141b9d71396e2 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/TempDiskLimiter.java @@ -0,0 +1,23 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +/** Capacity limiter for temporary disk occupancy (Byte), same rules as MEMORY. */ +public class TempDiskLimiter extends CapacityResourceLimiter {} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ThrottleQuotaLimit.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ThrottleQuotaLimit.java index 14095879ebcdf..eac31ce0d4d2e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ThrottleQuotaLimit.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/ThrottleQuotaLimit.java @@ -36,7 +36,8 @@ public ThrottleQuotaLimit() { } public void setQuotas(TSetThrottleQuotaReq req) { - if (!req.getThrottleQuota().getThrottleLimit().isEmpty()) { + if (req.getThrottleQuota().isSetThrottleLimit() + && !req.getThrottleQuota().getThrottleLimit().isEmpty()) { userQuotaLimiter.put( req.getUserName(), QuotaLimiter.fromThrottle(req.getThrottleQuota().getThrottleLimit())); } @@ -44,6 +45,12 @@ public void setQuotas(TSetThrottleQuotaReq req) { cpuLimit.put(req.getUserName(), req.getThrottleQuota().cpuLimit); } + public void removeQuota(String userName) { + userQuotaLimiter.remove(userName); + memLimit.remove(userName); + cpuLimit.remove(userName); + } + public Map getUserQuotaLimiter() { return userQuotaLimiter; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/UserResourceQuotaExceededException.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/UserResourceQuotaExceededException.java new file mode 100644 index 0000000000000..a8f9fd97c2aee --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/UserResourceQuotaExceededException.java @@ -0,0 +1,30 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.commons.exception.IoTDBException; +import org.apache.iotdb.rpc.TSStatusCode; + +public class UserResourceQuotaExceededException extends IoTDBException { + + public UserResourceQuotaExceededException(String message) { + super(message, TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode()); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/UserResourceQuotaManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/UserResourceQuotaManager.java new file mode 100644 index 0000000000000..cd7bf2d7079e4 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/UserResourceQuotaManager.java @@ -0,0 +1,484 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.common.rpc.thrift.TSetThrottleQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TSetUserResourceQuotaReq; +import org.apache.iotdb.common.rpc.thrift.TThrottleQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.commons.client.exception.ClientManagerException; +import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory; +import org.apache.iotdb.commons.concurrent.threadpool.ScheduledExecutorUtil; +import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.commons.quota.OperationType; +import org.apache.iotdb.commons.quota.ResourceQuotaRange; +import org.apache.iotdb.commons.quota.ResourceType; +import org.apache.iotdb.commons.quota.UserResourceQuota; +import org.apache.iotdb.commons.quota.UserResourceQuotaConverter; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.protocol.client.ConfigNodeClient; +import org.apache.iotdb.db.protocol.client.ConfigNodeClientManager; +import org.apache.iotdb.db.protocol.client.ConfigNodeInfo; +import org.apache.iotdb.db.queryengine.plan.execution.config.executor.ClusterConfigTaskExecutor; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.thrift.TException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.EnumMap; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +public class UserResourceQuotaManager { + + private static final Logger LOGGER = LoggerFactory.getLogger(UserResourceQuotaManager.class); + + private final Map userQuotas = new ConcurrentHashMap<>(); + private final Map nodeStates = new EnumMap<>(ResourceType.class); + private final Map limiters = new EnumMap<>(ResourceType.class); + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + private final int nodeId; + private static volatile boolean initialized = false; + + /** Low-frequency dedicated report; do not piggyback on main DataNode heartbeat. */ + private static final long USAGE_REPORT_INTERVAL_SECONDS = 10L; + + private UserResourceQuotaManager() { + nodeId = IoTDBDescriptor.getInstance().getConfig().getDataNodeId(); + long cpuCapacity = IoTDBDescriptor.getInstance().getConfig().getDnQuotaCpuSlots(); + long memCapacity = IoTDBDescriptor.getInstance().getConfig().getDnQuotaMemoryBytes(); + long tempDiskCapacity = IoTDBDescriptor.getInstance().getConfig().getDnQuotaTempDiskBytes(); + nodeStates.put(ResourceType.CPU, new NodeQuotaState(nodeId, cpuCapacity)); + nodeStates.put(ResourceType.MEMORY, new NodeQuotaState(nodeId, memCapacity)); + nodeStates.put(ResourceType.TEMP_DISK, new NodeQuotaState(nodeId, tempDiskCapacity)); + limiters.put(ResourceType.CPU, new CpuSlotLimiter()); + limiters.put(ResourceType.MEMORY, new MemoryBudgetLimiter()); + limiters.put(ResourceType.DISK_IO, new DiskIoLimiter()); + limiters.put(ResourceType.TEMP_DISK, new TempDiskLimiter()); + recover(); + startUsageReport(); + initialized = true; + } + + private void startUsageReport() { + ScheduledExecutorService executor = + IoTDBThreadPoolFactory.newSingleThreadScheduledExecutor("UserResourceQuota-UsageReport"); + ScheduledExecutorUtil.safelyScheduleWithFixedDelay( + executor, + this::reportUsageToConfigNode, + USAGE_REPORT_INTERVAL_SECONDS, + USAGE_REPORT_INTERVAL_SECONDS, + TimeUnit.SECONDS); + LOGGER.info( + String.format( + StorageEngineMessages + .LOG_USER_RESOURCE_USAGE_REPORT_STARTED_WITH_INTERVAL_ARG_SECONDS_C3CC4CC2, + USAGE_REPORT_INTERVAL_SECONDS)); + } + + private void reportUsageToConfigNode() { + if (!isEnabled()) { + return; + } + try (ConfigNodeClient client = + ConfigNodeClientManager.getInstance().borrowClient(ConfigNodeInfo.CONFIG_REGION_ID)) { + TSStatus status = client.reportUserResourceUsage(nodeId, snapshotUsage()); + if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.debug( + StorageEngineMessages.LOG_FAILED_TO_REPORT_USER_RESOURCE_USAGE_TO_CONFIGNODE_D08BB930 + + ": {}", + status); + } + } catch (ClientManagerException | TException e) { + LOGGER.debug( + StorageEngineMessages.LOG_FAILED_TO_REPORT_USER_RESOURCE_USAGE_TO_CONFIGNODE_D08BB930, e); + } + } + + private static class Holder { + private static final UserResourceQuotaManager INSTANCE = new UserResourceQuotaManager(); + } + + public static UserResourceQuotaManager getInstance() { + return Holder.INSTANCE; + } + + public static boolean isInitialized() { + return initialized; + } + + public AcquireResult acquire( + String user, + OperationType op, + ResourceType resource, + long amount, + AcquireContext ctx, + AcquirePolicy policy) { + if (!isEnabled() || isExempt(user)) { + return AcquireResult.success(new QuotaToken(this, user, op, resource, 0)); + } + // Unconfigured / unlimited users still go through capacity scheduling so that + // configured min guarantees are not starved (unlimited != skip node checks). + ResourceQuotaRange range = getRange(user, op, resource); + ResourceLimiter limiter = limiters.get(resource); + if (limiter == null || !limiter.isEnforced()) { + return AcquireResult.success(new QuotaToken(this, user, op, resource, 0)); + } + NodeQuotaState node = nodeStates.get(resource); + if (node == null) { + return AcquireResult.success(new QuotaToken(this, user, op, resource, 0)); + } + + long deadline = System.currentTimeMillis() + policy.getMaxWaitMs(); + String lastRejectReason = StorageEngineMessages.EXCEPTION_NODE_CAPACITY_EXCEEDED_89601D9A; + while (true) { + lock.writeLock().lock(); + try { + LimiterAcquireResult attempt = limiter.tryAcquire(user, amount, range, node); + if (attempt.isSuccess()) { + return AcquireResult.success(new QuotaToken(this, user, op, resource, amount)); + } + lastRejectReason = attempt.getRejectReason(); + } finally { + lock.writeLock().unlock(); + } + if (System.currentTimeMillis() >= deadline) { + String reason = + String.format( + StorageEngineMessages.EXCEPTION_USER_RESOURCE_QUOTA_WAIT_TIMEOUT_6F3A1C2D, + op.name().toLowerCase(), + resource.name().toLowerCase()) + + ": " + + lastRejectReason; + logAcquireRejected(user, op, resource, reason); + return AcquireResult.reject(reason); + } + try { + Thread.sleep(policy.getRetryIntervalMs()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + String reason = + String.format( + StorageEngineMessages.EXCEPTION_USER_RESOURCE_QUOTA_ACQUIRE_INTERRUPTED_31D4116D, + op.name().toLowerCase() + " " + resource.name().toLowerCase()); + logAcquireRejected(user, op, resource, reason); + return AcquireResult.reject(reason); + } + } + } + + private void logAcquireRejected( + String user, OperationType op, ResourceType resource, String reason) { + LOGGER.warn( + StorageEngineMessages.LOG_USER_RESOURCE_QUOTA_ACQUIRE_REJECTED_F5079ADB, + user, + op.name().toLowerCase(), + resource.name().toLowerCase(), + reason); + } + + public QuotaToken acquireOrThrow( + String user, + OperationType op, + ResourceType resource, + long amount, + AcquireContext ctx, + AcquirePolicy policy) + throws UserResourceQuotaExceededException { + AcquireResult result = acquire(user, op, resource, amount, ctx, policy); + if (!result.isSuccess()) { + throw new UserResourceQuotaExceededException(result.getRejectReason()); + } + return result.getToken(); + } + + public void release(QuotaToken token) { + if (!isEnabled() || token.getAmount() <= 0) { + return; + } + ResourceLimiter limiter = limiters.get(token.getResource()); + NodeQuotaState node = nodeStates.get(token.getResource()); + if (limiter == null || node == null) { + return; + } + lock.writeLock().lock(); + try { + limiter.release(token.getUser(), token.getAmount(), node); + } finally { + lock.writeLock().unlock(); + } + } + + public long getInUse(String user, OperationType op, ResourceType resource) { + NodeQuotaState node = nodeStates.get(resource); + if (node == null) { + return 0; + } + return node.inUse(user); + } + + public void updateQuota(String user, UserResourceQuota quota) { + updateQuotaInternal(user, quota, true); + } + + void updateQuotaWithoutThrottleSync(String user, UserResourceQuota quota) { + updateQuotaInternal(user, quota, false); + } + + private void updateQuotaInternal(String user, UserResourceQuota quota, boolean syncThrottle) { + if (quota == null) { + return; + } + lock.writeLock().lock(); + try { + UserResourceQuota existing = userQuotas.computeIfAbsent(user, k -> new UserResourceQuota()); + existing.mergeFrom(quota); + refreshNodeRanges(user, existing); + if (syncThrottle) { + syncThrottleQuota(user, existing); + } + } finally { + lock.writeLock().unlock(); + } + LOGGER.info(StorageEngineMessages.LOG_USER_RESOURCE_QUOTA_UPDATED_3C8F5A7B, user); + } + + public void setUserResourceQuota(TSetUserResourceQuotaReq req) { + if (isClearRequest(req.getUserResourceQuota())) { + clearUserQuota(req.getUserName()); + return; + } + updateQuota( + req.getUserName(), UserResourceQuotaConverter.fromThrift(req.getUserResourceQuota())); + } + + /** True when CN broadcasts an empty quota meaning DELETE USER QUOTA. */ + static boolean isClearRequest(TUserResourceQuota quota) { + if (quota == null) { + return true; + } + boolean readEmpty = + !quota.isSetReadQuota() || quota.getReadQuota() == null || quota.getReadQuota().isEmpty(); + boolean writeEmpty = + !quota.isSetWriteQuota() + || quota.getWriteQuota() == null + || quota.getWriteQuota().isEmpty(); + boolean throttleEmpty = + !quota.isSetThrottleLimit() + || quota.getThrottleLimit() == null + || quota.getThrottleLimit().isEmpty(); + return readEmpty && writeEmpty && throttleEmpty; + } + + public void clearUserQuota(String user) { + lock.writeLock().lock(); + try { + userQuotas.remove(user); + for (NodeQuotaState node : nodeStates.values()) { + // Keep inUse for in-flight tokens; only clear configured ranges (DELETE = unlimited). + node.clearUserRange(user); + } + DataNodeThrottleQuotaManager.getInstance().getThrottleQuotaLimit().removeQuota(user); + } finally { + lock.writeLock().unlock(); + } + LOGGER.info(StorageEngineMessages.LOG_USER_RESOURCE_QUOTA_UPDATED_3C8F5A7B, user); + } + + public long getMinGap(String user, OperationType op, ResourceType resource) { + ResourceQuotaRange range = getRange(user, op, resource); + if (range == null + || range.getMinValue() == IoTDBConstant.UNLIMITED_VALUE + || range.getMinValue() < 0) { + return 0; + } + long inUse = getInUse(user, op, resource); + return Math.max(0, range.getMinValue() - inUse); + } + + public UserResourceQuota getUserQuota(String user) { + return userQuotas.get(user); + } + + public Map getAllUserQuotas() { + return new HashMap<>(userQuotas); + } + + /** Snapshot current in-use for dedicated report RPC to ConfigNode (SHOW aggregation). */ + public org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot snapshotUsage() { + org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot snap = + new org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot(); + Map> read = new HashMap<>(); + Map> write = + new HashMap<>(); + lock.readLock().lock(); + try { + for (Map.Entry entry : nodeStates.entrySet()) { + org.apache.iotdb.common.rpc.thrift.TResourceType tType = toThriftType(entry.getKey()); + if (tType == null) { + continue; + } + for (Map.Entry usage : entry.getValue().getInUseByUser().entrySet()) { + if (usage.getValue() == null || usage.getValue() <= 0) { + continue; + } + read.computeIfAbsent(usage.getKey(), k -> new HashMap<>()).put(tType, usage.getValue()); + write.computeIfAbsent(usage.getKey(), k -> new HashMap<>()).put(tType, usage.getValue()); + } + } + } finally { + lock.readLock().unlock(); + } + if (!read.isEmpty()) { + snap.setReadInUse(read); + } + if (!write.isEmpty()) { + snap.setWriteInUse(write); + } + return snap; + } + + private static org.apache.iotdb.common.rpc.thrift.TResourceType toThriftType( + ResourceType resource) { + switch (resource) { + case CPU: + return org.apache.iotdb.common.rpc.thrift.TResourceType.CPU; + case MEMORY: + return org.apache.iotdb.common.rpc.thrift.TResourceType.MEMORY; + case TEMP_DISK: + return org.apache.iotdb.common.rpc.thrift.TResourceType.TEMP_DISK; + case DISK_IO: + return org.apache.iotdb.common.rpc.thrift.TResourceType.DISK_IO; + default: + return null; + } + } + + public int getNodeId() { + return nodeId; + } + + public NodeQuotaState getNodeState(ResourceType resource) { + return nodeStates.get(resource); + } + + private ResourceQuotaRange getRange(String user, OperationType op, ResourceType resource) { + UserResourceQuota quota = userQuotas.get(user); + if (quota == null) { + return null; + } + return quota.getRange(op, resource); + } + + private void refreshNodeRanges(String user, UserResourceQuota quota) { + for (OperationType op : OperationType.values()) { + Map ranges = + op == OperationType.READ ? quota.getReadQuota() : quota.getWriteQuota(); + for (Map.Entry entry : ranges.entrySet()) { + NodeQuotaState node = nodeStates.get(entry.getKey()); + if (node != null) { + node.updateRange(user, entry.getValue()); + } + } + } + } + + private void syncThrottleQuota(String user, UserResourceQuota quota) { + TThrottleQuota throttle = UserResourceQuotaConverter.toThrottleQuota(quota); + TSetThrottleQuotaReq req = new TSetThrottleQuotaReq(); + req.setUserName(user); + req.setThrottleQuota(throttle); + DataNodeThrottleQuotaManager.getInstance().getThrottleQuotaLimit().setQuotas(req); + } + + private void recover() { + TUserResourceQuotaResp resp = ClusterConfigTaskExecutor.getInstance().getUserResourceQuota(); + if (resp == null || resp.getUserResourceQuota() == null) { + return; + } + for (Map.Entry entry : resp.getUserResourceQuota().entrySet()) { + // Do not sync throttle during recover; DataNodeThrottleQuotaManager recovers separately. + updateQuotaWithoutThrottleSync( + entry.getKey(), UserResourceQuotaConverter.fromThrift(entry.getValue())); + } + } + + public QuotaTokenBundle acquireReadResources( + String user, long memoryBytes, AcquireContext ctx, AcquirePolicy policy) + throws UserResourceQuotaExceededException { + AcquireResult cpu = acquire(user, OperationType.READ, ResourceType.CPU, 1, ctx, policy); + if (!cpu.isSuccess()) { + throw new UserResourceQuotaExceededException(cpu.getRejectReason()); + } + AcquireResult mem = + acquire(user, OperationType.READ, ResourceType.MEMORY, memoryBytes, ctx, policy); + if (!mem.isSuccess()) { + release(cpu.getToken()); + throw new UserResourceQuotaExceededException(mem.getRejectReason()); + } + AcquireResult tempDisk = + acquire(user, OperationType.READ, ResourceType.TEMP_DISK, memoryBytes, ctx, policy); + if (!tempDisk.isSuccess()) { + release(mem.getToken()); + release(cpu.getToken()); + throw new UserResourceQuotaExceededException(tempDisk.getRejectReason()); + } + return new QuotaTokenBundle(cpu.getToken(), mem.getToken(), tempDisk.getToken()); + } + + public QuotaTokenBundle acquireWriteResources( + String user, long memoryBytes, AcquireContext ctx, AcquirePolicy policy) + throws UserResourceQuotaExceededException { + AcquireResult cpu = acquire(user, OperationType.WRITE, ResourceType.CPU, 1, ctx, policy); + if (!cpu.isSuccess()) { + throw new UserResourceQuotaExceededException(cpu.getRejectReason()); + } + AcquireResult mem = + acquire(user, OperationType.WRITE, ResourceType.MEMORY, memoryBytes, ctx, policy); + if (!mem.isSuccess()) { + release(cpu.getToken()); + throw new UserResourceQuotaExceededException(mem.getRejectReason()); + } + AcquireResult tempDisk = + acquire(user, OperationType.WRITE, ResourceType.TEMP_DISK, memoryBytes, ctx, policy); + if (!tempDisk.isSuccess()) { + release(mem.getToken()); + release(cpu.getToken()); + throw new UserResourceQuotaExceededException(tempDisk.getRejectReason()); + } + return new QuotaTokenBundle(cpu.getToken(), mem.getToken(), tempDisk.getToken()); + } + + private boolean isEnabled() { + return IoTDBDescriptor.getInstance().getConfig().isQuotaEnable(); + } + + private boolean isExempt(String user) { + return IoTDBConstant.PATH_ROOT.equals(user); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/WriteMemoryEstimator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/WriteMemoryEstimator.java new file mode 100644 index 0000000000000..98570302e20fa --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/rescon/quotas/WriteMemoryEstimator.java @@ -0,0 +1,129 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.db.queryengine.plan.statement.Statement; +import org.apache.iotdb.db.queryengine.plan.statement.StatementType; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowStatement; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsOfOneDeviceStatement; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertRowsStatement; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertStatement; +import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertTabletStatement; +import org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement; +import org.apache.iotdb.db.queryengine.plan.statement.pipe.PipeEnrichedStatement; +import org.apache.iotdb.db.utils.TypeInferenceUtils; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.utils.BitMap; + +public final class WriteMemoryEstimator { + + private WriteMemoryEstimator() {} + + public static long estimate(Statement s) { + if (s.getType() == StatementType.PIPE_ENRICHED) { + s = ((PipeEnrichedStatement) s).getInnerStatement(); + } + switch (s.getType()) { + case INSERT: + if (s instanceof InsertStatement) { + long size = 0; + InsertStatement insertStatement = (InsertStatement) s; + for (Object[] values : insertStatement.getValuesList()) { + size += calculationWrite(values); + } + return size; + } + if (s instanceof InsertRowStatement) { + return calculationWrite(((InsertRowStatement) s).getValues()); + } + return 0; + case BATCH_INSERT: + return estimateTablet((InsertTabletStatement) s); + case BATCH_INSERT_ONE_DEVICE: + long oneDeviceSize = 0; + for (InsertRowStatement row : + ((InsertRowsOfOneDeviceStatement) s).getInsertRowStatementList()) { + oneDeviceSize += calculationWrite(row.getValues()); + } + return oneDeviceSize; + case BATCH_INSERT_ROWS: + long rowsSize = 0; + for (InsertRowStatement row : ((InsertRowsStatement) s).getInsertRowStatementList()) { + rowsSize += calculationWrite(row.getValues()); + } + return rowsSize; + case MULTI_BATCH_INSERT: + if (s instanceof LoadTsFileStatement) { + long loadSize = 0; + LoadTsFileStatement load = (LoadTsFileStatement) s; + for (int i = 0; i < load.getResources().size(); i++) { + loadSize += load.getResources().get(i).getTsFileSize(); + } + return loadSize; + } + if (s instanceof InsertMultiTabletsStatement) { + long tabletSize = 0; + InsertMultiTabletsStatement multi = (InsertMultiTabletsStatement) s; + for (InsertTabletStatement tablet : multi.getInsertTabletStatementList()) { + tabletSize += estimateTablet(tablet); + } + return tabletSize; + } + return 0; + default: + return 0; + } + } + + private static long calculationWrite(Object[] values) { + long size = 0; + for (Object value : values) { + TSDataType dataType = TypeInferenceUtils.getPredictedDataType(value, true); + if (dataType != null) { + size += dataType.getDataTypeSize(); + } + } + return size; + } + + private static long estimateTablet(InsertTabletStatement tablet) { + long size = 0; + int rowCount = tablet.getRowCount(); + TSDataType[] dataTypes = tablet.getDataTypes(); + if (dataTypes != null) { + for (TSDataType dataType : dataTypes) { + if (dataType != null) { + size += (long) dataType.getDataTypeSize() * rowCount; + } + } + } + BitMap[] bitMaps = tablet.getBitMaps(); + if (bitMaps != null) { + for (BitMap bitMap : bitMaps) { + if (bitMap != null) { + size += bitMap.getSize(); + } + } + } + return size; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/ShowUserResourceQuotaTaskTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/ShowUserResourceQuotaTaskTest.java new file mode 100644 index 0000000000000..3d6168e66e79f --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/execution/config/sys/quota/ShowUserResourceQuotaTaskTest.java @@ -0,0 +1,95 @@ +/* + * 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.iotdb.db.queryengine.plan.execution.config.sys.quota; + +import org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange; +import org.apache.iotdb.common.rpc.thrift.TResourceType; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceUsageSnapshot; +import org.apache.iotdb.confignode.rpc.thrift.TUserResourceQuotaResp; +import org.apache.iotdb.db.queryengine.plan.execution.config.ConfigTaskResult; +import org.apache.iotdb.rpc.TSStatusCode; + +import com.google.common.util.concurrent.SettableFuture; +import org.apache.tsfile.read.common.block.TsBlock; +import org.junit.Assert; +import org.junit.Test; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.TimeUnit; + +public class ShowUserResourceQuotaTaskTest { + + @Test + public void testBuildTsBlockExpandsPerAliveDataNode() throws Exception { + TUserResourceQuota quota = new TUserResourceQuota(); + Map read = new HashMap<>(); + read.put(TResourceType.CPU, new TResourceQuotaRange(1, 4)); + quota.setReadQuota(read); + + TUserResourceQuotaResp resp = new TUserResourceQuotaResp(); + Map quotas = new HashMap<>(); + quotas.put("u1", quota); + resp.setUserResourceQuota(quotas); + + Map usageByNode = new HashMap<>(); + TUserResourceUsageSnapshot snap1 = new TUserResourceUsageSnapshot(); + Map> readInUse1 = new HashMap<>(); + Map u1Cpu = new HashMap<>(); + u1Cpu.put(TResourceType.CPU, 2L); + readInUse1.put("u1", u1Cpu); + snap1.setReadInUse(readInUse1); + usageByNode.put(1, snap1); + // Running DN without usage report yet: empty snapshot still expands a NodeID row set. + usageByNode.put(2, new TUserResourceUsageSnapshot()); + resp.setUsageByDataNode(usageByNode); + + SettableFuture future = SettableFuture.create(); + ShowUserResourceQuotaTask.buildTSBlock(resp, future); + ConfigTaskResult result = future.get(5, TimeUnit.SECONDS); + Assert.assertEquals(TSStatusCode.SUCCESS_STATUS, result.getStatusCode()); + TsBlock block = result.getResultSet(); + Assert.assertNotNull(block); + + Set nodeIds = new HashSet<>(); + int usedForNode1 = -1; + int usedForNode2 = -1; + for (int i = 0; i < block.getPositionCount(); i++) { + String nodeId = block.getColumn(1).getBinary(i).toString(); + String type = block.getColumn(3).getBinary(i).toString(); + if ("cpu".equals(type)) { + nodeIds.add(nodeId); + int used = Integer.parseInt(block.getColumn(6).getBinary(i).toString()); + if ("1".equals(nodeId)) { + usedForNode1 = used; + } else if ("2".equals(nodeId)) { + usedForNode2 = used; + } + } + } + Assert.assertTrue(nodeIds.contains("1")); + Assert.assertTrue(nodeIds.contains("2")); + Assert.assertEquals(2, usedForNode1); + Assert.assertEquals(0, usedForNode2); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/UserResourceQuotaParseSmokeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/UserResourceQuotaParseSmokeTest.java new file mode 100644 index 0000000000000..6c56502372309 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/parser/UserResourceQuotaParseSmokeTest.java @@ -0,0 +1,168 @@ +/* + * 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.iotdb.db.queryengine.plan.parser; + +import org.apache.iotdb.common.rpc.thrift.TResourceType; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.common.rpc.thrift.ThrottleType; +import org.apache.iotdb.commons.exception.SemanticException; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.plan.statement.Statement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.DeleteUserResourceQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.SetUserResourceQuotaStatement; +import org.apache.iotdb.db.queryengine.plan.statement.sys.quota.ShowUserResourceQuotaStatement; + +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +import java.time.ZoneId; + +public class UserResourceQuotaParseSmokeTest { + + private boolean oldEnable; + + @Before + public void setUp() { + oldEnable = IoTDBDescriptor.getInstance().getConfig().isQuotaEnable(); + IoTDBDescriptor.getInstance().getConfig().setQuotaEnable(true); + } + + @After + public void tearDown() { + IoTDBDescriptor.getInstance().getConfig().setQuotaEnable(oldEnable); + } + + @Test + public void validUnitsParse() { + // disk_io fixed unit: bytes/sec long; memory still uses size suffix (e.g. 8G). + Statement s = + StatementGenerator.createStatement( + "SET USER QUOTA ON u1 WITH read_cpu_min=1, read_cpu_max=4, write_memory_max=8G, write_disk_io_max=10485760", + ZoneId.systemDefault()); + Assert.assertTrue(s instanceof SetUserResourceQuotaStatement); + TUserResourceQuota q = ((SetUserResourceQuotaStatement) s).getUserResourceQuota(); + Assert.assertEquals(1, q.getReadQuota().get(TResourceType.CPU).getMinValue()); + Assert.assertEquals(4, q.getReadQuota().get(TResourceType.CPU).getMaxValue()); + Assert.assertEquals( + 8L * 1024 * 1024 * 1024, q.getWriteQuota().get(TResourceType.MEMORY).getMaxValue()); + Assert.assertTrue(q.getThrottleLimit().containsKey(ThrottleType.WRITE_SIZE)); + Assert.assertEquals( + 10L * 1024 * 1024, q.getThrottleLimit().get(ThrottleType.WRITE_SIZE).getSoftLimit()); + } + + @Test + public void tempDiskAttrParses() { + // temp_disk fixed unit: bytes long. + Statement s = + StatementGenerator.createStatement( + "SET USER QUOTA ON u1 WITH read_temp_disk_max=10737418240, write_temp_disk_min=1073741824", + ZoneId.systemDefault()); + Assert.assertTrue(s instanceof SetUserResourceQuotaStatement); + TUserResourceQuota q = ((SetUserResourceQuotaStatement) s).getUserResourceQuota(); + Assert.assertEquals( + 10L * 1024 * 1024 * 1024, q.getReadQuota().get(TResourceType.TEMP_DISK).getMaxValue()); + Assert.assertEquals( + 1L * 1024 * 1024 * 1024, q.getWriteQuota().get(TResourceType.TEMP_DISK).getMinValue()); + } + + @Test + public void invalidGbSuffixRejected() { + try { + StatementGenerator.createStatement( + "SET USER QUOTA ON u1 WITH write_memory_max=8GB", ZoneId.systemDefault()); + Assert.fail("expected parse failure for 8GB"); + } catch (SemanticException | NumberFormatException e) { + // expected: unit is last char only (B/K/M/G/T/P), so 8GB is invalid + } + } + + @Test + public void showUserQuotaParses() { + Statement s = StatementGenerator.createStatement("SHOW USER QUOTA u1", ZoneId.systemDefault()); + Assert.assertTrue(s instanceof ShowUserResourceQuotaStatement); + ShowUserResourceQuotaStatement show = (ShowUserResourceQuotaStatement) s; + Assert.assertEquals("u1", show.getUserName()); + Assert.assertFalse(show.isSummary()); + Assert.assertNull(show.getDataNodeId()); + } + + @Test + public void deleteUserQuotaParses() { + Statement s = + StatementGenerator.createStatement("DELETE USER QUOTA ON u1", ZoneId.systemDefault()); + Assert.assertTrue(s instanceof DeleteUserResourceQuotaStatement); + Assert.assertEquals("u1", ((DeleteUserResourceQuotaStatement) s).getUserName()); + } + + @Test + public void quotaDisabledRejected() { + IoTDBDescriptor.getInstance().getConfig().setQuotaEnable(false); + try { + StatementGenerator.createStatement( + "SET USER QUOTA ON u1 WITH read_cpu_max=1", ZoneId.systemDefault()); + Assert.fail("expected quota_enable=false to reject SET"); + } catch (SemanticException e) { + Assert.assertTrue(e.getMessage().toLowerCase().contains("enable")); + } + } + + @Test + public void rootUserRejected() { + try { + StatementGenerator.createStatement( + "SET USER QUOTA ON `root` WITH read_cpu_max=1", ZoneId.systemDefault()); + Assert.fail("expected root SET to fail"); + } catch (SemanticException e) { + // expected + } + try { + StatementGenerator.createStatement("DELETE USER QUOTA ON `root`", ZoneId.systemDefault()); + Assert.fail("expected root DELETE to fail"); + } catch (SemanticException e) { + // expected + } + } + + @Test + public void invalidAttrAndDiskIoRejected() { + try { + StatementGenerator.createStatement( + "SET USER QUOTA ON u1 WITH read_cpu_xx=1", ZoneId.systemDefault()); + Assert.fail("expected invalid attr"); + } catch (RuntimeException e) { + // SemanticException / IllegalArgumentException + } + try { + StatementGenerator.createStatement( + "SET USER QUOTA ON u1 WITH write_disk_io_max='10M/sec'", ZoneId.systemDefault()); + Assert.fail("expected invalid disk_io complex unit"); + } catch (RuntimeException e) { + // SemanticException: disk_io must be positive long bytes/sec + } + try { + StatementGenerator.createStatement( + "SET USER QUOTA ON u1 WITH write_temp_disk_max=8G", ZoneId.systemDefault()); + Assert.fail("expected invalid temp_disk size suffix"); + } catch (RuntimeException e) { + // SemanticException: temp_disk must be positive long bytes + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/quotas/CapacityResourceLimiterTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/quotas/CapacityResourceLimiterTest.java new file mode 100644 index 0000000000000..ee966842047ed --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/rescon/quotas/CapacityResourceLimiterTest.java @@ -0,0 +1,310 @@ +/* + * 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.iotdb.db.storageengine.rescon.quotas; + +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.commons.quota.OperationType; +import org.apache.iotdb.commons.quota.ResourceQuotaRange; +import org.apache.iotdb.commons.quota.ResourceType; +import org.apache.iotdb.commons.quota.UserResourceQuota; +import org.apache.iotdb.db.conf.IoTDBConfig; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.i18n.StorageEngineMessages; + +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +public class CapacityResourceLimiterTest { + + @Before + public void setUp() { + IoTDBDescriptor.getInstance().getConfig().setQuotaEnable(true); + } + + @After + public void tearDown() { + IoTDBDescriptor.getInstance().getConfig().setQuotaEnable(false); + } + + @Test + public void testMaxAndMinScheduling() { + CapacityResourceLimiter limiter = new CapacityResourceLimiter(); + NodeQuotaState node = new NodeQuotaState(1, 4); + node.updateRange("u1", new ResourceQuotaRange(1, 3)); + node.updateRange("u2", new ResourceQuotaRange(1, 2)); + + Assert.assertTrue( + limiter.tryAcquire("u1", 2, node.getRangeByUser().get("u1"), node).isSuccess()); + Assert.assertTrue( + limiter.tryAcquire("u2", 2, node.getRangeByUser().get("u2"), node).isSuccess()); + LimiterAcquireResult denied = + limiter.tryAcquire("u1", 1, node.getRangeByUser().get("u1"), node); + Assert.assertFalse(denied.isSuccess()); + Assert.assertEquals( + StorageEngineMessages.EXCEPTION_NODE_CAPACITY_EXCEEDED_89601D9A, denied.getRejectReason()); + + limiter.release("u2", 1, node); + Assert.assertTrue( + limiter.tryAcquire("u1", 1, node.getRangeByUser().get("u1"), node).isSuccess()); + } + + @Test + public void testRejectReasonUserMax() { + CapacityResourceLimiter limiter = new CapacityResourceLimiter(); + NodeQuotaState node = new NodeQuotaState(1, 100); + node.updateRange("u1", new ResourceQuotaRange(0, 2)); + Assert.assertTrue( + limiter.tryAcquire("u1", 2, node.getRangeByUser().get("u1"), node).isSuccess()); + LimiterAcquireResult denied = + limiter.tryAcquire("u1", 1, node.getRangeByUser().get("u1"), node); + Assert.assertFalse(denied.isSuccess()); + Assert.assertEquals( + StorageEngineMessages.EXCEPTION_USER_MAX_EXCEEDED_3D400C08, denied.getRejectReason()); + } + + @Test + public void testRejectReasonMinGap() { + CapacityResourceLimiter limiter = new CapacityResourceLimiter(); + NodeQuotaState node = new NodeQuotaState(1, 4); + node.updateRange("core", new ResourceQuotaRange(3, 4)); + LimiterAcquireResult denied = limiter.tryAcquire("free", 2, null, node); + Assert.assertFalse(denied.isSuccess()); + Assert.assertEquals( + StorageEngineMessages.EXCEPTION_MIN_GAP_RESERVATION_B84C4EE4, denied.getRejectReason()); + } + + @Test + public void testManagerAcquireRelease() { + UserResourceQuotaManager manager = UserResourceQuotaManager.getInstance(); + UserResourceQuota quota = new UserResourceQuota(); + quota.getReadQuota().put(ResourceType.CPU, new ResourceQuotaRange(0, 2)); + manager.updateQuota("quotaTestUser", quota); + + AcquireContext ctx = new AcquireContext().setStatementType("QUERY"); + AcquirePolicy policy = AcquirePolicy.defaults(); + policy.setMaxWaitMs(0); + QuotaToken token = + manager + .acquire("quotaTestUser", OperationType.READ, ResourceType.CPU, 1, ctx, policy) + .getToken(); + Assert.assertEquals(1, manager.getInUse("quotaTestUser", OperationType.READ, ResourceType.CPU)); + token.close(); + Assert.assertEquals(0, manager.getInUse("quotaTestUser", OperationType.READ, ResourceType.CPU)); + } + + @Test + public void testManagerRejectIncludesConcreteReason() { + UserResourceQuotaManager manager = UserResourceQuotaManager.getInstance(); + UserResourceQuota quota = new UserResourceQuota(); + quota.getReadQuota().put(ResourceType.CPU, new ResourceQuotaRange(0, 1)); + manager.updateQuota("rejectReasonUser", quota); + + AcquirePolicy policy = AcquirePolicy.defaults(); + policy.setMaxWaitMs(0); + Assert.assertTrue( + manager + .acquire( + "rejectReasonUser", + OperationType.READ, + ResourceType.CPU, + 1, + new AcquireContext(), + policy) + .isSuccess()); + AcquireResult rejected = + manager.acquire( + "rejectReasonUser", + OperationType.READ, + ResourceType.CPU, + 1, + new AcquireContext(), + policy); + Assert.assertFalse(rejected.isSuccess()); + Assert.assertTrue( + rejected + .getRejectReason() + .contains(StorageEngineMessages.EXCEPTION_USER_MAX_EXCEEDED_3D400C08)); + } + + @Test + public void testUnlimitedUserPassesWhenCapacityAllows() { + UserResourceQuotaManager manager = UserResourceQuotaManager.getInstance(); + AcquirePolicy policy = AcquirePolicy.defaults(); + policy.setMaxWaitMs(0); + QuotaToken token = + manager + .acquire( + "noQuotaUser", + OperationType.READ, + ResourceType.CPU, + 1, + new AcquireContext(), + policy) + .getToken(); + Assert.assertNotNull(token); + Assert.assertTrue(manager.getInUse("noQuotaUser", OperationType.READ, ResourceType.CPU) >= 1); + token.close(); + } + + @Test + public void testUnconfiguredUserRespectsMinGap() { + CapacityResourceLimiter limiter = new CapacityResourceLimiter(); + NodeQuotaState node = new NodeQuotaState(1, 4); + node.updateRange("core", new ResourceQuotaRange(3, 4)); + Assert.assertFalse(limiter.tryAcquire("free", 2, null, node).isSuccess()); + Assert.assertTrue( + limiter.tryAcquire("core", 3, node.getRangeByUser().get("core"), node).isSuccess()); + Assert.assertTrue(limiter.tryAcquire("free", 1, null, node).isSuccess()); + } + + @Test + public void testLowerMaxDoesNotKillExistingInUse() { + CapacityResourceLimiter limiter = new CapacityResourceLimiter(); + NodeQuotaState node = new NodeQuotaState(1, 10); + node.updateRange("u1", new ResourceQuotaRange(0, 8)); + Assert.assertTrue( + limiter.tryAcquire("u1", 6, node.getRangeByUser().get("u1"), node).isSuccess()); + node.updateRange("u1", new ResourceQuotaRange(0, 4)); + Assert.assertEquals(6, node.inUse("u1")); + LimiterAcquireResult denied = + limiter.tryAcquire("u1", 1, node.getRangeByUser().get("u1"), node); + Assert.assertFalse(denied.isSuccess()); + Assert.assertEquals( + StorageEngineMessages.EXCEPTION_USER_MAX_EXCEEDED_3D400C08, denied.getRejectReason()); + } + + @Test + public void testClearUserQuotaKeepsInUseForInFlightToken() { + UserResourceQuotaManager manager = UserResourceQuotaManager.getInstance(); + UserResourceQuota quota = new UserResourceQuota(); + quota.getReadQuota().put(ResourceType.CPU, new ResourceQuotaRange(0, 4)); + manager.updateQuota("clearKeepUser", quota); + + AcquirePolicy policy = AcquirePolicy.defaults(); + policy.setMaxWaitMs(0); + QuotaToken token = + manager + .acquire( + "clearKeepUser", + OperationType.READ, + ResourceType.CPU, + 2, + new AcquireContext(), + policy) + .getToken(); + Assert.assertEquals(2, manager.getInUse("clearKeepUser", OperationType.READ, ResourceType.CPU)); + + manager.clearUserQuota("clearKeepUser"); + Assert.assertNull(manager.getUserQuota("clearKeepUser")); + Assert.assertEquals(2, manager.getInUse("clearKeepUser", OperationType.READ, ResourceType.CPU)); + + token.close(); + Assert.assertEquals(0, manager.getInUse("clearKeepUser", OperationType.READ, ResourceType.CPU)); + } + + @Test + public void testIsClearRequest() { + Assert.assertTrue(UserResourceQuotaManager.isClearRequest(null)); + Assert.assertTrue(UserResourceQuotaManager.isClearRequest(new TUserResourceQuota())); + TUserResourceQuota quota = new TUserResourceQuota(); + quota.putToReadQuota( + org.apache.iotdb.common.rpc.thrift.TResourceType.CPU, + new org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange(0, 1)); + Assert.assertFalse(UserResourceQuotaManager.isClearRequest(quota)); + } + + @Test + public void testDefaultTempDiskBytesIsMinOfHundredGiBAndTenthDisk() { + IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + long computed = config.computeDefaultTempDiskBytes(); + long hundredGiB = 100L * 1024 * 1024 * 1024; + Assert.assertTrue(computed > 0); + Assert.assertTrue(computed <= hundredGiB); + long totalSpace = 0L; + String[] dirs = config.getDataDirs(); + if (dirs != null) { + for (String dir : dirs) { + if (dir != null) { + long space = new java.io.File(dir).getTotalSpace(); + if (space > 0) { + totalSpace += space; + } + } + } + } + if (totalSpace > 0) { + Assert.assertEquals(Math.min(hundredGiB, totalSpace / 10), computed); + } else { + Assert.assertEquals(hundredGiB, computed); + } + } + + @Test + public void testRootUserExempt() { + UserResourceQuotaManager manager = UserResourceQuotaManager.getInstance(); + UserResourceQuota quota = new UserResourceQuota(); + quota.getReadQuota().put(ResourceType.CPU, new ResourceQuotaRange(0, 1)); + manager.updateQuota(IoTDBConstant.PATH_ROOT, quota); + AcquirePolicy policy = AcquirePolicy.defaults(); + policy.setMaxWaitMs(0); + Assert.assertTrue( + manager + .acquire( + IoTDBConstant.PATH_ROOT, + OperationType.READ, + ResourceType.CPU, + 1, + new AcquireContext(), + policy) + .isSuccess()); + } + + @Test + public void testWriteTempDiskMaxRejectDoesNotAffectOtherUser() { + UserResourceQuotaManager manager = UserResourceQuotaManager.getInstance(); + UserResourceQuota q1 = new UserResourceQuota(); + q1.getWriteQuota().put(ResourceType.TEMP_DISK, new ResourceQuotaRange(0, 1)); + manager.updateQuota("td_u1", q1); + UserResourceQuota q2 = new UserResourceQuota(); + q2.getWriteQuota().put(ResourceType.TEMP_DISK, new ResourceQuotaRange(0, 10)); + manager.updateQuota("td_u2", q2); + + AcquirePolicy policy = AcquirePolicy.defaults(); + policy.setMaxWaitMs(0); + Assert.assertFalse( + manager + .acquire( + "td_u1", + OperationType.WRITE, + ResourceType.TEMP_DISK, + 2, + new AcquireContext(), + policy) + .isSuccess()); + AcquireResult ok = + manager.acquire( + "td_u2", OperationType.WRITE, ResourceType.TEMP_DISK, 2, new AcquireContext(), policy); + Assert.assertTrue(ok.isSuccess()); + ok.getToken().close(); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/OperationType.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/OperationType.java new file mode 100644 index 0000000000000..18ad33cfd6f1e --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/OperationType.java @@ -0,0 +1,25 @@ +/* + * 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.iotdb.commons.quota; + +public enum OperationType { + READ, + WRITE +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/ResourceQuotaRange.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/ResourceQuotaRange.java new file mode 100644 index 0000000000000..231fca25ff538 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/ResourceQuotaRange.java @@ -0,0 +1,74 @@ +/* + * 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.iotdb.commons.quota; + +import org.apache.iotdb.commons.conf.IoTDBConstant; + +import java.util.Objects; + +public class ResourceQuotaRange { + + private long minValue = IoTDBConstant.UNLIMITED_VALUE; + private long maxValue = IoTDBConstant.UNLIMITED_VALUE; + + public ResourceQuotaRange() {} + + public ResourceQuotaRange(long minValue, long maxValue) { + this.minValue = minValue; + this.maxValue = maxValue; + } + + public long getMinValue() { + return minValue; + } + + public void setMinValue(long minValue) { + this.minValue = minValue; + } + + public long getMaxValue() { + return maxValue; + } + + public void setMaxValue(long maxValue) { + this.maxValue = maxValue; + } + + public boolean isUnlimited() { + return minValue == IoTDBConstant.UNLIMITED_VALUE && maxValue == IoTDBConstant.UNLIMITED_VALUE; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + ResourceQuotaRange that = (ResourceQuotaRange) o; + return minValue == that.minValue && maxValue == that.maxValue; + } + + @Override + public int hashCode() { + return Objects.hash(minValue, maxValue); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/ResourceType.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/ResourceType.java new file mode 100644 index 0000000000000..474247ddd9af0 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/ResourceType.java @@ -0,0 +1,27 @@ +/* + * 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.iotdb.commons.quota; + +public enum ResourceType { + CPU, + MEMORY, + DISK_IO, + TEMP_DISK +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/UserResourceQuota.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/UserResourceQuota.java new file mode 100644 index 0000000000000..1a71aee49d881 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/UserResourceQuota.java @@ -0,0 +1,92 @@ +/* + * 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.iotdb.commons.quota; + +import org.apache.iotdb.common.rpc.thrift.TTimedQuota; +import org.apache.iotdb.common.rpc.thrift.ThrottleType; + +import java.util.EnumMap; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +public class UserResourceQuota { + + private final Map readQuota = new EnumMap<>(ResourceType.class); + private final Map writeQuota = + new EnumMap<>(ResourceType.class); + private final Map throttleLimit = new HashMap<>(); + + public Map getReadQuota() { + return readQuota; + } + + public Map getWriteQuota() { + return writeQuota; + } + + public Map getThrottleLimit() { + return throttleLimit; + } + + public ResourceQuotaRange getRange(OperationType op, ResourceType resource) { + Map quota = op == OperationType.READ ? readQuota : writeQuota; + return quota.get(resource); + } + + public void mergeFrom(UserResourceQuota other) { + if (other == null) { + return; + } + for (Map.Entry entry : other.readQuota.entrySet()) { + readQuota.put(entry.getKey(), copyRange(entry.getValue())); + } + for (Map.Entry entry : other.writeQuota.entrySet()) { + writeQuota.put(entry.getKey(), copyRange(entry.getValue())); + } + throttleLimit.putAll(other.throttleLimit); + } + + private static ResourceQuotaRange copyRange(ResourceQuotaRange range) { + if (range == null) { + return null; + } + return new ResourceQuotaRange(range.getMinValue(), range.getMaxValue()); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + UserResourceQuota that = (UserResourceQuota) o; + return Objects.equals(readQuota, that.readQuota) + && Objects.equals(writeQuota, that.writeQuota) + && Objects.equals(throttleLimit, that.throttleLimit); + } + + @Override + public int hashCode() { + return Objects.hash(readQuota, writeQuota, throttleLimit); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/UserResourceQuotaConverter.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/UserResourceQuotaConverter.java new file mode 100644 index 0000000000000..33dda8aa762d8 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/quota/UserResourceQuotaConverter.java @@ -0,0 +1,155 @@ +/* + * 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.iotdb.commons.quota; + +import org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange; +import org.apache.iotdb.common.rpc.thrift.TResourceType; +import org.apache.iotdb.common.rpc.thrift.TThrottleQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.commons.conf.IoTDBConstant; + +import java.util.EnumMap; +import java.util.HashMap; +import java.util.Map; + +public final class UserResourceQuotaConverter { + + private UserResourceQuotaConverter() {} + + public static UserResourceQuota fromThrift(TUserResourceQuota thrift) { + UserResourceQuota quota = new UserResourceQuota(); + if (thrift == null) { + return quota; + } + if (thrift.isSetReadQuota()) { + for (Map.Entry entry : thrift.getReadQuota().entrySet()) { + quota.getReadQuota().put(toResourceType(entry.getKey()), fromRange(entry.getValue())); + } + } + if (thrift.isSetWriteQuota()) { + for (Map.Entry entry : + thrift.getWriteQuota().entrySet()) { + quota.getWriteQuota().put(toResourceType(entry.getKey()), fromRange(entry.getValue())); + } + } + if (thrift.isSetThrottleLimit()) { + quota.getThrottleLimit().putAll(thrift.getThrottleLimit()); + } + return quota; + } + + public static TUserResourceQuota toThrift(UserResourceQuota quota) { + TUserResourceQuota thrift = new TUserResourceQuota(); + if (quota == null) { + return thrift; + } + if (!quota.getReadQuota().isEmpty()) { + Map readQuota = new EnumMap<>(TResourceType.class); + for (Map.Entry entry : quota.getReadQuota().entrySet()) { + readQuota.put(toThriftType(entry.getKey()), toRange(entry.getValue())); + } + thrift.setReadQuota(readQuota); + } + if (!quota.getWriteQuota().isEmpty()) { + Map writeQuota = new EnumMap<>(TResourceType.class); + for (Map.Entry entry : quota.getWriteQuota().entrySet()) { + writeQuota.put(toThriftType(entry.getKey()), toRange(entry.getValue())); + } + thrift.setWriteQuota(writeQuota); + } + if (!quota.getThrottleLimit().isEmpty()) { + thrift.setThrottleLimit(new HashMap<>(quota.getThrottleLimit())); + } + return thrift; + } + + public static UserResourceQuota fromThrottleQuota(TThrottleQuota throttleQuota) { + UserResourceQuota quota = new UserResourceQuota(); + if (throttleQuota == null) { + return quota; + } + if (throttleQuota.isSetCpuLimit() && throttleQuota.getCpuLimit() > 0) { + quota + .getReadQuota() + .put( + ResourceType.CPU, + new ResourceQuotaRange(IoTDBConstant.UNLIMITED_VALUE, throttleQuota.getCpuLimit())); + } + if (throttleQuota.isSetMemLimit() && throttleQuota.getMemLimit() > 0) { + quota + .getReadQuota() + .put( + ResourceType.MEMORY, + new ResourceQuotaRange(IoTDBConstant.UNLIMITED_VALUE, throttleQuota.getMemLimit())); + } + if (throttleQuota.isSetThrottleLimit()) { + quota.getThrottleLimit().putAll(throttleQuota.getThrottleLimit()); + } + return quota; + } + + public static TThrottleQuota toThrottleQuota(UserResourceQuota quota) { + TThrottleQuota thrift = new TThrottleQuota(); + if (quota == null) { + return thrift; + } + ResourceQuotaRange readCpu = quota.getReadQuota().get(ResourceType.CPU); + if (readCpu != null && readCpu.getMaxValue() > 0) { + thrift.setCpuLimit((int) readCpu.getMaxValue()); + } + ResourceQuotaRange readMem = quota.getReadQuota().get(ResourceType.MEMORY); + if (readMem != null && readMem.getMaxValue() > 0) { + thrift.setMemLimit(readMem.getMaxValue()); + } + if (!quota.getThrottleLimit().isEmpty()) { + thrift.setThrottleLimit(new HashMap<>(quota.getThrottleLimit())); + } else { + thrift.setThrottleLimit(new HashMap<>()); + } + return thrift; + } + + private static ResourceQuotaRange fromRange(TResourceQuotaRange range) { + if (range == null) { + return new ResourceQuotaRange(); + } + return new ResourceQuotaRange(range.getMinValue(), range.getMaxValue()); + } + + private static TResourceQuotaRange toRange(ResourceQuotaRange range) { + TResourceQuotaRange thrift = new TResourceQuotaRange(); + if (range == null) { + thrift.setMinValue(IoTDBConstant.UNLIMITED_VALUE); + thrift.setMaxValue(IoTDBConstant.UNLIMITED_VALUE); + return thrift; + } + thrift.setMinValue(range.getMinValue()); + thrift.setMaxValue(range.getMaxValue()); + return thrift; + } + + private static ResourceType toResourceType(TResourceType type) { + return ResourceType.valueOf(type.name()); + } + + private static TResourceType toThriftType(ResourceType type) { + return TResourceType.valueOf(type.name()); + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java index 18653a52c5708..aa92a31ec1015 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/schema/column/ColumnHeaderConstant.java @@ -316,7 +316,10 @@ private ColumnHeaderConstant() { // column names for show space quota public static final String QUOTA_TYPE = "QuotaType"; public static final String LIMIT = "Limit"; + public static final String MIN_LIMIT = "Min"; + public static final String MAX_LIMIT = "Max"; public static final String USED = "Used"; + public static final String MIN_GAP = "MinGap"; // column names for show throttle quota public static final String USER = "User"; @@ -709,6 +712,17 @@ private ColumnHeaderConstant() { new ColumnHeader(LIMIT, TSDataType.TEXT), new ColumnHeader(READ_WRITE, TSDataType.TEXT)); + public static final List showUserResourceQuotaColumnHeaders = + ImmutableList.of( + new ColumnHeader(USER, TSDataType.TEXT), + new ColumnHeader(NODE_ID, TSDataType.TEXT), + new ColumnHeader(READ_WRITE, TSDataType.TEXT), + new ColumnHeader(QUOTA_TYPE, TSDataType.TEXT), + new ColumnHeader(MIN_LIMIT, TSDataType.TEXT), + new ColumnHeader(MAX_LIMIT, TSDataType.TEXT), + new ColumnHeader(USED, TSDataType.TEXT), + new ColumnHeader(MIN_GAP, TSDataType.TEXT)); + public static final List showModelsColumnHeaders = ImmutableList.of( new ColumnHeader(MODEL_ID, TSDataType.TEXT), diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/quota/UserResourceQuotaConverterTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/quota/UserResourceQuotaConverterTest.java new file mode 100644 index 0000000000000..0431f241e4063 --- /dev/null +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/quota/UserResourceQuotaConverterTest.java @@ -0,0 +1,74 @@ +/* + * 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.iotdb.commons.quota; + +import org.apache.iotdb.common.rpc.thrift.TResourceQuotaRange; +import org.apache.iotdb.common.rpc.thrift.TResourceType; +import org.apache.iotdb.common.rpc.thrift.TThrottleQuota; +import org.apache.iotdb.common.rpc.thrift.TUserResourceQuota; +import org.apache.iotdb.commons.conf.IoTDBConstant; + +import org.junit.Assert; +import org.junit.Test; + +import java.util.EnumMap; +import java.util.Map; + +public class UserResourceQuotaConverterTest { + + @Test + public void testThrottleMigration() { + TThrottleQuota throttle = new TThrottleQuota(); + throttle.setCpuLimit(4); + throttle.setMemLimit(1024); + UserResourceQuota quota = UserResourceQuotaConverter.fromThrottleQuota(throttle); + Assert.assertEquals(4, quota.getReadQuota().get(ResourceType.CPU).getMaxValue()); + Assert.assertEquals(1024, quota.getReadQuota().get(ResourceType.MEMORY).getMaxValue()); + Assert.assertEquals(4, UserResourceQuotaConverter.toThrottleQuota(quota).getCpuLimit()); + } + + @Test + public void testThriftRoundTrip() { + TUserResourceQuota thrift = new TUserResourceQuota(); + Map read = new EnumMap<>(TResourceType.class); + read.put(TResourceType.CPU, new TResourceQuotaRange(1, 8)); + read.put(TResourceType.MEMORY, new TResourceQuotaRange(IoTDBConstant.UNLIMITED_VALUE, 2048)); + thrift.setReadQuota(read); + UserResourceQuota quota = UserResourceQuotaConverter.fromThrift(thrift); + TUserResourceQuota roundTrip = UserResourceQuotaConverter.toThrift(quota); + Assert.assertEquals(1, roundTrip.getReadQuota().get(TResourceType.CPU).getMinValue()); + Assert.assertEquals(8, roundTrip.getReadQuota().get(TResourceType.CPU).getMaxValue()); + } + + @Test + public void testPartialMergeKeepsUnmentionedDims() { + UserResourceQuota base = new UserResourceQuota(); + base.getReadQuota().put(ResourceType.CPU, new ResourceQuotaRange(1, 4)); + base.getWriteQuota().put(ResourceType.MEMORY, new ResourceQuotaRange(0, 1024)); + + UserResourceQuota patch = new UserResourceQuota(); + patch.getReadQuota().put(ResourceType.CPU, new ResourceQuotaRange(2, 8)); + + base.mergeFrom(patch); + Assert.assertEquals(2, base.getReadQuota().get(ResourceType.CPU).getMinValue()); + Assert.assertEquals(8, base.getReadQuota().get(ResourceType.CPU).getMaxValue()); + Assert.assertEquals(1024, base.getWriteQuota().get(ResourceType.MEMORY).getMaxValue()); + } +} diff --git a/iotdb-protocol/thrift-commons/src/main/thrift/common.thrift b/iotdb-protocol/thrift-commons/src/main/thrift/common.thrift index d52dc08813fbd..ff5fb34721473 100644 --- a/iotdb-protocol/thrift-commons/src/main/thrift/common.thrift +++ b/iotdb-protocol/thrift-commons/src/main/thrift/common.thrift @@ -197,6 +197,41 @@ struct TSetThrottleQuotaReq { 2: required TThrottleQuota throttleQuota } +enum TResourceType { + CPU, + MEMORY, + DISK_IO, + TEMP_DISK +} + +enum TOperationType { + READ, + WRITE +} + +struct TResourceQuotaRange { + 1: required i64 minValue + 2: required i64 maxValue +} + +struct TUserResourceQuota { + 1: optional map readQuota + 2: optional map writeQuota + 3: optional map throttleLimit +} + +// Per-DataNode in-use snapshot for SHOW USER QUOTA (heartbeat reported). +struct TUserResourceUsageSnapshot { + // username -> resourceType -> inUse + 1: optional map> readInUse + 2: optional map> writeInUse +} + +struct TSetUserResourceQuotaReq { + 1: required string userName + 2: required TUserResourceQuota userResourceQuota +} + struct TPipeHeartbeatResp { 1: required list pipeMetaList 2: optional list pipeCompletedList diff --git a/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift b/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift index a414584f80ec9..787928996ad7c 100644 --- a/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift +++ b/iotdb-protocol/thrift-confignode/src/main/thrift/confignode.thrift @@ -1176,6 +1176,19 @@ struct TShowThrottleReq { 1: optional string userName; } +struct TUserResourceQuotaResp { + 1: required common.TSStatus status + 2: optional map userResourceQuota + // dataNodeId -> in-use snapshot (from heartbeat); only alive DataNodes + 3: optional map usageByDataNode +} + +struct TShowUserResourceQuotaReq { + 1: optional string userName + 2: optional i32 dataNodeId + 3: optional bool summary +} + // ==================================================== // Activation @@ -2076,6 +2089,21 @@ service IConfigNodeRPCService { /** Get throttle quota information */ TThrottleQuotaResp getThrottleQuota() + /** Set user resource quota */ + common.TSStatus setUserResourceQuota(common.TSetUserResourceQuotaReq req) + + /** Show user resource quota */ + TUserResourceQuotaResp showUserResourceQuota(TShowUserResourceQuotaReq req) + + /** Get user resource quota */ + TUserResourceQuotaResp getUserResourceQuota() + + /** + * DataNode reports per-user resource in-use snapshot for SHOW USER QUOTA aggregation. + * Dedicated RPC (not main node heartbeat); low frequency is enough. + */ + common.TSStatus reportUserResourceUsage(i32 dataNodeId, common.TUserResourceUsageSnapshot usage) + /** Push heartbeat in shutdown */ common.TSStatus pushHeartbeat(i32 dataNodeId, common.TPipeHeartbeatResp resp) diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index 054b1628c6af0..17eb5a13068fa 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -1308,6 +1308,8 @@ service IDataNodeRPCService { **/ common.TSStatus setThrottleQuota(common.TSetThrottleQuotaReq req) + common.TSStatus setUserResourceQuota(common.TSetUserResourceQuotaReq req) + /** * Fetch fragment instance statistics for EXPLAIN ANALYZE */