From ac13213fe6630d7b834d7fc41672d6e16acc66d1 Mon Sep 17 00:00:00 2001 From: Ray Mattingly Date: Tue, 19 Sep 2023 19:11:55 -0400 Subject: [PATCH 1/5] Support quota user overrides --- .../hadoop/hbase/quotas/QuotaCache.java | 31 ++++- .../hbase/quotas/TestQuotaUserOverride.java | 111 +++++++++++++++++ .../TestQuotaUserOverrideConfiguration.java | 112 ++++++++++++++++++ 3 files changed, 253 insertions(+), 1 deletion(-) create mode 100644 hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java create mode 100644 hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java index 56253f7fcbb2..eac70fce83a3 100644 --- a/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java +++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java @@ -24,6 +24,7 @@ import java.util.EnumSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -35,8 +36,11 @@ import org.apache.hadoop.hbase.TableName; import org.apache.hadoop.hbase.client.Get; import org.apache.hadoop.hbase.client.RegionStatesCount; +import org.apache.hadoop.hbase.ipc.RpcCall; +import org.apache.hadoop.hbase.ipc.RpcServer; import org.apache.hadoop.hbase.regionserver.HRegionServer; import org.apache.hadoop.hbase.regionserver.RegionServerServices; +import org.apache.hadoop.hbase.util.Bytes; import org.apache.hadoop.hbase.util.EnvironmentEdgeManager; import org.apache.hadoop.security.UserGroupInformation; import org.apache.yetus.audience.InterfaceAudience; @@ -57,6 +61,11 @@ public class QuotaCache implements Stoppable { private static final Logger LOG = LoggerFactory.getLogger(QuotaCache.class); public static final String REFRESH_CONF_KEY = "hbase.quota.refresh.period"; + + // defines the request attribute key which, when provided, will override the request's username + // from the perspective of user quotas. This is also the default request attribute key + public static final String QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY = + "hbase.quota.user.override.key"; private static final int REFRESH_DEFAULT_PERIOD = 5 * 60000; // 5min private static final int EVICT_PERIOD_FACTOR = 5; // N * REFRESH_DEFAULT_PERIOD @@ -74,12 +83,15 @@ public class QuotaCache implements Stoppable { private final ConcurrentHashMap tableMachineQuotaFactors = new ConcurrentHashMap<>(); private final RegionServerServices rsServices; + private final String userOverrideRequestAttributeKey; private QuotaRefresherChore refreshChore; private boolean stopped = true; public QuotaCache(final RegionServerServices rsServices) { this.rsServices = rsServices; + this.userOverrideRequestAttributeKey = rsServices.getConfiguration() + .get(QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY); } public void start() throws IOException { @@ -125,7 +137,7 @@ public QuotaLimiter getUserLimiter(final UserGroupInformation ugi, final TableNa * @return the quota info associated to specified user */ public UserQuotaState getUserQuotaState(final UserGroupInformation ugi) { - return computeIfAbsent(userQuotaCache, ugi.getShortUserName(), UserQuotaState::new, + return computeIfAbsent(userQuotaCache, getQuotaUserName(ugi), UserQuotaState::new, this::triggerCacheRefresh); } @@ -160,6 +172,23 @@ protected boolean isExceedThrottleQuotaEnabled() { return exceedThrottleQuotaEnabled; } + /** + * Applies a request attribute user override if available, otherwise returns the UGI's short + * username + * @param ugi The request's UserGroupInformation + */ + private String getQuotaUserName(final UserGroupInformation ugi) { + Optional rpcCall = RpcServer.getCurrentCall(); + if ( + rpcCall.isPresent() + && rpcCall.get().getRequestAttributes().containsKey(userOverrideRequestAttributeKey) + ) { + return Bytes + .toString(rpcCall.get().getRequestAttributes().get(userOverrideRequestAttributeKey)); + } + return ugi.getShortUserName(); + } + /** * Returns the QuotaState requested. If the quota info is not in cache an empty one will be * returned and the quota request will be enqueued for the next cache refresh. diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java new file mode 100644 index 000000000000..53e25a2aec9d --- /dev/null +++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.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.hadoop.hbase.quotas; + +import static org.apache.hadoop.hbase.quotas.ThrottleQuotaTestUtil.doPuts; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.concurrent.TimeUnit; +import org.apache.hadoop.hbase.HBaseClassTestRule; +import org.apache.hadoop.hbase.HBaseTestingUtil; +import org.apache.hadoop.hbase.HConstants; +import org.apache.hadoop.hbase.TableName; +import org.apache.hadoop.hbase.client.Admin; +import org.apache.hadoop.hbase.client.Table; +import org.apache.hadoop.hbase.security.User; +import org.apache.hadoop.hbase.testclassification.MediumTests; +import org.apache.hadoop.hbase.testclassification.RegionServerTests; +import org.apache.hadoop.hbase.util.Bytes; +import org.apache.hadoop.hbase.util.EnvironmentEdgeManager; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +@Category({ RegionServerTests.class, MediumTests.class }) +public class TestQuotaUserOverride { + + @ClassRule + public static final HBaseClassTestRule CLASS_RULE = + HBaseClassTestRule.forClass(TestQuotaUserOverride.class); + + private static final HBaseTestingUtil TEST_UTIL = new HBaseTestingUtil(); + private static final byte[] FAMILY = Bytes.toBytes("cf"); + private static final byte[] QUALIFIER = Bytes.toBytes("q"); + private static final int NUM_SERVERS = 1; + + private static final TableName TABLE_NAME = TableName.valueOf("TestQuotaUserOverride"); + + @BeforeClass + public static void setUpBeforeClass() throws Exception { + TEST_UTIL.getConfiguration().setBoolean(QuotaUtil.QUOTA_CONF_KEY, true); + TEST_UTIL.getConfiguration().setInt(QuotaCache.REFRESH_CONF_KEY, 1_000); + TEST_UTIL.getConfiguration().setInt("hbase.client.pause", 1); + TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, 0); + TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_OPERATION_TIMEOUT, 500); + TEST_UTIL.startMiniCluster(NUM_SERVERS); + TEST_UTIL.waitTableAvailable(QuotaTableUtil.QUOTA_TABLE_NAME); + QuotaCache.TEST_FORCE_REFRESH = true; + TEST_UTIL.createTable(TABLE_NAME, FAMILY); + } + + @AfterClass + public static void tearDownAfterClass() throws Exception { + EnvironmentEdgeManager.reset(); + TEST_UTIL.shutdownMiniCluster(); + } + + @Test + public void testUserGlobalThrottleWithOverride() throws Exception { + final Admin admin = TEST_UTIL.getAdmin(); + final String userOverrideWithQuota = User.getCurrent().getShortName() + "123"; + final String userOverrideWithoutQuota = User.getCurrent().getShortName() + "456"; + + // Add 6req/min limit + admin.setQuota(QuotaSettingsFactory.throttleUser(userOverrideWithQuota, + ThrottleType.REQUEST_NUMBER, 6, TimeUnit.MINUTES)); + Thread.sleep(60_000); + + Table tableWithThrottle = admin.getConnection().getTableBuilder(TABLE_NAME, null) + .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, + Bytes.toBytes(userOverrideWithQuota)) + .build(); + Table tableWithoutThrottle = admin.getConnection().getTableBuilder(TABLE_NAME, null) + .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, + Bytes.toBytes(userOverrideWithoutQuota)) + .build(); + + // warm things up + doPuts(10, FAMILY, QUALIFIER, tableWithThrottle); + doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle); + + // should reject some requests + assertTrue(10 > doPuts(10, FAMILY, QUALIFIER, tableWithThrottle)); + // should accept all puts + assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle)); + + // Remove all the limits + admin.setQuota(QuotaSettingsFactory.unthrottleUser(userOverrideWithQuota)); + Thread.sleep(60_000); + assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithThrottle)); + assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle)); + } + +} diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java new file mode 100644 index 000000000000..1eac049677e5 --- /dev/null +++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java @@ -0,0 +1,112 @@ +/* + * 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.hadoop.hbase.quotas; + +import static org.apache.hadoop.hbase.quotas.ThrottleQuotaTestUtil.doPuts; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import java.util.concurrent.TimeUnit; +import org.apache.hadoop.hbase.HBaseClassTestRule; +import org.apache.hadoop.hbase.HBaseTestingUtil; +import org.apache.hadoop.hbase.HConstants; +import org.apache.hadoop.hbase.TableName; +import org.apache.hadoop.hbase.client.Admin; +import org.apache.hadoop.hbase.client.Table; +import org.apache.hadoop.hbase.security.User; +import org.apache.hadoop.hbase.testclassification.MediumTests; +import org.apache.hadoop.hbase.testclassification.RegionServerTests; +import org.apache.hadoop.hbase.util.Bytes; +import org.apache.hadoop.hbase.util.EnvironmentEdgeManager; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +@Category({ RegionServerTests.class, MediumTests.class }) +public class TestQuotaUserOverrideConfiguration { + + @ClassRule + public static final HBaseClassTestRule CLASS_RULE = + HBaseClassTestRule.forClass(TestQuotaUserOverrideConfiguration.class); + + private static final HBaseTestingUtil TEST_UTIL = new HBaseTestingUtil(); + private static final byte[] FAMILY = Bytes.toBytes("cf"); + private static final byte[] QUALIFIER = Bytes.toBytes("q"); + private static final int NUM_SERVERS = 1; + private static final String CUSTOM_OVERRIDE_KEY = "foo"; + + private static final TableName TABLE_NAME = + TableName.valueOf("TestQuotaUserOverrideConfiguration"); + + @BeforeClass + public static void setUpBeforeClass() throws Exception { + TEST_UTIL.getConfiguration().setBoolean(QuotaUtil.QUOTA_CONF_KEY, true); + TEST_UTIL.getConfiguration().setInt(QuotaCache.REFRESH_CONF_KEY, 1_000); + TEST_UTIL.getConfiguration().setInt("hbase.client.pause", 1); + TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, 0); + TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_OPERATION_TIMEOUT, 500); + TEST_UTIL.getConfiguration().set(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, + CUSTOM_OVERRIDE_KEY); + TEST_UTIL.startMiniCluster(NUM_SERVERS); + TEST_UTIL.waitTableAvailable(QuotaTableUtil.QUOTA_TABLE_NAME); + QuotaCache.TEST_FORCE_REFRESH = true; + TEST_UTIL.createTable(TABLE_NAME, FAMILY); + } + + @AfterClass + public static void tearDownAfterClass() throws Exception { + EnvironmentEdgeManager.reset(); + TEST_UTIL.shutdownMiniCluster(); + } + + @Test + public void testUserGlobalThrottleWithCustomOverride() throws Exception { + final Admin admin = TEST_UTIL.getAdmin(); + final String userOverrideWithQuota = User.getCurrent().getShortName() + "123"; + + // Add 6req/min limit + admin.setQuota(QuotaSettingsFactory.throttleUser(userOverrideWithQuota, + ThrottleType.REQUEST_NUMBER, 6, TimeUnit.MINUTES)); + Thread.sleep(60_000); + + Table tableWithThrottle = admin.getConnection().getTableBuilder(TABLE_NAME, null) + .setRequestAttribute(CUSTOM_OVERRIDE_KEY, Bytes.toBytes(userOverrideWithQuota)).build(); + Table tableWithoutThrottle = admin.getConnection().getTableBuilder(TABLE_NAME, null) + .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, + Bytes.toBytes(userOverrideWithQuota)) + .build(); + + // warm things up + doPuts(10, FAMILY, QUALIFIER, tableWithThrottle); + doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle); + + // should reject some requests + assertTrue(10 > doPuts(10, FAMILY, QUALIFIER, tableWithThrottle)); + // should accept all puts + assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle)); + + // Remove all the limits + admin.setQuota(QuotaSettingsFactory.unthrottleUser(userOverrideWithQuota)); + Thread.sleep(60_000); + assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithThrottle)); + assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle)); + } + +} From 951aec6e086f48532299b1f3b8df971deb7dd5b6 Mon Sep 17 00:00:00 2001 From: Ray Mattingly Date: Mon, 25 Sep 2023 16:46:21 -0400 Subject: [PATCH 2/5] initial PR feedback --- .../apache/hadoop/hbase/quotas/TestQuotaUserOverride.java | 5 ++--- .../hbase/quotas/TestQuotaUserOverrideConfiguration.java | 5 ++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java index 53e25a2aec9d..35561413202a 100644 --- a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java +++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java @@ -81,13 +81,12 @@ public void testUserGlobalThrottleWithOverride() throws Exception { // Add 6req/min limit admin.setQuota(QuotaSettingsFactory.throttleUser(userOverrideWithQuota, ThrottleType.REQUEST_NUMBER, 6, TimeUnit.MINUTES)); - Thread.sleep(60_000); - Table tableWithThrottle = admin.getConnection().getTableBuilder(TABLE_NAME, null) + Table tableWithThrottle = TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null) .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, Bytes.toBytes(userOverrideWithQuota)) .build(); - Table tableWithoutThrottle = admin.getConnection().getTableBuilder(TABLE_NAME, null) + Table tableWithoutThrottle = TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null) .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, Bytes.toBytes(userOverrideWithoutQuota)) .build(); diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java index 1eac049677e5..8421d64120be 100644 --- a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java +++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java @@ -84,11 +84,10 @@ public void testUserGlobalThrottleWithCustomOverride() throws Exception { // Add 6req/min limit admin.setQuota(QuotaSettingsFactory.throttleUser(userOverrideWithQuota, ThrottleType.REQUEST_NUMBER, 6, TimeUnit.MINUTES)); - Thread.sleep(60_000); - Table tableWithThrottle = admin.getConnection().getTableBuilder(TABLE_NAME, null) + Table tableWithThrottle = TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null) .setRequestAttribute(CUSTOM_OVERRIDE_KEY, Bytes.toBytes(userOverrideWithQuota)).build(); - Table tableWithoutThrottle = admin.getConnection().getTableBuilder(TABLE_NAME, null) + Table tableWithoutThrottle = TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null) .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, Bytes.toBytes(userOverrideWithQuota)) .build(); From 3528ca5b2cc48443576bca447fa6bc20430781f3 Mon Sep 17 00:00:00 2001 From: Ray Mattingly Date: Tue, 26 Sep 2023 16:50:16 -0400 Subject: [PATCH 3/5] removing default quota override key, unifying test --- .../hadoop/hbase/quotas/QuotaCache.java | 10 +- .../hbase/quotas/TestQuotaUserOverride.java | 15 +-- .../TestQuotaUserOverrideConfiguration.java | 111 ------------------ 3 files changed, 15 insertions(+), 121 deletions(-) delete mode 100644 hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java index eac70fce83a3..7fdf79820ff6 100644 --- a/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java +++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java @@ -63,7 +63,7 @@ public class QuotaCache implements Stoppable { public static final String REFRESH_CONF_KEY = "hbase.quota.refresh.period"; // defines the request attribute key which, when provided, will override the request's username - // from the perspective of user quotas. This is also the default request attribute key + // from the perspective of user quotas public static final String QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY = "hbase.quota.user.override.key"; private static final int REFRESH_DEFAULT_PERIOD = 5 * 60000; // 5min @@ -90,8 +90,8 @@ public class QuotaCache implements Stoppable { public QuotaCache(final RegionServerServices rsServices) { this.rsServices = rsServices; - this.userOverrideRequestAttributeKey = rsServices.getConfiguration() - .get(QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY); + this.userOverrideRequestAttributeKey = + rsServices.getConfiguration().get(QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY); } public void start() throws IOException { @@ -178,6 +178,10 @@ protected boolean isExceedThrottleQuotaEnabled() { * @param ugi The request's UserGroupInformation */ private String getQuotaUserName(final UserGroupInformation ugi) { + if (userOverrideRequestAttributeKey == null) { + return ugi.getShortUserName(); + } + Optional rpcCall = RpcServer.getCurrentCall(); if ( rpcCall.isPresent() diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java index 35561413202a..c369c05796d5 100644 --- a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java +++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java @@ -50,8 +50,10 @@ public class TestQuotaUserOverride { private static final byte[] FAMILY = Bytes.toBytes("cf"); private static final byte[] QUALIFIER = Bytes.toBytes("q"); private static final int NUM_SERVERS = 1; + private static final String CUSTOM_OVERRIDE_KEY = "foo"; - private static final TableName TABLE_NAME = TableName.valueOf("TestQuotaUserOverride"); + private static final TableName TABLE_NAME = + TableName.valueOf("TestQuotaUserOverride"); @BeforeClass public static void setUpBeforeClass() throws Exception { @@ -60,6 +62,8 @@ public static void setUpBeforeClass() throws Exception { TEST_UTIL.getConfiguration().setInt("hbase.client.pause", 1); TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, 0); TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_OPERATION_TIMEOUT, 500); + TEST_UTIL.getConfiguration().set(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, + CUSTOM_OVERRIDE_KEY); TEST_UTIL.startMiniCluster(NUM_SERVERS); TEST_UTIL.waitTableAvailable(QuotaTableUtil.QUOTA_TABLE_NAME); QuotaCache.TEST_FORCE_REFRESH = true; @@ -73,22 +77,19 @@ public static void tearDownAfterClass() throws Exception { } @Test - public void testUserGlobalThrottleWithOverride() throws Exception { + public void testUserGlobalThrottleWithCustomOverride() throws Exception { final Admin admin = TEST_UTIL.getAdmin(); final String userOverrideWithQuota = User.getCurrent().getShortName() + "123"; - final String userOverrideWithoutQuota = User.getCurrent().getShortName() + "456"; // Add 6req/min limit admin.setQuota(QuotaSettingsFactory.throttleUser(userOverrideWithQuota, ThrottleType.REQUEST_NUMBER, 6, TimeUnit.MINUTES)); Table tableWithThrottle = TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null) - .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, - Bytes.toBytes(userOverrideWithQuota)) - .build(); + .setRequestAttribute(CUSTOM_OVERRIDE_KEY, Bytes.toBytes(userOverrideWithQuota)).build(); Table tableWithoutThrottle = TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null) .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, - Bytes.toBytes(userOverrideWithoutQuota)) + Bytes.toBytes(userOverrideWithQuota)) .build(); // warm things up diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java deleted file mode 100644 index 8421d64120be..000000000000 --- a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverrideConfiguration.java +++ /dev/null @@ -1,111 +0,0 @@ -/* - * 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.hadoop.hbase.quotas; - -import static org.apache.hadoop.hbase.quotas.ThrottleQuotaTestUtil.doPuts; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertTrue; - -import java.util.concurrent.TimeUnit; -import org.apache.hadoop.hbase.HBaseClassTestRule; -import org.apache.hadoop.hbase.HBaseTestingUtil; -import org.apache.hadoop.hbase.HConstants; -import org.apache.hadoop.hbase.TableName; -import org.apache.hadoop.hbase.client.Admin; -import org.apache.hadoop.hbase.client.Table; -import org.apache.hadoop.hbase.security.User; -import org.apache.hadoop.hbase.testclassification.MediumTests; -import org.apache.hadoop.hbase.testclassification.RegionServerTests; -import org.apache.hadoop.hbase.util.Bytes; -import org.apache.hadoop.hbase.util.EnvironmentEdgeManager; -import org.junit.AfterClass; -import org.junit.BeforeClass; -import org.junit.ClassRule; -import org.junit.Test; -import org.junit.experimental.categories.Category; - -@Category({ RegionServerTests.class, MediumTests.class }) -public class TestQuotaUserOverrideConfiguration { - - @ClassRule - public static final HBaseClassTestRule CLASS_RULE = - HBaseClassTestRule.forClass(TestQuotaUserOverrideConfiguration.class); - - private static final HBaseTestingUtil TEST_UTIL = new HBaseTestingUtil(); - private static final byte[] FAMILY = Bytes.toBytes("cf"); - private static final byte[] QUALIFIER = Bytes.toBytes("q"); - private static final int NUM_SERVERS = 1; - private static final String CUSTOM_OVERRIDE_KEY = "foo"; - - private static final TableName TABLE_NAME = - TableName.valueOf("TestQuotaUserOverrideConfiguration"); - - @BeforeClass - public static void setUpBeforeClass() throws Exception { - TEST_UTIL.getConfiguration().setBoolean(QuotaUtil.QUOTA_CONF_KEY, true); - TEST_UTIL.getConfiguration().setInt(QuotaCache.REFRESH_CONF_KEY, 1_000); - TEST_UTIL.getConfiguration().setInt("hbase.client.pause", 1); - TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, 0); - TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_OPERATION_TIMEOUT, 500); - TEST_UTIL.getConfiguration().set(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, - CUSTOM_OVERRIDE_KEY); - TEST_UTIL.startMiniCluster(NUM_SERVERS); - TEST_UTIL.waitTableAvailable(QuotaTableUtil.QUOTA_TABLE_NAME); - QuotaCache.TEST_FORCE_REFRESH = true; - TEST_UTIL.createTable(TABLE_NAME, FAMILY); - } - - @AfterClass - public static void tearDownAfterClass() throws Exception { - EnvironmentEdgeManager.reset(); - TEST_UTIL.shutdownMiniCluster(); - } - - @Test - public void testUserGlobalThrottleWithCustomOverride() throws Exception { - final Admin admin = TEST_UTIL.getAdmin(); - final String userOverrideWithQuota = User.getCurrent().getShortName() + "123"; - - // Add 6req/min limit - admin.setQuota(QuotaSettingsFactory.throttleUser(userOverrideWithQuota, - ThrottleType.REQUEST_NUMBER, 6, TimeUnit.MINUTES)); - - Table tableWithThrottle = TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null) - .setRequestAttribute(CUSTOM_OVERRIDE_KEY, Bytes.toBytes(userOverrideWithQuota)).build(); - Table tableWithoutThrottle = TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null) - .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, - Bytes.toBytes(userOverrideWithQuota)) - .build(); - - // warm things up - doPuts(10, FAMILY, QUALIFIER, tableWithThrottle); - doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle); - - // should reject some requests - assertTrue(10 > doPuts(10, FAMILY, QUALIFIER, tableWithThrottle)); - // should accept all puts - assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle)); - - // Remove all the limits - admin.setQuota(QuotaSettingsFactory.unthrottleUser(userOverrideWithQuota)); - Thread.sleep(60_000); - assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithThrottle)); - assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle)); - } - -} From 761b4357c8e4aaa92d0a1a7b8281a56cab1b3ee3 Mon Sep 17 00:00:00 2001 From: Ray Mattingly Date: Tue, 26 Sep 2023 17:02:22 -0400 Subject: [PATCH 4/5] add table without override to test --- .../hadoop/hbase/quotas/TestQuotaUserOverride.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java index c369c05796d5..75b3cc3ca84a 100644 --- a/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java +++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestQuotaUserOverride.java @@ -52,8 +52,7 @@ public class TestQuotaUserOverride { private static final int NUM_SERVERS = 1; private static final String CUSTOM_OVERRIDE_KEY = "foo"; - private static final TableName TABLE_NAME = - TableName.valueOf("TestQuotaUserOverride"); + private static final TableName TABLE_NAME = TableName.valueOf("TestQuotaUserOverride"); @BeforeClass public static void setUpBeforeClass() throws Exception { @@ -91,21 +90,27 @@ public void testUserGlobalThrottleWithCustomOverride() throws Exception { .setRequestAttribute(QuotaCache.QUOTA_USER_REQUEST_ATTRIBUTE_OVERRIDE_KEY, Bytes.toBytes(userOverrideWithQuota)) .build(); + Table tableWithoutThrottle2 = + TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null).build(); // warm things up doPuts(10, FAMILY, QUALIFIER, tableWithThrottle); doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle); + doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle2); // should reject some requests assertTrue(10 > doPuts(10, FAMILY, QUALIFIER, tableWithThrottle)); // should accept all puts assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle)); + // should accept all puts + assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle2)); // Remove all the limits admin.setQuota(QuotaSettingsFactory.unthrottleUser(userOverrideWithQuota)); Thread.sleep(60_000); assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithThrottle)); assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle)); + assertEquals(10, doPuts(10, FAMILY, QUALIFIER, tableWithoutThrottle2)); } } From 345e5c849411d22c5091f7e895bd6d2f057c2949 Mon Sep 17 00:00:00 2001 From: Ray Mattingly Date: Wed, 27 Sep 2023 09:23:59 -0400 Subject: [PATCH 5/5] fetch single request attribute --- .../java/org/apache/hadoop/hbase/ipc/RpcCall.java | 12 +++++++++++- .../org/apache/hadoop/hbase/ipc/ServerCall.java | 13 +++++++++++++ .../apache/hadoop/hbase/quotas/QuotaCache.java | 15 ++++++++------- .../hbase/namequeues/TestNamedQueueRecorder.java | 5 +++++ .../store/region/TestRegionProcedureStore.java | 5 +++++ 5 files changed, 42 insertions(+), 8 deletions(-) diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/ipc/RpcCall.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/ipc/RpcCall.java index 0555202f88b9..260d6e1a9803 100644 --- a/hbase-server/src/main/java/org/apache/hadoop/hbase/ipc/RpcCall.java +++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/ipc/RpcCall.java @@ -92,11 +92,21 @@ public interface RpcCall extends RpcCallContext { Map getConnectionAttributes(); /** - * Returns the map of attributes specified when building the request. + * Returns the map of attributes specified when building the request. This map is lazily evaluated + * so if you only need a single attribute then it may be cheaper to use + * {@link #getRequestAttribute(String)} * @see org.apache.hadoop.hbase.client.TableBuilder#setRequestAttribute(String, byte[]) */ Map getRequestAttributes(); + /** + * Returns a single request attribute value, or null if no value is present. If you need many + * request attributes then you should fetch the lazily evaluated map via + * {@link #getRequestAttributes()} + * @see org.apache.hadoop.hbase.client.TableBuilder#setRequestAttribute(String, byte[]) + */ + byte[] getRequestAttribute(String key); + /** Returns Port of remote address in this call */ int getRemotePort(); diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/ipc/ServerCall.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/ipc/ServerCall.java index 66a2e44fac19..ed688977b963 100644 --- a/hbase-server/src/main/java/org/apache/hadoop/hbase/ipc/ServerCall.java +++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/ipc/ServerCall.java @@ -234,6 +234,19 @@ public Map getRequestAttributes() { return this.requestAttributes; } + @Override + public byte[] getRequestAttribute(String key) { + if (this.requestAttributes == null) { + for (HBaseProtos.NameBytesPair nameBytesPair : header.getAttributeList()) { + if (nameBytesPair.getName().equals(key)) { + return nameBytesPair.getValue().toByteArray(); + } + } + return null; + } + return this.requestAttributes.get(key); + } + @Override public int getPriority() { return this.header.getPriority(); diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java index 7fdf79820ff6..8434e1a4fd08 100644 --- a/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java +++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/QuotaCache.java @@ -183,14 +183,15 @@ private String getQuotaUserName(final UserGroupInformation ugi) { } Optional rpcCall = RpcServer.getCurrentCall(); - if ( - rpcCall.isPresent() - && rpcCall.get().getRequestAttributes().containsKey(userOverrideRequestAttributeKey) - ) { - return Bytes - .toString(rpcCall.get().getRequestAttributes().get(userOverrideRequestAttributeKey)); + if (!rpcCall.isPresent()) { + return ugi.getShortUserName(); + } + + byte[] override = rpcCall.get().getRequestAttribute(userOverrideRequestAttributeKey); + if (override == null) { + return ugi.getShortUserName(); } - return ugi.getShortUserName(); + return Bytes.toString(override); } /** diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/namequeues/TestNamedQueueRecorder.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/namequeues/TestNamedQueueRecorder.java index af6c51260fd5..35a1757115c9 100644 --- a/hbase-server/src/test/java/org/apache/hadoop/hbase/namequeues/TestNamedQueueRecorder.java +++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/namequeues/TestNamedQueueRecorder.java @@ -771,6 +771,11 @@ public Map getRequestAttributes() { pair -> pair.getValue().toByteArray())); } + @Override + public byte[] getRequestAttribute(String key) { + return null; + } + @Override public int getRemotePort() { return 0; diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/procedure2/store/region/TestRegionProcedureStore.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/procedure2/store/region/TestRegionProcedureStore.java index 83f788ba1518..305f0e29e952 100644 --- a/hbase-server/src/test/java/org/apache/hadoop/hbase/procedure2/store/region/TestRegionProcedureStore.java +++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/procedure2/store/region/TestRegionProcedureStore.java @@ -232,6 +232,11 @@ public Map getRequestAttributes() { return null; } + @Override + public byte[] getRequestAttribute(String key) { + return null; + } + @Override public int getRemotePort() { return 0;