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