Source code

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.jupiter.api.Assertions.assertEquals;
021import static org.junit.jupiter.api.Assertions.assertInstanceOf;
022import static org.junit.jupiter.api.Assertions.assertThrows;
023
024import java.io.IOException;
025import java.util.ArrayList;
026import java.util.List;
027import java.util.UUID;
028import java.util.concurrent.TimeUnit;
029import org.apache.hadoop.hbase.HBaseTestingUtil;
030import org.apache.hadoop.hbase.HConstants;
031import org.apache.hadoop.hbase.TableName;
032import org.apache.hadoop.hbase.client.Admin;
033import org.apache.hadoop.hbase.client.CheckAndMutate;
034import org.apache.hadoop.hbase.client.Get;
035import org.apache.hadoop.hbase.client.Increment;
036import org.apache.hadoop.hbase.client.Put;
037import org.apache.hadoop.hbase.client.RowMutations;
038import org.apache.hadoop.hbase.client.Scan;
039import org.apache.hadoop.hbase.client.Table;
040import org.apache.hadoop.hbase.ipc.HBaseRpcControllerImpl;
041import org.apache.hadoop.hbase.regionserver.RSRpcServices;
042import org.apache.hadoop.hbase.security.User;
043import org.apache.hadoop.hbase.testclassification.MediumTests;
044import org.apache.hadoop.hbase.testclassification.RegionServerTests;
045import org.apache.hadoop.hbase.util.Bytes;
046import org.apache.hadoop.hbase.util.EnvironmentEdgeManager;
047import org.junit.jupiter.api.AfterAll;
048import org.junit.jupiter.api.BeforeAll;
049import org.junit.jupiter.api.Tag;
050import org.junit.jupiter.api.Test;
051import org.slf4j.Logger;
052import org.slf4j.LoggerFactory;
053
054import org.apache.hbase.thirdparty.com.google.protobuf.ServiceException;
055
056import org.apache.hadoop.hbase.shaded.protobuf.RequestConverter;
057import org.apache.hadoop.hbase.shaded.protobuf.generated.ClientProtos.ScanResponse;
058
059@Tag(RegionServerTests.TAG)
060@Tag(MediumTests.TAG)
061public class TestAtomicReadQuota {
062  private static final Logger LOG = LoggerFactory.getLogger(TestAtomicReadQuota.class);
063  private static final HBaseTestingUtil TEST_UTIL = new HBaseTestingUtil();
064  private static final TableName TABLE_NAME = TableName.valueOf(UUID.randomUUID().toString());
065  private static final byte[] FAMILY = Bytes.toBytes("cf");
066  private static final byte[] QUALIFIER = Bytes.toBytes("q");
067
068  @AfterAll
069  public static void tearDown() throws Exception {
070    ThrottleQuotaTestUtil.clearQuotaCache(TEST_UTIL);
071    EnvironmentEdgeManager.reset();
072    TEST_UTIL.deleteTable(TABLE_NAME);
073    TEST_UTIL.shutdownMiniCluster();
074  }
075
076  @BeforeAll
077  public static void setUpBeforeClass() throws Exception {
078    TEST_UTIL.getConfiguration().setBoolean(QuotaUtil.QUOTA_CONF_KEY, true);
079    TEST_UTIL.getConfiguration().setInt(QuotaCache.REFRESH_CONF_KEY, 1000);
080    TEST_UTIL.startMiniCluster(1);
081    TEST_UTIL.waitTableAvailable(QuotaTableUtil.QUOTA_TABLE_NAME);
082    TEST_UTIL.createTable(TABLE_NAME, FAMILY);
083    TEST_UTIL.waitTableAvailable(TABLE_NAME);
084  }
085
086  @Test
087  public void testCloseScannerWhenReadQuotaIsExhausted() throws Exception {
088    setupGenericQuota();
089    try {
090      RSRpcServices rpcServices = TEST_UTIL.getHBaseCluster().getRegionServer(0).getRSRpcServices();
091      int scannerCount = rpcServices.getScannersCount();
092      byte[] regionName = TEST_UTIL.getAdmin().getRegions(TABLE_NAME).get(0).getRegionName();
093      ScanResponse opened = rpcServices.scan(new HBaseRpcControllerImpl(),
094        RequestConverter.buildScanRequest(regionName, new Scan(), 0, false));
095      long scannerId = opened.getScannerId();
096      assertEquals(scannerCount + 1, rpcServices.getScannersCount());
097
098      ServiceException throttled =
099        assertThrows(ServiceException.class, () -> rpcServices.scan(new HBaseRpcControllerImpl(),
100          RequestConverter.buildScanRequest(scannerId, 1, true, false)));
101      assertInstanceOf(RpcThrottlingException.class, throttled.getCause());
102      assertEquals(scannerCount + 1, rpcServices.getScannersCount());
103
104      ServiceException renewThrottled =
105        assertThrows(ServiceException.class, () -> rpcServices.scan(new HBaseRpcControllerImpl(),
106          RequestConverter.buildScanRequest(scannerId, 0, true, 0, false, true, -1)));
107      assertInstanceOf(RpcThrottlingException.class, renewThrottled.getCause());
108
109      ServiceException openThrottled =
110        assertThrows(ServiceException.class, () -> rpcServices.scan(new HBaseRpcControllerImpl(),
111          RequestConverter.buildScanRequest(regionName, new Scan(), 0, true)));
112      assertInstanceOf(RpcThrottlingException.class, openThrottled.getCause());
113
114      rpcServices.scan(new HBaseRpcControllerImpl(),
115        RequestConverter.buildScanRequest(scannerId, 0, true, false));
116      assertEquals(scannerCount, rpcServices.getScannersCount());
117    } finally {
118      cleanupQuota();
119    }
120  }
121
122  @Test
123  public void testIncrementCountedAgainstReadCapacity() throws Exception {
124    setupGenericQuota();
125
126    Increment inc = new Increment(Bytes.toBytes(UUID.randomUUID().toString()));
127    inc.addColumn(FAMILY, QUALIFIER, 1);
128    testThrottle(table -> table.increment(inc));
129  }
130
131  @Test
132  public void testConditionalRowMutationsCountedAgainstReadCapacity() throws Exception {
133    setupGenericQuota();
134
135    byte[] row = Bytes.toBytes(UUID.randomUUID().toString());
136    Increment inc = new Increment(row);
137    inc.addColumn(FAMILY, Bytes.toBytes("doot"), 1);
138    Put put = new Put(row);
139    put.addColumn(FAMILY, Bytes.toBytes("doot"), Bytes.toBytes("v"));
140
141    RowMutations rowMutations = new RowMutations(row);
142    rowMutations.add(inc);
143    rowMutations.add(put);
144    testThrottle(table -> table.mutateRow(rowMutations));
145  }
146
147  @Test
148  public void testNonConditionalRowMutationsOmittedFromReadCapacity() throws Exception {
149    setupGenericQuota();
150
151    byte[] row = Bytes.toBytes(UUID.randomUUID().toString());
152    Put put = new Put(row);
153    put.addColumn(FAMILY, Bytes.toBytes("doot"), Bytes.toBytes("v"));
154
155    RowMutations rowMutations = new RowMutations(row);
156    rowMutations.add(put);
157    try (Table table = getTable()) {
158      for (int i = 0; i < 100; i++) {
159        table.mutateRow(rowMutations);
160      }
161    }
162  }
163
164  @Test
165  public void testNonAtomicPutOmittedFromReadCapacity() throws Exception {
166    setupGenericQuota();
167    runNonAtomicPuts();
168  }
169
170  @Test
171  public void testNonAtomicMultiPutOmittedFromReadCapacity() throws Exception {
172    setupGenericQuota();
173    runNonAtomicPuts();
174  }
175
176  @Test
177  public void testCheckAndMutateCountedAgainstReadCapacity() throws Exception {
178    setupGenericQuota();
179
180    byte[] row = Bytes.toBytes(UUID.randomUUID().toString());
181    byte[] value = Bytes.toBytes("v");
182    Put put = new Put(row);
183    put.addColumn(FAMILY, Bytes.toBytes("doot"), value);
184    CheckAndMutate checkAndMutate =
185      CheckAndMutate.newBuilder(row).ifEquals(FAMILY, QUALIFIER, value).build(put);
186
187    testThrottle(table -> table.checkAndMutate(checkAndMutate));
188  }
189
190  @Test
191  public void testAtomicBatchCountedAgainstReadCapacity() throws Exception {
192    setupGenericQuota();
193
194    byte[] row = Bytes.toBytes(UUID.randomUUID().toString());
195    Increment inc = new Increment(row);
196    inc.addColumn(FAMILY, Bytes.toBytes("doot"), 1);
197
198    List<Increment> incs = new ArrayList<>(2);
199    incs.add(inc);
200    incs.add(inc);
201
202    testThrottle(table -> {
203      Object[] results = new Object[] {};
204      table.batch(incs, results);
205      return results;
206    });
207  }
208
209  @Test
210  public void testAtomicBatchCountedAgainstAtomicOnlyReqNum() throws Exception {
211    setupAtomicOnlyReqNumQuota();
212
213    byte[] row = Bytes.toBytes(UUID.randomUUID().toString());
214    Increment inc = new Increment(row);
215    inc.addColumn(FAMILY, Bytes.toBytes("doot"), 1);
216
217    List<Increment> incs = new ArrayList<>(2);
218    incs.add(inc);
219    incs.add(inc);
220
221    testThrottle(table -> {
222      Object[] results = new Object[] {};
223      table.batch(incs, results);
224      return results;
225    });
226  }
227
228  @Test
229  public void testAtomicBatchCountedAgainstAtomicOnlyReadSize() throws Exception {
230    setupAtomicOnlyReadSizeQuota();
231
232    byte[] row = Bytes.toBytes(UUID.randomUUID().toString());
233    Increment inc = new Increment(row);
234    inc.addColumn(FAMILY, Bytes.toBytes("doot"), 1);
235
236    List<Increment> incs = new ArrayList<>(2);
237    incs.add(inc);
238    incs.add(inc);
239
240    testThrottle(table -> {
241      Object[] results = new Object[] {};
242      table.batch(incs, results);
243      return results;
244    });
245  }
246
247  @Test
248  public void testNonAtomicWritesIgnoredByAtomicOnlyReqNum() throws Exception {
249    setupAtomicOnlyReqNumQuota();
250    runNonAtomicPuts();
251  }
252
253  @Test
254  public void testNonAtomicWritesIgnoredByAtomicOnlyReadSize() throws Exception {
255    setupAtomicOnlyReadSizeQuota();
256    runNonAtomicPuts();
257  }
258
259  @Test
260  public void testNonAtomicReadsIgnoredByAtomicOnlyReqNum() throws Exception {
261    setupAtomicOnlyReqNumQuota();
262    runNonAtomicReads();
263  }
264
265  @Test
266  public void testNonAtomicReadsIgnoredByAtomicOnlyReadSize() throws Exception {
267    setupAtomicOnlyReadSizeQuota();
268    runNonAtomicReads();
269  }
270
271  private void runNonAtomicPuts() throws Exception {
272    Put put1 = new Put(Bytes.toBytes(UUID.randomUUID().toString()));
273    put1.addColumn(FAMILY, Bytes.toBytes("doot"), Bytes.toBytes("v"));
274    Put put2 = new Put(Bytes.toBytes(UUID.randomUUID().toString()));
275    put2.addColumn(FAMILY, Bytes.toBytes("doot"), Bytes.toBytes("v"));
276
277    Increment inc = new Increment(Bytes.toBytes(UUID.randomUUID().toString()));
278    inc.addColumn(FAMILY, Bytes.toBytes("doot"), 1);
279
280    List<Put> puts = new ArrayList<>(2);
281    puts.add(put1);
282    puts.add(put2);
283
284    try (Table table = getTable()) {
285      for (int i = 0; i < 100; i++) {
286        table.put(puts);
287      }
288    }
289  }
290
291  private void runNonAtomicReads() throws Exception {
292    try (Table table = getTable()) {
293      byte[] row = Bytes.toBytes(UUID.randomUUID().toString());
294      Get get = new Get(row);
295      table.get(get);
296    }
297  }
298
299  private void setupGenericQuota() throws Exception {
300    try (Admin admin = TEST_UTIL.getAdmin()) {
301      admin.setQuota(QuotaSettingsFactory.throttleUser(User.getCurrent().getShortName(),
302        ThrottleType.READ_NUMBER, 1, TimeUnit.MINUTES));
303    }
304    ThrottleQuotaTestUtil.triggerUserCacheRefresh(TEST_UTIL, false, TABLE_NAME);
305  }
306
307  private void setupAtomicOnlyReqNumQuota() throws Exception {
308    try (Admin admin = TEST_UTIL.getAdmin()) {
309      admin.setQuota(QuotaSettingsFactory.throttleUser(User.getCurrent().getShortName(),
310        ThrottleType.ATOMIC_REQUEST_NUMBER, 1, TimeUnit.MINUTES));
311    }
312    ThrottleQuotaTestUtil.triggerUserCacheRefresh(TEST_UTIL, false, TABLE_NAME);
313  }
314
315  private void setupAtomicOnlyReadSizeQuota() throws Exception {
316    try (Admin admin = TEST_UTIL.getAdmin()) {
317      admin.setQuota(QuotaSettingsFactory.throttleUser(User.getCurrent().getShortName(),
318        ThrottleType.ATOMIC_READ_SIZE, 1, TimeUnit.MINUTES));
319    }
320    ThrottleQuotaTestUtil.triggerUserCacheRefresh(TEST_UTIL, false, TABLE_NAME);
321  }
322
323  private void cleanupQuota() throws Exception {
324    try (Admin admin = TEST_UTIL.getAdmin()) {
325      admin.setQuota(QuotaSettingsFactory.unthrottleUser(User.getCurrent().getShortName()));
326    }
327    ThrottleQuotaTestUtil.triggerUserCacheRefresh(TEST_UTIL, true, TABLE_NAME);
328  }
329
330  private void testThrottle(ThrowingFunction<Table, ?> request) throws Exception {
331    try (Table table = getTable()) {
332      // we have a read quota configured, so this should fail
333      TEST_UTIL.waitFor(60_000, () -> {
334        try {
335          request.run(table);
336          return false;
337        } catch (Exception e) {
338          boolean success = e.getCause() instanceof RpcThrottlingException;
339          if (!success) {
340            LOG.error("Unexpected exception", e);
341          }
342          return success;
343        }
344      });
345    } finally {
346      cleanupQuota();
347    }
348  }
349
350  private Table getTable() throws IOException {
351    TEST_UTIL.getConfiguration().setInt("hbase.client.pause", 100);
352    TEST_UTIL.getConfiguration().setInt(HConstants.HBASE_CLIENT_RETRIES_NUMBER, 1);
353    return TEST_UTIL.getConnection().getTableBuilder(TABLE_NAME, null).setOperationTimeout(250)
354      .build();
355  }
356
357  @FunctionalInterface
358  private interface ThrowingFunction<I, O> {
359    O run(I input) throws Exception;
360  }
361
362}