View Javadoc

1   /*
2    * Licensed to the Apache Software Foundation (ASF) under one or more
3    * contributor license agreements.  See the NOTICE file distributed with
4    * this work for additional information regarding copyright ownership.
5    * The ASF licenses this file to you under the Apache License, Version 2.0
6    * (the "License"); you may not use this file except in compliance with
7    * the License.  You may obtain a copy of the License at
8    *
9    * http://www.apache.org/licenses/LICENSE-2.0
10   *
11   * Unless required by applicable law or agreed to in writing, software
12   * distributed under the License is distributed on an "AS IS" BASIS,
13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14   * See the License for the specific language governing permissions and
15   * limitations under the License.
16   */
17  package org.apache.hadoop.hbase.quotas;
18  
19  import static com.google.common.collect.Iterables.size;
20  import static org.junit.Assert.assertEquals;
21  import static org.junit.Assert.assertNull;
22  import static org.mockito.Matchers.any;
23  import static org.mockito.Mockito.mock;
24  import static org.mockito.Mockito.when;
25  
26  import java.util.HashMap;
27  import java.util.Map;
28  import java.util.concurrent.atomic.AtomicReference;
29  
30  import org.apache.hadoop.hbase.HRegionInfo;
31  import org.apache.hadoop.hbase.TableName;
32  import org.apache.hadoop.hbase.client.Connection;
33  import org.apache.hadoop.hbase.quotas.SpaceQuotaSnapshot.SpaceQuotaStatus;
34  import org.apache.hadoop.hbase.protobuf.ProtobufUtil;
35  import org.apache.hadoop.hbase.protobuf.generated.QuotaProtos;
36  import org.apache.hadoop.hbase.protobuf.generated.QuotaProtos.Quotas;
37  import org.apache.hadoop.hbase.protobuf.generated.QuotaProtos.SpaceQuota;
38  import org.apache.hadoop.hbase.testclassification.SmallTests;
39  import org.apache.hadoop.hbase.util.Bytes;
40  import org.junit.Before;
41  import org.junit.Test;
42  import org.junit.experimental.categories.Category;
43  import org.mockito.invocation.InvocationOnMock;
44  import org.mockito.stubbing.Answer;
45  
46  /**
47   * Test class for {@link TableQuotaSnapshotStore}.
48   */
49  @Category(SmallTests.class)
50  public class TestTableQuotaViolationStore {
51    private static final long ONE_MEGABYTE = 1024L * 1024L;
52  
53    private Connection conn;
54    private QuotaObserverChore chore;
55    private Map<HRegionInfo, Long> regionReports;
56    private TableQuotaSnapshotStore store;
57  
58    @Before
59    public void setup() {
60      conn = mock(Connection.class);
61      chore = mock(QuotaObserverChore.class);
62      regionReports = new HashMap<>();
63      store = new TableQuotaSnapshotStore(conn, chore, regionReports);
64    }
65  
66    @Test
67    public void testFilterRegionsByTable() throws Exception {
68      TableName tn1 = TableName.valueOf("foo");
69      TableName tn2 = TableName.valueOf("bar");
70      TableName tn3 = TableName.valueOf("ns", "foo");
71  
72      assertEquals(0, size(store.filterBySubject(tn1)));
73  
74      for (int i = 0; i < 5; i++) {
75        regionReports.put(new HRegionInfo(tn1, Bytes.toBytes(i), Bytes.toBytes(i+1)), 0L);
76      }
77      for (int i = 0; i < 3; i++) {
78        regionReports.put(new HRegionInfo(tn2, Bytes.toBytes(i), Bytes.toBytes(i+1)), 0L);
79      }
80      for (int i = 0; i < 10; i++) {
81        regionReports.put(new HRegionInfo(tn3, Bytes.toBytes(i), Bytes.toBytes(i+1)), 0L);
82      }
83      assertEquals(18, regionReports.size());
84      assertEquals(5, size(store.filterBySubject(tn1)));
85      assertEquals(3, size(store.filterBySubject(tn2)));
86      assertEquals(10, size(store.filterBySubject(tn3)));
87    }
88  
89    @Test
90    public void testTargetViolationState() {
91      TableName tn1 = TableName.valueOf("violation1");
92      TableName tn2 = TableName.valueOf("observance1");
93      TableName tn3 = TableName.valueOf("observance2");
94      SpaceQuota quota = SpaceQuota.newBuilder()
95          .setSoftLimit(1024L * 1024L)
96          .setViolationPolicy(ProtobufUtil.toProtoViolationPolicy(SpaceViolationPolicy.DISABLE))
97          .build();
98  
99      // Create some junk data to filter. Makes sure it's so large that it would
100     // immediately violate the quota.
101     for (int i = 0; i < 3; i++) {
102       regionReports.put(new HRegionInfo(tn2, Bytes.toBytes(i), Bytes.toBytes(i + 1)),
103           5L * ONE_MEGABYTE);
104       regionReports.put(new HRegionInfo(tn3, Bytes.toBytes(i), Bytes.toBytes(i + 1)),
105           5L * ONE_MEGABYTE);
106     }
107 
108     regionReports.put(new HRegionInfo(tn1, Bytes.toBytes(0), Bytes.toBytes(1)), 1024L * 512L);
109     regionReports.put(new HRegionInfo(tn1, Bytes.toBytes(1), Bytes.toBytes(2)), 1024L * 256L);
110 
111     SpaceQuotaSnapshot tn1Snapshot = new SpaceQuotaSnapshot(
112         SpaceQuotaStatus.notInViolation(), 1024L * 768L, 1024L * 1024L);
113 
114     // Below the quota
115     assertEquals(tn1Snapshot, store.getTargetState(tn1, quota));
116 
117     regionReports.put(new HRegionInfo(tn1, Bytes.toBytes(2), Bytes.toBytes(3)), 1024L * 256L);
118     tn1Snapshot = new SpaceQuotaSnapshot(SpaceQuotaStatus.notInViolation(), 1024L * 1024L, 1024L * 1024L);
119 
120     // Equal to the quota is still in observance
121     assertEquals(tn1Snapshot, store.getTargetState(tn1, quota));
122 
123     regionReports.put(new HRegionInfo(tn1, Bytes.toBytes(3), Bytes.toBytes(4)), 1024L);
124     tn1Snapshot = new SpaceQuotaSnapshot(
125         new SpaceQuotaStatus(SpaceViolationPolicy.DISABLE), 1024L * 1024L + 1024L, 1024L * 1024L);
126 
127     // Exceeds the quota, should be in violation
128     assertEquals(tn1Snapshot, store.getTargetState(tn1, quota));
129   }
130 
131   @Test
132   public void testGetSpaceQuota() throws Exception {
133     TableQuotaSnapshotStore mockStore = mock(TableQuotaSnapshotStore.class);
134     when(mockStore.getSpaceQuota(any(TableName.class))).thenCallRealMethod();
135 
136     Quotas quotaWithSpace = Quotas.newBuilder().setSpace(
137         SpaceQuota.newBuilder()
138             .setSoftLimit(1024L)
139             .setViolationPolicy(QuotaProtos.SpaceViolationPolicy.DISABLE)
140             .build())
141         .build();
142     Quotas quotaWithoutSpace = Quotas.newBuilder().build();
143 
144     final AtomicReference<Quotas> quotaRef = new AtomicReference<>();
145     when(mockStore.getQuotaForTable(any(TableName.class))).then(new Answer<Quotas>() {
146       @Override
147       public Quotas answer(InvocationOnMock invocation) throws Throwable {
148         return quotaRef.get();
149       }
150     });
151 
152     quotaRef.set(quotaWithSpace);
153     assertEquals(quotaWithSpace.getSpace(), mockStore.getSpaceQuota(TableName.valueOf("foo")));
154     quotaRef.set(quotaWithoutSpace);
155     assertNull(mockStore.getSpaceQuota(TableName.valueOf("foo")));
156   }
157 }