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.assertFalse;
021import static org.junit.jupiter.api.Assertions.assertTrue;
022import static org.junit.jupiter.api.Assertions.fail;
023
024import java.util.concurrent.TimeUnit;
025import org.apache.hadoop.conf.Configuration;
026import org.apache.hadoop.hbase.HBaseConfiguration;
027import org.apache.hadoop.hbase.TableName;
028import org.apache.hadoop.hbase.testclassification.RegionServerTests;
029import org.apache.hadoop.hbase.testclassification.SmallTests;
030import org.junit.jupiter.api.Tag;
031import org.junit.jupiter.api.Test;
032import org.junit.jupiter.api.TestInfo;
033
034import org.apache.hadoop.hbase.shaded.protobuf.ProtobufUtil;
035import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.Quotas;
036import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.Throttle;
037
038@Tag(RegionServerTests.TAG)
039@Tag(SmallTests.TAG)
040public class TestQuotaState {
041
042  private static final TableName UNKNOWN_TABLE_NAME = TableName.valueOf("unknownTable");
043
044  private static final Configuration conf = HBaseConfiguration.create();
045
046  @Test
047  public void testQuotaStateBypass() {
048    QuotaState quotaInfo = new QuotaState();
049    assertTrue(quotaInfo.isBypass());
050    assertNoopLimiter(quotaInfo.getGlobalLimiter());
051
052    UserQuotaState userQuotaState = new UserQuotaState();
053    assertTrue(userQuotaState.isBypass());
054    assertNoopLimiter(userQuotaState.getTableLimiter(UNKNOWN_TABLE_NAME));
055  }
056
057  @Test
058  public void testSimpleQuotaStateOperation(TestInfo testInfo) {
059    final TableName tableName = TableName.valueOf(testInfo.getTestMethod().get().getName());
060    final int NUM_GLOBAL_THROTTLE = 3;
061    final int NUM_TABLE_THROTTLE = 2;
062
063    UserQuotaState quotaInfo = new UserQuotaState();
064    assertTrue(quotaInfo.isBypass());
065
066    // Set global quota
067    quotaInfo.setQuotas(conf, buildReqNumThrottle(NUM_GLOBAL_THROTTLE));
068    assertFalse(quotaInfo.isBypass());
069
070    // Set table quota
071    quotaInfo.setQuotas(conf, tableName, buildReqNumThrottle(NUM_TABLE_THROTTLE));
072    assertFalse(quotaInfo.isBypass());
073    assertTrue(quotaInfo.getGlobalLimiter() == quotaInfo.getTableLimiter(UNKNOWN_TABLE_NAME));
074    assertThrottleException(quotaInfo.getTableLimiter(UNKNOWN_TABLE_NAME), NUM_GLOBAL_THROTTLE);
075    assertThrottleException(quotaInfo.getTableLimiter(tableName), NUM_TABLE_THROTTLE);
076  }
077
078  @Test
079  public void testQuotaStateUpdateGlobalThrottle() {
080    final int NUM_GLOBAL_THROTTLE_1 = 3;
081    final int NUM_GLOBAL_THROTTLE_2 = 11;
082
083    QuotaState quotaInfo = new QuotaState();
084    assertTrue(quotaInfo.isBypass());
085
086    // Add global throttle
087    QuotaState otherQuotaState = new QuotaState();
088    otherQuotaState.setQuotas(conf, buildReqNumThrottle(NUM_GLOBAL_THROTTLE_1));
089    assertFalse(otherQuotaState.isBypass());
090
091    quotaInfo.update(otherQuotaState);
092    assertFalse(quotaInfo.isBypass());
093    assertThrottleException(quotaInfo.getGlobalLimiter(), NUM_GLOBAL_THROTTLE_1);
094
095    // Update global Throttle
096    otherQuotaState = new QuotaState();
097    otherQuotaState.setQuotas(conf, buildReqNumThrottle(NUM_GLOBAL_THROTTLE_2));
098    assertFalse(otherQuotaState.isBypass());
099
100    quotaInfo.update(otherQuotaState);
101    assertFalse(quotaInfo.isBypass());
102    assertThrottleException(quotaInfo.getGlobalLimiter(),
103      NUM_GLOBAL_THROTTLE_2 - NUM_GLOBAL_THROTTLE_1);
104
105    // Remove global throttle
106    otherQuotaState = new QuotaState();
107    assertTrue(otherQuotaState.isBypass());
108
109    quotaInfo.update(otherQuotaState);
110    assertTrue(quotaInfo.isBypass());
111    assertNoopLimiter(quotaInfo.getGlobalLimiter());
112  }
113
114  @Test
115  public void testQuotaStateUpdateTableThrottle(TestInfo testInfo) {
116    final String methodName = testInfo.getTestMethod().get().getName();
117    final TableName tableNameA = TableName.valueOf(methodName + "A");
118    final TableName tableNameB = TableName.valueOf(methodName + "B");
119    final TableName tableNameC = TableName.valueOf(methodName + "C");
120    final int TABLE_A_THROTTLE_1 = 3;
121    final int TABLE_A_THROTTLE_2 = 11;
122    final int TABLE_B_THROTTLE = 4;
123    final int TABLE_C_THROTTLE = 5;
124
125    UserQuotaState quotaInfo = new UserQuotaState();
126    assertTrue(quotaInfo.isBypass());
127
128    // Add A B table limiters
129    UserQuotaState otherQuotaState = new UserQuotaState();
130    otherQuotaState.setQuotas(conf, tableNameA, buildReqNumThrottle(TABLE_A_THROTTLE_1));
131    otherQuotaState.setQuotas(conf, tableNameB, buildReqNumThrottle(TABLE_B_THROTTLE));
132    assertFalse(otherQuotaState.isBypass());
133
134    quotaInfo.update(otherQuotaState);
135    assertFalse(quotaInfo.isBypass());
136    assertThrottleException(quotaInfo.getTableLimiter(tableNameA), TABLE_A_THROTTLE_1);
137    assertThrottleException(quotaInfo.getTableLimiter(tableNameB), TABLE_B_THROTTLE);
138    assertNoopLimiter(quotaInfo.getTableLimiter(tableNameC));
139
140    // Add C, Remove B, Update A table limiters
141    otherQuotaState = new UserQuotaState();
142    otherQuotaState.setQuotas(conf, tableNameA, buildReqNumThrottle(TABLE_A_THROTTLE_2));
143    otherQuotaState.setQuotas(conf, tableNameC, buildReqNumThrottle(TABLE_C_THROTTLE));
144    assertFalse(otherQuotaState.isBypass());
145
146    quotaInfo.update(otherQuotaState);
147    assertFalse(quotaInfo.isBypass());
148    assertThrottleException(quotaInfo.getTableLimiter(tableNameA),
149      TABLE_A_THROTTLE_2 - TABLE_A_THROTTLE_1);
150    assertThrottleException(quotaInfo.getTableLimiter(tableNameC), TABLE_C_THROTTLE);
151    assertNoopLimiter(quotaInfo.getTableLimiter(tableNameB));
152
153    // Remove table limiters
154    otherQuotaState = new UserQuotaState();
155    assertTrue(otherQuotaState.isBypass());
156
157    quotaInfo.update(otherQuotaState);
158    assertTrue(quotaInfo.isBypass());
159    assertNoopLimiter(quotaInfo.getTableLimiter(UNKNOWN_TABLE_NAME));
160  }
161
162  @Test
163  public void testTableThrottleWithBatch() {
164    final TableName TABLE_A = TableName.valueOf("TableA");
165    final int TABLE_A_THROTTLE_1 = 3;
166
167    UserQuotaState quotaInfo = new UserQuotaState();
168    assertTrue(quotaInfo.isBypass());
169
170    // Add A table limiters
171    UserQuotaState otherQuotaState = new UserQuotaState();
172    otherQuotaState.setQuotas(conf, TABLE_A, buildReqNumThrottle(TABLE_A_THROTTLE_1));
173    assertFalse(otherQuotaState.isBypass());
174
175    quotaInfo.update(otherQuotaState);
176    assertFalse(quotaInfo.isBypass());
177    QuotaLimiter limiter = quotaInfo.getTableLimiter(TABLE_A);
178    try {
179      limiter.checkQuota(TABLE_A_THROTTLE_1 + 1, TABLE_A_THROTTLE_1 + 1, 0, 0, 1, 0, false, 0L);
180      fail("Should have thrown RpcThrottlingException");
181    } catch (RpcThrottlingException e) {
182      // expected
183    }
184  }
185
186  private Quotas buildReqNumThrottle(final long limit) {
187    return Quotas.newBuilder()
188      .setThrottle(Throttle.newBuilder()
189        .setReqNum(ProtobufUtil.toTimedQuota(limit, TimeUnit.MINUTES, QuotaScope.MACHINE)).build())
190      .build();
191  }
192
193  private void assertThrottleException(final QuotaLimiter limiter, final int availReqs) {
194    assertNoThrottleException(limiter, availReqs);
195    try {
196      limiter.checkQuota(1, 1, 0, 0, 1, 0, false, 0L);
197      fail("Should have thrown RpcThrottlingException");
198    } catch (RpcThrottlingException e) {
199      // expected
200    }
201  }
202
203  private void assertNoThrottleException(final QuotaLimiter limiter, final int availReqs) {
204    for (int i = 0; i < availReqs; ++i) {
205      try {
206        limiter.checkQuota(1, 1, 0, 0, 1, 0, false, 0L);
207      } catch (RpcThrottlingException e) {
208        fail("Unexpected RpcThrottlingException after " + i + " requests. limit=" + availReqs);
209      }
210      limiter.grabQuota(1, 1, 0, 0, 1, 0, false, 0L);
211    }
212  }
213
214  private void assertNoopLimiter(final QuotaLimiter limiter) {
215    assertTrue(limiter == NoopQuotaLimiter.get());
216    assertNoThrottleException(limiter, 100);
217  }
218}