From ac13213fe6630d7b834d7fc41672d6e16acc66d1 Mon Sep 17 00:00:00 2001 From: Ray Mattingly Date: Tue, 19 Sep 2023 19:11:55 -0400 Subject: [PATCH] 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)); + } + +}