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.quotas;
019
020import static org.junit.Assert.assertEquals;
021import static org.junit.Assert.assertFalse;
022import static org.junit.Assert.assertNull;
023import static org.junit.Assert.assertTrue;
024import static org.junit.Assert.fail;
025
026import java.io.IOException;
027import java.util.ArrayList;
028import java.util.List;
029import java.util.Objects;
030import java.util.concurrent.TimeUnit;
031
032import org.apache.hadoop.hbase.Cell;
033import org.apache.hadoop.hbase.CellScanner;
034import org.apache.hadoop.hbase.HBaseClassTestRule;
035import org.apache.hadoop.hbase.HBaseTestingUtility;
036import org.apache.hadoop.hbase.HConstants;
037import org.apache.hadoop.hbase.TableName;
038import org.apache.hadoop.hbase.client.Admin;
039import org.apache.hadoop.hbase.client.ColumnFamilyDescriptorBuilder;
040import org.apache.hadoop.hbase.client.Put;
041import org.apache.hadoop.hbase.client.Result;
042import org.apache.hadoop.hbase.client.ResultScanner;
043import org.apache.hadoop.hbase.client.Scan;
044import org.apache.hadoop.hbase.client.Table;
045import org.apache.hadoop.hbase.client.TableDescriptor;
046import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
047import org.apache.hadoop.hbase.security.User;
048import org.apache.hadoop.hbase.testclassification.ClientTests;
049import org.apache.hadoop.hbase.testclassification.MediumTests;
050import org.apache.hadoop.hbase.util.Bytes;
051import org.apache.hadoop.hbase.util.JVMClusterUtil;
052import org.junit.After;
053import org.junit.AfterClass;
054import org.junit.Assert;
055import org.junit.BeforeClass;
056import org.junit.ClassRule;
057import org.junit.Test;
058import org.junit.experimental.categories.Category;
059import org.slf4j.Logger;
060import org.slf4j.LoggerFactory;
061
062import org.apache.hbase.thirdparty.com.google.common.collect.Iterables;
063
064import org.apache.hadoop.hbase.shaded.protobuf.ProtobufUtil;
065import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos;
066import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.Quotas;
067import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.SpaceLimitRequest;
068
069/**
070 * minicluster tests that validate that quota  entries are properly set in the quota table
071 */
072@Category({ClientTests.class, MediumTests.class})
073public class TestQuotaAdmin {
074
075  @ClassRule
076  public static final HBaseClassTestRule CLASS_RULE =
077      HBaseClassTestRule.forClass(TestQuotaAdmin.class);
078
079  private static final Logger LOG = LoggerFactory.getLogger(TestQuotaAdmin.class);
080
081  private final static HBaseTestingUtility TEST_UTIL = new HBaseTestingUtility();
082
083  @BeforeClass
084  public static void setUpBeforeClass() throws Exception {
085    TEST_UTIL.getConfiguration().setBoolean(QuotaUtil.QUOTA_CONF_KEY, true);
086    TEST_UTIL.getConfiguration().setInt(QuotaCache.REFRESH_CONF_KEY, 2000);
087    TEST_UTIL.getConfiguration().setInt("hbase.hstore.compactionThreshold", 10);
088    TEST_UTIL.getConfiguration().setInt("hbase.regionserver.msginterval", 100);
089    TEST_UTIL.getConfiguration().setInt("hbase.client.pause", 250);
090    TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, 6);
091    TEST_UTIL.getConfiguration().setBoolean("hbase.master.enabletable.roundrobin", true);
092    TEST_UTIL.startMiniCluster(1);
093    TEST_UTIL.waitTableAvailable(QuotaTableUtil.QUOTA_TABLE_NAME);
094  }
095
096  @After
097  public void clearQuotaTable() throws Exception {
098    if (TEST_UTIL.getAdmin().tableExists(QuotaUtil.QUOTA_TABLE_NAME)) {
099      TEST_UTIL.getAdmin().disableTable(QuotaUtil.QUOTA_TABLE_NAME);
100      TEST_UTIL.getAdmin().truncateTable(QuotaUtil.QUOTA_TABLE_NAME, false);
101    }
102  }
103
104  @AfterClass
105  public static void tearDownAfterClass() throws Exception {
106    TEST_UTIL.shutdownMiniCluster();
107  }
108
109  @Test
110  public void testThrottleType() throws Exception {
111    Admin admin = TEST_UTIL.getAdmin();
112    String userName = User.getCurrent().getShortName();
113
114    admin.setQuota(
115        QuotaSettingsFactory.throttleUser(userName, ThrottleType.READ_NUMBER, 6, TimeUnit.MINUTES));
116    admin.setQuota(QuotaSettingsFactory
117        .throttleUser(userName, ThrottleType.WRITE_NUMBER, 12, TimeUnit.MINUTES));
118    admin.setQuota(QuotaSettingsFactory.bypassGlobals(userName, true));
119
120    try (QuotaRetriever scanner = QuotaRetriever.open(TEST_UTIL.getConfiguration())) {
121      int countThrottle = 0;
122      int countGlobalBypass = 0;
123      for (QuotaSettings settings: scanner) {
124        switch (settings.getQuotaType()) {
125          case THROTTLE:
126            ThrottleSettings throttle = (ThrottleSettings)settings;
127            if (throttle.getSoftLimit() == 6) {
128              assertEquals(ThrottleType.READ_NUMBER, throttle.getThrottleType());
129            } else if (throttle.getSoftLimit() == 12) {
130              assertEquals(ThrottleType.WRITE_NUMBER, throttle.getThrottleType());
131            } else {
132              fail("should not come here, because don't set quota with this limit");
133            }
134            assertEquals(userName, throttle.getUserName());
135            assertEquals(null, throttle.getTableName());
136            assertEquals(null, throttle.getNamespace());
137            assertEquals(TimeUnit.MINUTES, throttle.getTimeUnit());
138            countThrottle++;
139            break;
140          case GLOBAL_BYPASS:
141            countGlobalBypass++;
142            break;
143          default:
144            fail("unexpected settings type: " + settings.getQuotaType());
145        }
146      }
147      assertEquals(2, countThrottle);
148      assertEquals(1, countGlobalBypass);
149    }
150
151    admin.setQuota(QuotaSettingsFactory.unthrottleUser(userName));
152    assertNumResults(1, null);
153    admin.setQuota(QuotaSettingsFactory.bypassGlobals(userName, false));
154    assertNumResults(0, null);
155  }
156
157  @Test
158  public void testSimpleScan() throws Exception {
159    Admin admin = TEST_UTIL.getAdmin();
160    String userName = User.getCurrent().getShortName();
161
162    admin.setQuota(QuotaSettingsFactory
163      .throttleUser(userName, ThrottleType.REQUEST_NUMBER, 6, TimeUnit.MINUTES));
164    admin.setQuota(QuotaSettingsFactory.bypassGlobals(userName, true));
165
166    try (QuotaRetriever scanner = QuotaRetriever.open(TEST_UTIL.getConfiguration())) {
167      int countThrottle = 0;
168      int countGlobalBypass = 0;
169      for (QuotaSettings settings: scanner) {
170        LOG.debug(Objects.toString(settings));
171        switch (settings.getQuotaType()) {
172          case THROTTLE:
173            ThrottleSettings throttle = (ThrottleSettings)settings;
174            assertEquals(userName, throttle.getUserName());
175            assertEquals(null, throttle.getTableName());
176            assertEquals(null, throttle.getNamespace());
177            assertEquals(null, throttle.getRegionServer());
178            assertEquals(6, throttle.getSoftLimit());
179            assertEquals(TimeUnit.MINUTES, throttle.getTimeUnit());
180            countThrottle++;
181            break;
182          case GLOBAL_BYPASS:
183            countGlobalBypass++;
184            break;
185          default:
186            fail("unexpected settings type: " + settings.getQuotaType());
187        }
188      }
189      assertEquals(1, countThrottle);
190      assertEquals(1, countGlobalBypass);
191    }
192
193    admin.setQuota(QuotaSettingsFactory.unthrottleUser(userName));
194    assertNumResults(1, null);
195    admin.setQuota(QuotaSettingsFactory.bypassGlobals(userName, false));
196    assertNumResults(0, null);
197  }
198
199  @Test
200  public void testMultiQuotaThrottling() throws Exception {
201    byte[] FAMILY = Bytes.toBytes("testFamily");
202    byte[] ROW = Bytes.toBytes("testRow");
203    byte[] QUALIFIER = Bytes.toBytes("testQualifier");
204    byte[] VALUE = Bytes.toBytes("testValue");
205
206    Admin admin = TEST_UTIL.getAdmin();
207    TableName tableName = TableName.valueOf("testMultiQuotaThrottling");
208    TableDescriptor desc = TableDescriptorBuilder.newBuilder(tableName)
209        .setColumnFamily(ColumnFamilyDescriptorBuilder.of(FAMILY)).build();
210    admin.createTable(desc);
211
212    // Set up the quota.
213    admin.setQuota(QuotaSettingsFactory.throttleTable(tableName, ThrottleType.WRITE_NUMBER, 6,
214        TimeUnit.SECONDS));
215
216    Thread.sleep(1000);
217    TEST_UTIL.getRSForFirstRegionInTable(tableName).getRegionServerRpcQuotaManager().
218        getQuotaCache().triggerCacheRefresh();
219    Thread.sleep(1000);
220
221    Table t =  TEST_UTIL.getConnection().getTable(tableName);
222    try {
223      int size = 5;
224      List actions = new ArrayList();
225      Object[] results = new Object[size];
226
227      for (int i = 0; i < size; i++) {
228        Put put1 = new Put(ROW);
229        put1.addColumn(FAMILY, QUALIFIER, VALUE);
230        actions.add(put1);
231      }
232      t.batch(actions, results);
233      t.batch(actions, results);
234    } catch (IOException e) {
235      fail("Not supposed to get ThrottlingExcepiton " + e);
236    } finally {
237      t.close();
238    }
239  }
240
241
242  @Test
243  public void testQuotaRetrieverFilter() throws Exception {
244    Admin admin = TEST_UTIL.getAdmin();
245    TableName[] tables = new TableName[] {
246      TableName.valueOf("T0"), TableName.valueOf("T01"), TableName.valueOf("NS0:T2"),
247    };
248    String[] namespaces = new String[] { "NS0", "NS01", "NS2" };
249    String[] users = new String[] { "User0", "User01", "User2" };
250
251    for (String user: users) {
252      admin.setQuota(QuotaSettingsFactory
253        .throttleUser(user, ThrottleType.REQUEST_NUMBER, 1, TimeUnit.MINUTES));
254
255      for (TableName table: tables) {
256        admin.setQuota(QuotaSettingsFactory
257          .throttleUser(user, table, ThrottleType.REQUEST_NUMBER, 2, TimeUnit.MINUTES));
258      }
259
260      for (String ns: namespaces) {
261        admin.setQuota(QuotaSettingsFactory
262          .throttleUser(user, ns, ThrottleType.REQUEST_NUMBER, 3, TimeUnit.MINUTES));
263      }
264    }
265    assertNumResults(21, null);
266
267    for (TableName table: tables) {
268      admin.setQuota(QuotaSettingsFactory
269        .throttleTable(table, ThrottleType.REQUEST_NUMBER, 4, TimeUnit.MINUTES));
270    }
271    assertNumResults(24, null);
272
273    for (String ns: namespaces) {
274      admin.setQuota(QuotaSettingsFactory
275        .throttleNamespace(ns, ThrottleType.REQUEST_NUMBER, 5, TimeUnit.MINUTES));
276    }
277    assertNumResults(27, null);
278
279    assertNumResults(7, new QuotaFilter().setUserFilter("User0"));
280    assertNumResults(0, new QuotaFilter().setUserFilter("User"));
281    assertNumResults(21, new QuotaFilter().setUserFilter("User.*"));
282    assertNumResults(3, new QuotaFilter().setUserFilter("User.*").setTableFilter("T0"));
283    assertNumResults(3, new QuotaFilter().setUserFilter("User.*").setTableFilter("NS.*"));
284    assertNumResults(0, new QuotaFilter().setUserFilter("User.*").setTableFilter("T"));
285    assertNumResults(6, new QuotaFilter().setUserFilter("User.*").setTableFilter("T.*"));
286    assertNumResults(3, new QuotaFilter().setUserFilter("User.*").setNamespaceFilter("NS0"));
287    assertNumResults(0, new QuotaFilter().setUserFilter("User.*").setNamespaceFilter("NS"));
288    assertNumResults(9, new QuotaFilter().setUserFilter("User.*").setNamespaceFilter("NS.*"));
289    assertNumResults(6, new QuotaFilter().setUserFilter("User.*")
290                                            .setTableFilter("T0").setNamespaceFilter("NS0"));
291    assertNumResults(1, new QuotaFilter().setTableFilter("T0"));
292    assertNumResults(0, new QuotaFilter().setTableFilter("T"));
293    assertNumResults(2, new QuotaFilter().setTableFilter("T.*"));
294    assertNumResults(3, new QuotaFilter().setTableFilter(".*T.*"));
295    assertNumResults(1, new QuotaFilter().setNamespaceFilter("NS0"));
296    assertNumResults(0, new QuotaFilter().setNamespaceFilter("NS"));
297    assertNumResults(3, new QuotaFilter().setNamespaceFilter("NS.*"));
298
299    for (String user: users) {
300      admin.setQuota(QuotaSettingsFactory.unthrottleUser(user));
301      for (TableName table: tables) {
302        admin.setQuota(QuotaSettingsFactory.unthrottleUser(user, table));
303      }
304      for (String ns: namespaces) {
305        admin.setQuota(QuotaSettingsFactory.unthrottleUser(user, ns));
306      }
307    }
308    assertNumResults(6, null);
309
310    for (TableName table: tables) {
311      admin.setQuota(QuotaSettingsFactory.unthrottleTable(table));
312    }
313    assertNumResults(3, null);
314
315    for (String ns: namespaces) {
316      admin.setQuota(QuotaSettingsFactory.unthrottleNamespace(ns));
317    }
318    assertNumResults(0, null);
319  }
320
321  @Test
322  public void testSetGetRemoveSpaceQuota() throws Exception {
323    Admin admin = TEST_UTIL.getAdmin();
324    final TableName tn = TableName.valueOf("sq_table1");
325    final long sizeLimit = 1024L * 1024L * 1024L * 1024L * 5L; // 5TB
326    final SpaceViolationPolicy violationPolicy = SpaceViolationPolicy.NO_WRITES;
327    QuotaSettings settings = QuotaSettingsFactory.limitTableSpace(tn, sizeLimit, violationPolicy);
328    admin.setQuota(settings);
329
330    // Verify the Quotas in the table
331    try (Table quotaTable = TEST_UTIL.getConnection().getTable(QuotaTableUtil.QUOTA_TABLE_NAME)) {
332      ResultScanner scanner = quotaTable.getScanner(new Scan());
333      try {
334        Result r = Iterables.getOnlyElement(scanner);
335        CellScanner cells = r.cellScanner();
336        assertTrue("Expected to find a cell", cells.advance());
337        assertSpaceQuota(sizeLimit, violationPolicy, cells.current());
338      } finally {
339        scanner.close();
340      }
341    }
342
343    // Verify we can retrieve it via the QuotaRetriever API
344    QuotaRetriever scanner = QuotaRetriever.open(admin.getConfiguration());
345    try {
346      assertSpaceQuota(sizeLimit, violationPolicy, Iterables.getOnlyElement(scanner));
347    } finally {
348      scanner.close();
349    }
350
351    // Now, remove the quota
352    QuotaSettings removeQuota = QuotaSettingsFactory.removeTableSpaceLimit(tn);
353    admin.setQuota(removeQuota);
354
355    // Verify that the record doesn't exist in the table
356    try (Table quotaTable = TEST_UTIL.getConnection().getTable(QuotaTableUtil.QUOTA_TABLE_NAME)) {
357      ResultScanner rs = quotaTable.getScanner(new Scan());
358      try {
359        assertNull("Did not expect to find a quota entry", rs.next());
360      } finally {
361        rs.close();
362      }
363    }
364
365    // Verify that we can also not fetch it via the API
366    scanner = QuotaRetriever.open(admin.getConfiguration());
367    try {
368      assertNull("Did not expect to find a quota entry", scanner.next());
369    } finally {
370      scanner.close();
371    }
372  }
373
374  @Test
375  public void testSetModifyRemoveSpaceQuota() throws Exception {
376    Admin admin = TEST_UTIL.getAdmin();
377    final TableName tn = TableName.valueOf("sq_table2");
378    final long originalSizeLimit = 1024L * 1024L * 1024L * 1024L * 5L; // 5TB
379    final SpaceViolationPolicy violationPolicy = SpaceViolationPolicy.NO_WRITES;
380    QuotaSettings settings = QuotaSettingsFactory.limitTableSpace(tn, originalSizeLimit,
381        violationPolicy);
382    admin.setQuota(settings);
383
384    // Verify the Quotas in the table
385    try (Table quotaTable = TEST_UTIL.getConnection().getTable(QuotaTableUtil.QUOTA_TABLE_NAME)) {
386      ResultScanner scanner = quotaTable.getScanner(new Scan());
387      try {
388        Result r = Iterables.getOnlyElement(scanner);
389        CellScanner cells = r.cellScanner();
390        assertTrue("Expected to find a cell", cells.advance());
391        assertSpaceQuota(originalSizeLimit, violationPolicy, cells.current());
392      } finally {
393        scanner.close();
394      }
395    }
396
397    // Verify we can retrieve it via the QuotaRetriever API
398    QuotaRetriever quotaScanner = QuotaRetriever.open(admin.getConfiguration());
399    try {
400      assertSpaceQuota(originalSizeLimit, violationPolicy, Iterables.getOnlyElement(quotaScanner));
401    } finally {
402      quotaScanner.close();
403    }
404
405    // Setting a new size and policy should be reflected
406    final long newSizeLimit = 1024L * 1024L * 1024L * 1024L; // 1TB
407    final SpaceViolationPolicy newViolationPolicy = SpaceViolationPolicy.NO_WRITES_COMPACTIONS;
408    QuotaSettings newSettings = QuotaSettingsFactory.limitTableSpace(tn, newSizeLimit,
409        newViolationPolicy);
410    admin.setQuota(newSettings);
411
412    // Verify the new Quotas in the table
413    try (Table quotaTable = TEST_UTIL.getConnection().getTable(QuotaTableUtil.QUOTA_TABLE_NAME)) {
414      ResultScanner scanner = quotaTable.getScanner(new Scan());
415      try {
416        Result r = Iterables.getOnlyElement(scanner);
417        CellScanner cells = r.cellScanner();
418        assertTrue("Expected to find a cell", cells.advance());
419        assertSpaceQuota(newSizeLimit, newViolationPolicy, cells.current());
420      } finally {
421        scanner.close();
422      }
423    }
424
425    // Verify we can retrieve the new quota via the QuotaRetriever API
426    quotaScanner = QuotaRetriever.open(admin.getConfiguration());
427    try {
428      assertSpaceQuota(newSizeLimit, newViolationPolicy, Iterables.getOnlyElement(quotaScanner));
429    } finally {
430      quotaScanner.close();
431    }
432
433    // Now, remove the quota
434    QuotaSettings removeQuota = QuotaSettingsFactory.removeTableSpaceLimit(tn);
435    admin.setQuota(removeQuota);
436
437    // Verify that the record doesn't exist in the table
438    try (Table quotaTable = TEST_UTIL.getConnection().getTable(QuotaTableUtil.QUOTA_TABLE_NAME)) {
439      ResultScanner scanner = quotaTable.getScanner(new Scan());
440      try {
441        assertNull("Did not expect to find a quota entry", scanner.next());
442      } finally {
443        scanner.close();
444      }
445    }
446
447    // Verify that we can also not fetch it via the API
448    quotaScanner = QuotaRetriever.open(admin.getConfiguration());
449    try {
450      assertNull("Did not expect to find a quota entry", quotaScanner.next());
451    } finally {
452      quotaScanner.close();
453    }
454  }
455
456  private void assertNumResults(int expected, final QuotaFilter filter) throws Exception {
457    assertEquals(expected, countResults(filter));
458  }
459
460  @Test
461  public void testSetGetRemoveRPCQuota() throws Exception {
462    testSetGetRemoveRPCQuota(ThrottleType.REQUEST_SIZE);
463    testSetGetRemoveRPCQuota(ThrottleType.REQUEST_CAPACITY_UNIT);
464  }
465
466  private void testSetGetRemoveRPCQuota(ThrottleType throttleType) throws Exception {
467    Admin admin = TEST_UTIL.getAdmin();
468    final TableName tn = TableName.valueOf("sq_table1");
469    QuotaSettings settings =
470        QuotaSettingsFactory.throttleTable(tn, throttleType, 2L, TimeUnit.HOURS);
471    admin.setQuota(settings);
472
473    // Verify the Quota in the table
474    verifyRecordPresentInQuotaTable(throttleType, 2L, TimeUnit.HOURS);
475
476    // Verify we can retrieve it via the QuotaRetriever API
477    verifyFetchableViaAPI(admin, throttleType, 2L, TimeUnit.HOURS);
478
479    // Now, remove the quota
480    QuotaSettings removeQuota = QuotaSettingsFactory.unthrottleTable(tn);
481    admin.setQuota(removeQuota);
482
483    // Verify that the record doesn't exist in the table
484    verifyRecordNotPresentInQuotaTable();
485
486    // Verify that we can also not fetch it via the API
487    verifyNotFetchableViaAPI(admin);
488  }
489
490  @Test
491  public void testSetModifyRemoveRPCQuota() throws Exception {
492    Admin admin = TEST_UTIL.getAdmin();
493    final TableName tn = TableName.valueOf("sq_table1");
494    QuotaSettings settings =
495        QuotaSettingsFactory.throttleTable(tn, ThrottleType.REQUEST_SIZE, 2L, TimeUnit.HOURS);
496    admin.setQuota(settings);
497
498    // Verify the Quota in the table
499    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_SIZE, 2L, TimeUnit.HOURS);
500
501    // Verify we can retrieve it via the QuotaRetriever API
502    verifyFetchableViaAPI(admin, ThrottleType.REQUEST_SIZE, 2L, TimeUnit.HOURS);
503
504    // Setting a limit and time unit should be reflected
505    QuotaSettings newSettings =
506        QuotaSettingsFactory.throttleTable(tn, ThrottleType.REQUEST_SIZE, 3L, TimeUnit.DAYS);
507    admin.setQuota(newSettings);
508
509    // Verify the new Quota in the table
510    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_SIZE, 3L, TimeUnit.DAYS);
511
512    // Verify we can retrieve the new quota via the QuotaRetriever API
513    verifyFetchableViaAPI(admin, ThrottleType.REQUEST_SIZE, 3L, TimeUnit.DAYS);
514
515    // Now, remove the quota
516    QuotaSettings removeQuota = QuotaSettingsFactory.unthrottleTable(tn);
517    admin.setQuota(removeQuota);
518
519    // Verify that the record doesn't exist in the table
520    verifyRecordNotPresentInQuotaTable();
521
522    // Verify that we can also not fetch it via the API
523    verifyNotFetchableViaAPI(admin);
524
525  }
526
527  @Test
528  public void testSetAndRemoveRegionServerQuota() throws Exception {
529    Admin admin = TEST_UTIL.getAdmin();
530    String regionServer = QuotaTableUtil.QUOTA_REGION_SERVER_ROW_KEY;
531    QuotaFilter rsFilter = new QuotaFilter().setRegionServerFilter(regionServer);
532
533    admin.setQuota(QuotaSettingsFactory.throttleRegionServer(regionServer,
534      ThrottleType.REQUEST_NUMBER, 10, TimeUnit.MINUTES));
535    assertNumResults(1, rsFilter);
536    // Verify the Quota in the table
537    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_NUMBER, 10, TimeUnit.MINUTES);
538
539    admin.setQuota(QuotaSettingsFactory.throttleRegionServer(regionServer,
540      ThrottleType.REQUEST_NUMBER, 20, TimeUnit.MINUTES));
541    assertNumResults(1, rsFilter);
542    // Verify the Quota in the table
543    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_NUMBER, 20, TimeUnit.MINUTES);
544
545    admin.setQuota(QuotaSettingsFactory.throttleRegionServer(regionServer, ThrottleType.READ_NUMBER,
546      30, TimeUnit.SECONDS));
547    int count = 0;
548    QuotaRetriever scanner = QuotaRetriever.open(TEST_UTIL.getConfiguration(), rsFilter);
549    try {
550      for (QuotaSettings settings : scanner) {
551        assertTrue(settings.getQuotaType() == QuotaType.THROTTLE);
552        ThrottleSettings throttleSettings = (ThrottleSettings) settings;
553        assertEquals(regionServer, throttleSettings.getRegionServer());
554        count++;
555        if (throttleSettings.getThrottleType() == ThrottleType.REQUEST_NUMBER) {
556          assertEquals(20, throttleSettings.getSoftLimit());
557          assertEquals(TimeUnit.MINUTES, throttleSettings.getTimeUnit());
558        } else if (throttleSettings.getThrottleType() == ThrottleType.READ_NUMBER) {
559          assertEquals(30, throttleSettings.getSoftLimit());
560          assertEquals(TimeUnit.SECONDS, throttleSettings.getTimeUnit());
561        }
562      }
563    } finally {
564      scanner.close();
565    }
566    assertEquals(2, count);
567
568    admin.setQuota(QuotaSettingsFactory.unthrottleRegionServer(regionServer));
569    assertNumResults(0, new QuotaFilter().setRegionServerFilter(regionServer));
570  }
571
572  @Test
573  public void testRpcThrottleWhenStartup() throws IOException, InterruptedException {
574    TEST_UTIL.getAdmin().switchRpcThrottle(false);
575    assertFalse(TEST_UTIL.getAdmin().isRpcThrottleEnabled());
576    TEST_UTIL.killMiniHBaseCluster();
577
578    TEST_UTIL.startMiniHBaseCluster();
579    assertFalse(TEST_UTIL.getAdmin().isRpcThrottleEnabled());
580    for (JVMClusterUtil.RegionServerThread rs : TEST_UTIL.getHBaseCluster()
581        .getRegionServerThreads()) {
582      RegionServerRpcQuotaManager quotaManager =
583          rs.getRegionServer().getRegionServerRpcQuotaManager();
584      assertFalse(quotaManager.isRpcThrottleEnabled());
585    }
586    // enable rpc throttle
587    TEST_UTIL.getAdmin().switchRpcThrottle(true);
588    assertTrue(TEST_UTIL.getAdmin().isRpcThrottleEnabled());
589  }
590
591  @Test
592  public void testSwitchRpcThrottle() throws IOException {
593    Admin admin = TEST_UTIL.getAdmin();
594    testSwitchRpcThrottle(admin, true, true);
595    testSwitchRpcThrottle(admin, true, false);
596    testSwitchRpcThrottle(admin, false, false);
597    testSwitchRpcThrottle(admin, false, true);
598  }
599
600  @Test
601  public void testSwitchExceedThrottleQuota() throws IOException {
602    String regionServer = QuotaTableUtil.QUOTA_REGION_SERVER_ROW_KEY;
603    Admin admin = TEST_UTIL.getAdmin();
604
605    try {
606      admin.exceedThrottleQuotaSwitch(true);
607      fail("should not come here, because can't enable exceed throttle quota "
608          + "if there is no region server quota");
609    } catch (IOException e) {
610      LOG.warn("Expected exception", e);
611    }
612
613    admin.setQuota(QuotaSettingsFactory.throttleRegionServer(regionServer,
614      ThrottleType.WRITE_NUMBER, 100, TimeUnit.SECONDS));
615    try {
616      admin.exceedThrottleQuotaSwitch(true);
617      fail("should not come here, because can't enable exceed throttle quota "
618          + "if there is no read region server quota");
619    } catch (IOException e) {
620      LOG.warn("Expected exception", e);
621    }
622
623    admin.setQuota(QuotaSettingsFactory.throttleRegionServer(regionServer, ThrottleType.READ_NUMBER,
624      20, TimeUnit.MINUTES));
625    try {
626      admin.exceedThrottleQuotaSwitch(true);
627      fail("should not come here, because can't enable exceed throttle quota "
628          + "because not all region server quota are in seconds time unit");
629    } catch (IOException e) {
630      LOG.warn("Expected exception", e);
631    }
632    admin.setQuota(QuotaSettingsFactory.throttleRegionServer(regionServer, ThrottleType.READ_NUMBER,
633      20, TimeUnit.SECONDS));
634
635    assertFalse(admin.exceedThrottleQuotaSwitch(true));
636    assertTrue(admin.exceedThrottleQuotaSwitch(true));
637    assertTrue(admin.exceedThrottleQuotaSwitch(false));
638    assertFalse(admin.exceedThrottleQuotaSwitch(false));
639    assertEquals(2, admin.getQuota(new QuotaFilter()).size());
640    admin.setQuota(QuotaSettingsFactory.unthrottleRegionServer(regionServer));
641  }
642
643  @Test
644  public void testQuotaScope() throws Exception {
645    Admin admin = TEST_UTIL.getAdmin();
646    String user = "user1";
647    String namespace = "testQuotaScope_ns";
648    TableName tableName = TableName.valueOf("testQuotaScope");
649    QuotaFilter filter = new QuotaFilter();
650
651    // set CLUSTER quota scope for namespace
652    admin.setQuota(QuotaSettingsFactory.throttleNamespace(namespace, ThrottleType.REQUEST_NUMBER,
653      10, TimeUnit.MINUTES, QuotaScope.CLUSTER));
654    assertNumResults(1, filter);
655    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_NUMBER, 10, TimeUnit.MINUTES,
656      QuotaScope.CLUSTER);
657    admin.setQuota(QuotaSettingsFactory.throttleNamespace(namespace, ThrottleType.REQUEST_NUMBER,
658      10, TimeUnit.MINUTES, QuotaScope.MACHINE));
659    assertNumResults(1, filter);
660    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_NUMBER, 10, TimeUnit.MINUTES,
661      QuotaScope.MACHINE);
662    admin.setQuota(QuotaSettingsFactory.unthrottleNamespace(namespace));
663    assertNumResults(0, filter);
664
665    // set CLUSTER quota scope for table
666    admin.setQuota(QuotaSettingsFactory.throttleTable(tableName, ThrottleType.REQUEST_NUMBER, 10,
667      TimeUnit.MINUTES, QuotaScope.CLUSTER));
668    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_NUMBER, 10, TimeUnit.MINUTES,
669      QuotaScope.CLUSTER);
670    admin.setQuota(QuotaSettingsFactory.unthrottleTable(tableName));
671
672    // set CLUSTER quota scope for user
673    admin.setQuota(QuotaSettingsFactory.throttleUser(user, ThrottleType.REQUEST_NUMBER, 10,
674      TimeUnit.MINUTES, QuotaScope.CLUSTER));
675    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_NUMBER, 10, TimeUnit.MINUTES,
676      QuotaScope.CLUSTER);
677    admin.setQuota(QuotaSettingsFactory.unthrottleUser(user));
678
679      // set CLUSTER quota scope for user and table
680    admin.setQuota(QuotaSettingsFactory.throttleUser(user, tableName, ThrottleType.REQUEST_NUMBER,
681      10, TimeUnit.MINUTES, QuotaScope.CLUSTER));
682    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_NUMBER, 10, TimeUnit.MINUTES,
683      QuotaScope.CLUSTER);
684    admin.setQuota(QuotaSettingsFactory.unthrottleUser(user));
685
686      // set CLUSTER quota scope for user and namespace
687    admin.setQuota(QuotaSettingsFactory.throttleUser(user, namespace, ThrottleType.REQUEST_NUMBER,
688      10, TimeUnit.MINUTES, QuotaScope.CLUSTER));
689    verifyRecordPresentInQuotaTable(ThrottleType.REQUEST_NUMBER, 10, TimeUnit.MINUTES,
690      QuotaScope.CLUSTER);
691    admin.setQuota(QuotaSettingsFactory.unthrottleUser(user));
692  }
693
694  private void testSwitchRpcThrottle(Admin admin, boolean oldRpcThrottle, boolean newRpcThrottle)
695      throws IOException {
696    boolean state = admin.switchRpcThrottle(newRpcThrottle);
697    Assert.assertEquals(oldRpcThrottle, state);
698    Assert.assertEquals(newRpcThrottle, admin.isRpcThrottleEnabled());
699    TEST_UTIL.getHBaseCluster().getRegionServerThreads().stream()
700        .forEach(rs -> Assert.assertEquals(newRpcThrottle,
701          rs.getRegionServer().getRegionServerRpcQuotaManager().isRpcThrottleEnabled()));
702  }
703
704  private void verifyRecordPresentInQuotaTable(ThrottleType type, long limit, TimeUnit tu)
705      throws Exception {
706    verifyRecordPresentInQuotaTable(type, limit, tu, QuotaScope.MACHINE);
707  }
708
709  private void verifyRecordPresentInQuotaTable(ThrottleType type, long limit, TimeUnit tu,
710      QuotaScope scope) throws Exception {
711    // Verify the RPC Quotas in the table
712    try (Table quotaTable = TEST_UTIL.getConnection().getTable(QuotaTableUtil.QUOTA_TABLE_NAME);
713        ResultScanner scanner = quotaTable.getScanner(new Scan())) {
714      Result r = Iterables.getOnlyElement(scanner);
715      CellScanner cells = r.cellScanner();
716      assertTrue("Expected to find a cell", cells.advance());
717      assertRPCQuota(type, limit, tu, scope, cells.current());
718    }
719  }
720
721  private void verifyRecordNotPresentInQuotaTable() throws Exception {
722    // Verify that the record doesn't exist in the QuotaTableUtil.QUOTA_TABLE_NAME
723    try (Table quotaTable = TEST_UTIL.getConnection().getTable(QuotaTableUtil.QUOTA_TABLE_NAME);
724        ResultScanner scanner = quotaTable.getScanner(new Scan())) {
725      assertNull("Did not expect to find a quota entry", scanner.next());
726    }
727  }
728
729  private void verifyFetchableViaAPI(Admin admin, ThrottleType type, long limit, TimeUnit tu)
730      throws Exception {
731    // Verify we can retrieve the new quota via the QuotaRetriever API
732    try (QuotaRetriever quotaScanner = QuotaRetriever.open(admin.getConfiguration())) {
733      assertRPCQuota(type, limit, tu, Iterables.getOnlyElement(quotaScanner));
734    }
735  }
736
737  private void verifyNotFetchableViaAPI(Admin admin) throws Exception {
738    // Verify that we can also not fetch it via the API
739    try (QuotaRetriever quotaScanner = QuotaRetriever.open(admin.getConfiguration())) {
740      assertNull("Did not expect to find a quota entry", quotaScanner.next());
741    }
742  }
743
744  private void assertRPCQuota(ThrottleType type, long limit, TimeUnit tu, QuotaScope scope,
745      Cell cell) throws Exception {
746    Quotas q = QuotaTableUtil
747        .quotasFromData(cell.getValueArray(), cell.getValueOffset(), cell.getValueLength());
748    assertTrue("Quota should have rpc quota defined", q.hasThrottle());
749
750    QuotaProtos.Throttle rpcQuota = q.getThrottle();
751    QuotaProtos.TimedQuota t = null;
752
753    switch (type) {
754      case REQUEST_SIZE:
755        assertTrue(rpcQuota.hasReqSize());
756        t = rpcQuota.getReqSize();
757        break;
758      case READ_NUMBER:
759        assertTrue(rpcQuota.hasReadNum());
760        t = rpcQuota.getReadNum();
761        break;
762      case READ_SIZE:
763        assertTrue(rpcQuota.hasReadSize());
764        t = rpcQuota.getReadSize();
765        break;
766      case REQUEST_NUMBER:
767        assertTrue(rpcQuota.hasReqNum());
768        t = rpcQuota.getReqNum();
769        break;
770      case WRITE_NUMBER:
771        assertTrue(rpcQuota.hasWriteNum());
772        t = rpcQuota.getWriteNum();
773        break;
774      case WRITE_SIZE:
775        assertTrue(rpcQuota.hasWriteSize());
776        t = rpcQuota.getWriteSize();
777        break;
778      case REQUEST_CAPACITY_UNIT:
779        assertTrue(rpcQuota.hasReqCapacityUnit());
780        t = rpcQuota.getReqCapacityUnit();
781        break;
782      case READ_CAPACITY_UNIT:
783        assertTrue(rpcQuota.hasReadCapacityUnit());
784        t = rpcQuota.getReadCapacityUnit();
785        break;
786      case WRITE_CAPACITY_UNIT:
787        assertTrue(rpcQuota.hasWriteCapacityUnit());
788        t = rpcQuota.getWriteCapacityUnit();
789        break;
790      default:
791    }
792
793    assertEquals(scope, ProtobufUtil.toQuotaScope(t.getScope()));
794    assertEquals(t.getSoftLimit(), limit);
795    assertEquals(t.getTimeUnit(), ProtobufUtil.toProtoTimeUnit(tu));
796  }
797
798  private void assertRPCQuota(ThrottleType type, long limit, TimeUnit tu,
799      QuotaSettings actualSettings) throws Exception {
800    assertTrue(
801        "The actual QuotaSettings was not an instance of " + ThrottleSettings.class + " but of "
802            + actualSettings.getClass(), actualSettings instanceof ThrottleSettings);
803    QuotaProtos.ThrottleRequest throttleRequest = ((ThrottleSettings) actualSettings).getProto();
804    assertEquals(limit, throttleRequest.getTimedQuota().getSoftLimit());
805    assertEquals(ProtobufUtil.toProtoTimeUnit(tu), throttleRequest.getTimedQuota().getTimeUnit());
806    assertEquals(ProtobufUtil.toProtoThrottleType(type), throttleRequest.getType());
807  }
808
809  private void assertSpaceQuota(
810      long sizeLimit, SpaceViolationPolicy violationPolicy, Cell cell) throws Exception {
811    Quotas q = QuotaTableUtil.quotasFromData(
812        cell.getValueArray(), cell.getValueOffset(), cell.getValueLength());
813    assertTrue("Quota should have space quota defined", q.hasSpace());
814    QuotaProtos.SpaceQuota spaceQuota = q.getSpace();
815    assertEquals(sizeLimit, spaceQuota.getSoftLimit());
816    assertEquals(violationPolicy, ProtobufUtil.toViolationPolicy(spaceQuota.getViolationPolicy()));
817  }
818
819  private void assertSpaceQuota(
820      long sizeLimit, SpaceViolationPolicy violationPolicy, QuotaSettings actualSettings) {
821    assertTrue("The actual QuotaSettings was not an instance of " + SpaceLimitSettings.class
822        + " but of " + actualSettings.getClass(), actualSettings instanceof SpaceLimitSettings);
823    SpaceLimitRequest spaceLimitRequest = ((SpaceLimitSettings) actualSettings).getProto();
824    assertEquals(sizeLimit, spaceLimitRequest.getQuota().getSoftLimit());
825    assertEquals(violationPolicy,
826        ProtobufUtil.toViolationPolicy(spaceLimitRequest.getQuota().getViolationPolicy()));
827  }
828
829  private int countResults(final QuotaFilter filter) throws Exception {
830    QuotaRetriever scanner = QuotaRetriever.open(TEST_UTIL.getConfiguration(), filter);
831    try {
832      int count = 0;
833      for (QuotaSettings settings: scanner) {
834        LOG.debug(Objects.toString(settings));
835        count++;
836      }
837      return count;
838    } finally {
839      scanner.close();
840    }
841  }
842}