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