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 org.junit.Assert.fail;
20
21 import java.io.IOException;
22 import java.util.ArrayList;
23 import java.util.List;
24 import java.util.Objects;
25 import java.util.Random;
26 import java.util.Set;
27 import java.util.Map.Entry;
28 import java.util.concurrent.atomic.AtomicLong;
29
30 import org.apache.commons.logging.Log;
31 import org.apache.commons.logging.LogFactory;
32 import org.apache.hadoop.hbase.HBaseTestingUtility;
33 import org.apache.hadoop.hbase.HColumnDescriptor;
34 import org.apache.hadoop.hbase.HTableDescriptor;
35 import org.apache.hadoop.hbase.NamespaceDescriptor;
36 import org.apache.hadoop.hbase.TableName;
37 import org.apache.hadoop.hbase.Waiter.Predicate;
38 import org.apache.hadoop.hbase.classification.InterfaceAudience;
39 import org.apache.hadoop.hbase.client.Admin;
40 import org.apache.hadoop.hbase.client.Connection;
41 import org.apache.hadoop.hbase.client.Put;
42 import org.apache.hadoop.hbase.client.Table;
43 import org.apache.hadoop.hbase.util.Bytes;
44 import org.junit.rules.TestName;
45
46 import com.google.common.collect.HashMultimap;
47 import com.google.common.collect.Iterables;
48 import com.google.common.collect.Multimap;
49
50 @InterfaceAudience.Private
51 public class SpaceQuotaHelperForTests {
52 private static final Log LOG = LogFactory.getLog(SpaceQuotaHelperForTests.class);
53
54 public static final int SIZE_PER_VALUE = 256;
55 public static final String F1 = "f1";
56 public static final long ONE_KILOBYTE = 1024L;
57 public static final long ONE_MEGABYTE = ONE_KILOBYTE * ONE_KILOBYTE;
58
59 private final HBaseTestingUtility testUtil;
60 private final TestName testName;
61 private final AtomicLong counter;
62
63 public SpaceQuotaHelperForTests(
64 HBaseTestingUtility testUtil, TestName testName, AtomicLong counter) {
65 this.testUtil = Objects.requireNonNull(testUtil);
66 this.testName = Objects.requireNonNull(testName);
67 this.counter = Objects.requireNonNull(counter);
68 }
69
70
71
72
73
74
75
76
77 long listNumDefinedQuotas(Connection conn) throws IOException {
78 QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration());
79 try {
80 return Iterables.size(scanner);
81 } finally {
82 if (scanner != null) {
83 scanner.close();
84 }
85 }
86 }
87
88
89
90
91 void removeAllQuotas(Connection conn) throws IOException {
92 QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration());
93 try {
94 for (QuotaSettings quotaSettings : scanner) {
95 final String namespace = quotaSettings.getNamespace();
96 final TableName tableName = quotaSettings.getTableName();
97 if (namespace != null) {
98 LOG.debug("Deleting quota for namespace: " + namespace);
99 QuotaUtil.deleteNamespaceQuota(conn, namespace);
100 } else {
101 assert tableName != null;
102 LOG.debug("Deleting quota for table: "+ tableName);
103 QuotaUtil.deleteTableQuota(conn, tableName);
104 }
105 }
106 } finally {
107 if (scanner != null) {
108 scanner.close();
109 }
110 }
111 }
112
113 QuotaSettings getTableSpaceQuota(Connection conn, TableName tn) throws IOException {
114 try (QuotaRetriever scanner = QuotaRetriever.open(
115 conn.getConfiguration(), new QuotaFilter().setTableFilter(tn.getNameAsString()))) {
116 for (QuotaSettings setting : scanner) {
117 if (setting.getTableName().equals(tn) && setting.getQuotaType() == QuotaType.SPACE) {
118 return setting;
119 }
120 }
121 return null;
122 }
123 }
124
125
126
127
128 void waitForQuotaTable(Connection conn) throws IOException {
129 waitForQuotaTable(conn, 30_000);
130 }
131
132
133
134
135 void waitForQuotaTable(final Connection conn, long timeout) throws IOException {
136 testUtil.waitFor(timeout, 1000, new Predicate<IOException>() {
137 @Override
138 public boolean evaluate() throws IOException {
139 return conn.getAdmin().tableExists(QuotaUtil.QUOTA_TABLE_NAME);
140 }
141 });
142 }
143
144 void writeData(TableName tn, long sizeInBytes) throws IOException {
145 writeData(testUtil.getConnection(), tn, sizeInBytes);
146 }
147
148 void writeData(Connection conn, TableName tn, long sizeInBytes) throws IOException {
149 final Table table = conn.getTable(tn);
150 try {
151 List<Put> updates = new ArrayList<>();
152 long bytesToWrite = sizeInBytes;
153 long rowKeyId = 0L;
154 final StringBuilder sb = new StringBuilder();
155 final Random r = new Random();
156 while (bytesToWrite > 0L) {
157 sb.setLength(0);
158 sb.append(Long.toString(rowKeyId));
159
160 Put p = new Put(Bytes.toBytes(sb.reverse().toString()));
161 byte[] value = new byte[SIZE_PER_VALUE];
162 r.nextBytes(value);
163 p.addColumn(Bytes.toBytes(F1), Bytes.toBytes("q1"), value);
164 updates.add(p);
165
166
167 if (updates.size() > 50) {
168 table.put(updates);
169 updates.clear();
170 }
171
172
173 bytesToWrite -= SIZE_PER_VALUE;
174 rowKeyId++;
175 }
176
177
178 if (!updates.isEmpty()) {
179 table.put(updates);
180 }
181
182 LOG.debug("Data was written to HBase");
183
184 testUtil.getHBaseAdmin().flush(tn);
185 LOG.debug("Data flushed to disk");
186 } finally {
187 table.close();
188 }
189 }
190
191 Multimap<TableName, QuotaSettings> createTablesWithSpaceQuotas() throws Exception {
192 final Admin admin = testUtil.getHBaseAdmin();
193 final Multimap<TableName, QuotaSettings> tablesWithQuotas = HashMultimap.create();
194
195 final TableName tn1 = createTable();
196 final TableName tn2 = createTable();
197
198 NamespaceDescriptor nd = NamespaceDescriptor.create("ns" + counter.getAndIncrement()).build();
199 admin.createNamespace(nd);
200 final TableName tn3 = createTableInNamespace(nd);
201 final TableName tn4 = createTableInNamespace(nd);
202 final TableName tn5 = createTableInNamespace(nd);
203
204 final long sizeLimit1 = 1024L * 1024L * 1024L * 1024L * 5L;
205 final SpaceViolationPolicy violationPolicy1 = SpaceViolationPolicy.NO_WRITES;
206 QuotaSettings qs1 = QuotaSettingsFactory.limitTableSpace(tn1, sizeLimit1, violationPolicy1);
207 tablesWithQuotas.put(tn1, qs1);
208 admin.setQuota(qs1);
209
210 final long sizeLimit2 = 1024L * 1024L * 1024L * 200L;
211 final SpaceViolationPolicy violationPolicy2 = SpaceViolationPolicy.NO_WRITES_COMPACTIONS;
212 QuotaSettings qs2 = QuotaSettingsFactory.limitTableSpace(tn2, sizeLimit2, violationPolicy2);
213 tablesWithQuotas.put(tn2, qs2);
214 admin.setQuota(qs2);
215
216 final long sizeLimit3 = 1024L * 1024L * 1024L * 1024L * 100L;
217 final SpaceViolationPolicy violationPolicy3 = SpaceViolationPolicy.NO_INSERTS;
218 QuotaSettings qs3 = QuotaSettingsFactory.limitNamespaceSpace(
219 nd.getName(), sizeLimit3, violationPolicy3);
220 tablesWithQuotas.put(tn3, qs3);
221 tablesWithQuotas.put(tn4, qs3);
222 tablesWithQuotas.put(tn5, qs3);
223 admin.setQuota(qs3);
224
225 final long sizeLimit4 = 1024L * 1024L * 1024L * 5L;
226 final SpaceViolationPolicy violationPolicy4 = SpaceViolationPolicy.NO_INSERTS;
227 QuotaSettings qs4 = QuotaSettingsFactory.limitTableSpace(tn5, sizeLimit4, violationPolicy4);
228
229
230 tablesWithQuotas.put(tn5, qs4);
231 admin.setQuota(qs4);
232
233 return tablesWithQuotas;
234 }
235
236 TableName createTable() throws Exception {
237 return createTableWithRegions(1);
238 }
239
240 TableName createTableWithRegions(int numRegions) throws Exception {
241 return createTableWithRegions(NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions);
242 }
243
244 TableName createTableWithRegions(Admin admin, int numRegions) throws Exception {
245 return createTableWithRegions(
246 testUtil.getHBaseAdmin(), NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions);
247 }
248
249 TableName createTableWithRegions(String namespace, int numRegions) throws Exception {
250 return createTableWithRegions(testUtil.getHBaseAdmin(), namespace, numRegions);
251 }
252
253 TableName createTableWithRegions(Admin admin, String namespace, int numRegions) throws Exception {
254 final TableName tn = TableName.valueOf(
255 namespace, testName.getMethodName() + counter.getAndIncrement());
256
257
258 if (admin.tableExists(tn)) {
259 admin.disableTable(tn);
260 admin.deleteTable(tn);
261 }
262
263
264 HTableDescriptor tableDesc = new HTableDescriptor(tn);
265 tableDesc.addFamily(new HColumnDescriptor(F1));
266 if (numRegions == 1) {
267 admin.createTable(tableDesc);
268 } else {
269 admin.createTable(tableDesc, Bytes.toBytes("0"), Bytes.toBytes("9"), numRegions);
270 }
271 return tn;
272 }
273
274 TableName createTableInNamespace(NamespaceDescriptor nd) throws Exception {
275 final Admin admin = testUtil.getHBaseAdmin();
276 final TableName tn = TableName.valueOf(nd.getName(),
277 testName.getMethodName() + counter.getAndIncrement());
278
279
280 if (admin.tableExists(tn)) {
281 admin.disableTable(tn);
282 admin.deleteTable(tn);
283 }
284
285
286 HTableDescriptor tableDesc = new HTableDescriptor(tn);
287 tableDesc.addFamily(new HColumnDescriptor(F1));
288
289 admin.createTable(tableDesc);
290 return tn;
291 }
292
293 void partitionTablesByQuotaTarget(Multimap<TableName,QuotaSettings> quotas,
294 Set<TableName> tablesWithTableQuota, Set<TableName> tablesWithNamespaceQuota) {
295
296 for (Entry<TableName, QuotaSettings> entry : quotas.entries()) {
297 SpaceLimitSettings settings = (SpaceLimitSettings) entry.getValue();
298 TableName tn = entry.getKey();
299 if (settings.getTableName() != null) {
300 tablesWithTableQuota.add(tn);
301 }
302 if (settings.getNamespace() != null) {
303 tablesWithNamespaceQuota.add(tn);
304 }
305
306 if (settings.getTableName() == null && settings.getNamespace() == null) {
307 fail("Unexpected table name with null tableName and namespace: " + tn);
308 }
309 }
310 }
311 }