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