1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
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
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
100
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
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
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
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 }