001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hbase.client;
019
020import static org.apache.hadoop.hbase.client.AsyncConnectionConfiguration.START_LOG_ERRORS_AFTER_COUNT_KEY;
021import static org.junit.jupiter.api.Assertions.assertEquals;
022import static org.junit.jupiter.api.Assertions.fail;
023
024import java.util.Objects;
025import java.util.concurrent.CompletableFuture;
026import java.util.concurrent.TimeUnit;
027import java.util.function.Supplier;
028import org.apache.hadoop.hbase.HBaseParameterizedTestTemplate;
029import org.apache.hadoop.hbase.HConstants;
030import org.apache.hadoop.hbase.TableName;
031import org.apache.hadoop.hbase.quotas.QuotaCache;
032import org.apache.hadoop.hbase.quotas.QuotaFilter;
033import org.apache.hadoop.hbase.quotas.QuotaSettings;
034import org.apache.hadoop.hbase.quotas.QuotaSettingsFactory;
035import org.apache.hadoop.hbase.quotas.QuotaTableUtil;
036import org.apache.hadoop.hbase.quotas.QuotaUtil;
037import org.apache.hadoop.hbase.quotas.ThrottleType;
038import org.apache.hadoop.hbase.security.User;
039import org.apache.hadoop.hbase.testclassification.ClientTests;
040import org.apache.hadoop.hbase.testclassification.MediumTests;
041import org.junit.jupiter.api.AfterAll;
042import org.junit.jupiter.api.BeforeAll;
043import org.junit.jupiter.api.Tag;
044import org.junit.jupiter.api.TestTemplate;
045
046@Tag(ClientTests.TAG)
047@Tag(MediumTests.TAG)
048@HBaseParameterizedTestTemplate(name = "{index}: policy = {0}")
049public class TestAsyncQuotaAdminApi extends TestAsyncAdminBase {
050
051  public TestAsyncQuotaAdminApi(Supplier<AsyncAdmin> admin) {
052    super(admin);
053  }
054
055  @BeforeAll
056  public static void setUpBeforeClass() throws Exception {
057    TEST_UTIL.getConfiguration().setBoolean(QuotaUtil.QUOTA_CONF_KEY, true);
058    TEST_UTIL.getConfiguration().setInt(QuotaCache.REFRESH_CONF_KEY, 2000);
059    TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_RPC_TIMEOUT_KEY, 60000);
060    TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_OPERATION_TIMEOUT, 120000);
061    TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, 2);
062    TEST_UTIL.getConfiguration().setInt(START_LOG_ERRORS_AFTER_COUNT_KEY, 0);
063    TEST_UTIL.startMiniCluster(1);
064    TEST_UTIL.waitTableAvailable(QuotaTableUtil.QUOTA_TABLE_NAME);
065    ASYNC_CONN = ConnectionFactory.createAsyncConnection(TEST_UTIL.getConfiguration()).get();
066  }
067
068  @AfterAll
069  public static void tearDownAfterClass() throws Exception {
070    TestAsyncAdminBase.tearDownAfterClass();
071  }
072
073  @TestTemplate
074  public void testThrottleType() throws Exception {
075    String userName = User.getCurrent().getShortName();
076
077    admin
078      .setQuota(
079        QuotaSettingsFactory.throttleUser(userName, ThrottleType.READ_NUMBER, 6, TimeUnit.MINUTES))
080      .get();
081    admin.setQuota(
082      QuotaSettingsFactory.throttleUser(userName, ThrottleType.WRITE_NUMBER, 12, TimeUnit.MINUTES))
083      .get();
084    admin.setQuota(QuotaSettingsFactory.bypassGlobals(userName, true)).get();
085
086    int countThrottle = 0;
087    int countGlobalBypass = 0;
088    for (QuotaSettings settings : admin.getQuota(null).get()) {
089      switch (settings.getQuotaType()) {
090        case THROTTLE:
091          countThrottle++;
092          break;
093        case GLOBAL_BYPASS:
094          countGlobalBypass++;
095          break;
096        default:
097          fail("unexpected settings type: " + settings.getQuotaType());
098      }
099    }
100    assertEquals(2, countThrottle);
101    assertEquals(1, countGlobalBypass);
102
103    admin.setQuota(QuotaSettingsFactory.unthrottleUser(userName)).get();
104    assertNumResults(1, null);
105    admin.setQuota(QuotaSettingsFactory.bypassGlobals(userName, false)).get();
106    assertNumResults(0, null);
107  }
108
109  @TestTemplate
110  public void testQuotaRetrieverFilter() throws Exception {
111    TableName[] tables = new TableName[] { TableName.valueOf("T0"), TableName.valueOf("T01"),
112      TableName.valueOf("NS0:T2"), };
113    String[] namespaces = new String[] { "NS0", "NS01", "NS2" };
114    String[] users = new String[] { "User0", "User01", "User2" };
115
116    for (String user : users) {
117      admin
118        .setQuota(
119          QuotaSettingsFactory.throttleUser(user, ThrottleType.REQUEST_NUMBER, 1, TimeUnit.MINUTES))
120        .get();
121
122      for (TableName table : tables) {
123        admin.setQuota(QuotaSettingsFactory.throttleUser(user, table, ThrottleType.REQUEST_NUMBER,
124          2, TimeUnit.MINUTES)).get();
125      }
126
127      for (String ns : namespaces) {
128        admin.setQuota(QuotaSettingsFactory.throttleUser(user, ns, ThrottleType.REQUEST_NUMBER, 3,
129          TimeUnit.MINUTES)).get();
130      }
131    }
132    assertNumResults(21, null);
133
134    for (TableName table : tables) {
135      admin.setQuota(
136        QuotaSettingsFactory.throttleTable(table, ThrottleType.REQUEST_NUMBER, 4, TimeUnit.MINUTES))
137        .get();
138    }
139    assertNumResults(24, null);
140
141    for (String ns : namespaces) {
142      admin.setQuota(QuotaSettingsFactory.throttleNamespace(ns, ThrottleType.REQUEST_NUMBER, 5,
143        TimeUnit.MINUTES)).get();
144    }
145    assertNumResults(27, null);
146
147    assertNumResults(7, new QuotaFilter().setUserFilter("User0"));
148    assertNumResults(0, new QuotaFilter().setUserFilter("User"));
149    assertNumResults(21, new QuotaFilter().setUserFilter("User.*"));
150    assertNumResults(3, new QuotaFilter().setUserFilter("User.*").setTableFilter("T0"));
151    assertNumResults(3, new QuotaFilter().setUserFilter("User.*").setTableFilter("NS.*"));
152    assertNumResults(0, new QuotaFilter().setUserFilter("User.*").setTableFilter("T"));
153    assertNumResults(6, new QuotaFilter().setUserFilter("User.*").setTableFilter("T.*"));
154    assertNumResults(3, new QuotaFilter().setUserFilter("User.*").setNamespaceFilter("NS0"));
155    assertNumResults(0, new QuotaFilter().setUserFilter("User.*").setNamespaceFilter("NS"));
156    assertNumResults(9, new QuotaFilter().setUserFilter("User.*").setNamespaceFilter("NS.*"));
157    assertNumResults(6,
158      new QuotaFilter().setUserFilter("User.*").setTableFilter("T0").setNamespaceFilter("NS0"));
159    assertNumResults(1, new QuotaFilter().setTableFilter("T0"));
160    assertNumResults(0, new QuotaFilter().setTableFilter("T"));
161    assertNumResults(2, new QuotaFilter().setTableFilter("T.*"));
162    assertNumResults(3, new QuotaFilter().setTableFilter(".*T.*"));
163    assertNumResults(1, new QuotaFilter().setNamespaceFilter("NS0"));
164    assertNumResults(0, new QuotaFilter().setNamespaceFilter("NS"));
165    assertNumResults(3, new QuotaFilter().setNamespaceFilter("NS.*"));
166
167    for (String user : users) {
168      admin.setQuota(QuotaSettingsFactory.unthrottleUser(user)).get();
169      for (TableName table : tables) {
170        admin.setQuota(QuotaSettingsFactory.unthrottleUser(user, table)).get();
171      }
172      for (String ns : namespaces) {
173        admin.setQuota(QuotaSettingsFactory.unthrottleUser(user, ns)).get();
174      }
175    }
176    assertNumResults(6, null);
177
178    for (TableName table : tables) {
179      admin.setQuota(QuotaSettingsFactory.unthrottleTable(table)).get();
180    }
181    assertNumResults(3, null);
182
183    for (String ns : namespaces) {
184      admin.setQuota(QuotaSettingsFactory.unthrottleNamespace(ns)).get();
185    }
186    assertNumResults(0, null);
187  }
188
189  @TestTemplate
190  public void testSwitchRpcThrottle() throws Exception {
191    CompletableFuture<Boolean> future1 = ASYNC_CONN.getAdmin().switchRpcThrottle(true);
192    assertEquals(true, future1.get().booleanValue());
193    CompletableFuture<Boolean> future2 = ASYNC_CONN.getAdmin().isRpcThrottleEnabled();
194    assertEquals(true, future2.get().booleanValue());
195  }
196
197  @TestTemplate
198  public void testSwitchExceedThrottleQuota() throws Exception {
199    AsyncAdmin admin = ASYNC_CONN.getAdmin();
200    assertEquals(false, admin.exceedThrottleQuotaSwitch(false).get().booleanValue());
201  }
202
203  private void assertNumResults(int expected, final QuotaFilter filter) throws Exception {
204    assertEquals(expected, countResults(filter));
205  }
206
207  private int countResults(final QuotaFilter filter) throws Exception {
208    int count = 0;
209    for (QuotaSettings settings : admin.getQuota(filter).get()) {
210      LOG.debug(Objects.toString(settings));
211      count++;
212    }
213    return count;
214  }
215}