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.assertTrue;
022import static org.junit.Assert.fail;
023
024import java.io.IOException;
025import java.util.ArrayList;
026import java.util.Collection;
027import java.util.List;
028import java.util.Map.Entry;
029import java.util.Objects;
030import java.util.Set;
031import java.util.concurrent.atomic.AtomicLong;
032import org.apache.hadoop.conf.Configuration;
033import org.apache.hadoop.fs.FileSystem;
034import org.apache.hadoop.fs.Path;
035import org.apache.hadoop.hbase.HBaseTestingUtility;
036import org.apache.hadoop.hbase.HConstants;
037import org.apache.hadoop.hbase.MiniHBaseCluster;
038import org.apache.hadoop.hbase.NamespaceDescriptor;
039import org.apache.hadoop.hbase.TableName;
040import org.apache.hadoop.hbase.TableNotEnabledException;
041import org.apache.hadoop.hbase.Waiter.Predicate;
042import org.apache.hadoop.hbase.client.Admin;
043import org.apache.hadoop.hbase.client.Append;
044import org.apache.hadoop.hbase.client.ClientServiceCallable;
045import org.apache.hadoop.hbase.client.ColumnFamilyDescriptorBuilder;
046import org.apache.hadoop.hbase.client.Connection;
047import org.apache.hadoop.hbase.client.Delete;
048import org.apache.hadoop.hbase.client.Increment;
049import org.apache.hadoop.hbase.client.Mutation;
050import org.apache.hadoop.hbase.client.Put;
051import org.apache.hadoop.hbase.client.Result;
052import org.apache.hadoop.hbase.client.ResultScanner;
053import org.apache.hadoop.hbase.client.Scan;
054import org.apache.hadoop.hbase.client.SecureBulkLoadClient;
055import org.apache.hadoop.hbase.client.Table;
056import org.apache.hadoop.hbase.client.TableDescriptor;
057import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
058import org.apache.hadoop.hbase.ipc.RpcControllerFactory;
059import org.apache.hadoop.hbase.regionserver.HRegion;
060import org.apache.hadoop.hbase.regionserver.HStore;
061import org.apache.hadoop.hbase.regionserver.HStoreFile;
062import org.apache.hadoop.hbase.regionserver.TestHRegionServerBulkLoad;
063import org.apache.hadoop.hbase.util.Bytes;
064import org.apache.hadoop.hbase.util.Pair;
065import org.apache.hadoop.util.StringUtils;
066import org.apache.yetus.audience.InterfaceAudience;
067import org.junit.rules.TestName;
068import org.slf4j.Logger;
069import org.slf4j.LoggerFactory;
070
071import org.apache.hbase.thirdparty.com.google.common.collect.HashMultimap;
072import org.apache.hbase.thirdparty.com.google.common.collect.Iterables;
073import org.apache.hbase.thirdparty.com.google.common.collect.Multimap;
074
075@InterfaceAudience.Private
076public class SpaceQuotaHelperForTests {
077  private static final Logger LOG = LoggerFactory.getLogger(SpaceQuotaHelperForTests.class);
078
079  public static final int SIZE_PER_VALUE = 256;
080  public static final String F1 = "f1";
081  public static final long ONE_KILOBYTE = 1024L;
082  public static final long ONE_MEGABYTE = ONE_KILOBYTE * ONE_KILOBYTE;
083  public static final long ONE_GIGABYTE = ONE_MEGABYTE * ONE_KILOBYTE;
084
085  private final HBaseTestingUtility testUtil;
086  private final TestName testName;
087  private final AtomicLong counter;
088  private static final int NUM_RETRIES = 10;
089
090  public SpaceQuotaHelperForTests(HBaseTestingUtility testUtil, TestName testName,
091    AtomicLong counter) {
092    this.testUtil = Objects.requireNonNull(testUtil);
093    this.testName = Objects.requireNonNull(testName);
094    this.counter = Objects.requireNonNull(counter);
095  }
096
097  //
098  // Static helpers
099  //
100
101  static void updateConfigForQuotas(Configuration conf) {
102    // Increase the frequency of some of the chores for responsiveness of the test
103    conf.setInt(FileSystemUtilizationChore.FS_UTILIZATION_CHORE_DELAY_KEY, 1000);
104    conf.setInt(FileSystemUtilizationChore.FS_UTILIZATION_CHORE_PERIOD_KEY, 1000);
105    conf.setInt(QuotaObserverChore.QUOTA_OBSERVER_CHORE_DELAY_KEY, 1000);
106    conf.setInt(QuotaObserverChore.QUOTA_OBSERVER_CHORE_PERIOD_KEY, 1000);
107    conf.setInt(SpaceQuotaRefresherChore.POLICY_REFRESHER_CHORE_DELAY_KEY, 1000);
108    conf.setInt(SpaceQuotaRefresherChore.POLICY_REFRESHER_CHORE_PERIOD_KEY, 1000);
109    conf.setInt(SnapshotQuotaObserverChore.SNAPSHOT_QUOTA_CHORE_DELAY_KEY, 1000);
110    conf.setInt(SnapshotQuotaObserverChore.SNAPSHOT_QUOTA_CHORE_PERIOD_KEY, 1000);
111    conf.setInt(RegionSizeReportingChore.REGION_SIZE_REPORTING_CHORE_PERIOD_KEY, 1000);
112    conf.setInt(RegionSizeReportingChore.REGION_SIZE_REPORTING_CHORE_DELAY_KEY, 1000);
113    // The period at which we check for compacted files that should be deleted from HDFS
114    conf.setInt("hbase.hfile.compaction.discharger.interval", 5 * 1000);
115    conf.setBoolean(QuotaUtil.QUOTA_CONF_KEY, true);
116  }
117
118  //
119  // Helpers
120  //
121
122  /**
123   * Returns the number of quotas defined in the HBase quota table.
124   */
125  long listNumDefinedQuotas(Connection conn) throws IOException {
126    QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration());
127    try {
128      return Iterables.size(scanner);
129    } finally {
130      if (scanner != null) {
131        scanner.close();
132      }
133    }
134  }
135
136  /**
137   * Writes the given mutation into a table until it violates the given policy. Verifies that the
138   * policy has been violated & then returns the name of the table created & written into.
139   */
140  TableName writeUntilViolationAndVerifyViolation(SpaceViolationPolicy policyToViolate, Mutation m)
141    throws Exception {
142    final TableName tn = writeUntilViolation(policyToViolate);
143    verifyViolation(policyToViolate, tn, m);
144    return tn;
145  }
146
147  /**
148   * Writes the given mutation into a table until it violates the given policy. Returns the name of
149   * the table created & written into.
150   */
151  TableName writeUntilViolation(SpaceViolationPolicy policyToViolate) throws Exception {
152    TableName tn = createTableWithRegions(10);
153    setQuotaLimit(tn, policyToViolate, 2L);
154    // Write more data than should be allowed and flush it to disk
155    writeData(tn, 3L * SpaceQuotaHelperForTests.ONE_MEGABYTE);
156
157    // This should be sufficient time for the chores to run and see the change.
158    Thread.sleep(5000);
159
160    return tn;
161  }
162
163  TableName writeUntilViolationAndVerifyViolationInNamespace(String ns,
164    SpaceViolationPolicy policyToViolate, Mutation m) throws Exception {
165    final TableName tn = writeUntilViolationInNamespace(ns, policyToViolate);
166    verifyViolation(policyToViolate, tn, m);
167    return tn;
168  }
169
170  TableName writeUntilViolationInNamespace(String ns, SpaceViolationPolicy policyToViolate)
171    throws Exception {
172    TableName tn = createTableWithRegions(ns, 10);
173
174    setQuotaLimit(ns, policyToViolate, 4L);
175
176    // Write more data than should be allowed and flush it to disk
177    writeData(tn, 5L * SpaceQuotaHelperForTests.ONE_MEGABYTE);
178
179    // This should be sufficient time for the chores to run and see the change.
180    Thread.sleep(5000);
181
182    return tn;
183  }
184
185  /**
186   * Verifies that the given policy on the given table has been violated
187   */
188  void verifyViolation(SpaceViolationPolicy policyToViolate, TableName tn, Mutation m)
189    throws Exception {
190    // But let's try a few times to get the exception before failing
191    boolean sawError = false;
192    String msg = "";
193    for (int i = 0; i < NUM_RETRIES && !sawError; i++) {
194      try (Table table = testUtil.getConnection().getTable(tn)) {
195        if (m instanceof Put) {
196          table.put((Put) m);
197        } else if (m instanceof Delete) {
198          table.delete((Delete) m);
199        } else if (m instanceof Append) {
200          table.append((Append) m);
201        } else if (m instanceof Increment) {
202          table.increment((Increment) m);
203        } else {
204          fail(
205            "Failed to apply " + m.getClass().getSimpleName() + " to the table. Programming error");
206        }
207        LOG.info("Did not reject the " + m.getClass().getSimpleName() + ", will sleep and retry");
208        Thread.sleep(2000);
209      } catch (Exception e) {
210        msg = StringUtils.stringifyException(e);
211        if (
212          (policyToViolate.equals(SpaceViolationPolicy.DISABLE)
213            && e instanceof TableNotEnabledException) || msg.contains(policyToViolate.name())
214        ) {
215          LOG.info("Got the expected exception={}", msg);
216          sawError = true;
217          break;
218        } else {
219          LOG.warn("Did not get the expected exception, will sleep and retry", e);
220          Thread.sleep(2000);
221        }
222      }
223    }
224    if (!sawError) {
225      try (Table quotaTable = testUtil.getConnection().getTable(QuotaUtil.QUOTA_TABLE_NAME)) {
226        ResultScanner scanner = quotaTable.getScanner(new Scan());
227        Result result = null;
228        LOG.info("Dumping contents of hbase:quota table");
229        while ((result = scanner.next()) != null) {
230          LOG.info(Bytes.toString(result.getRow()) + " => " + result.toString());
231        }
232        scanner.close();
233      }
234    } else {
235      if (policyToViolate.equals(SpaceViolationPolicy.DISABLE)) {
236        assertTrue(
237          msg.contains("TableNotEnabledException") || msg.contains(policyToViolate.name()));
238      } else {
239        assertTrue("Expected exception message to contain the word '" + policyToViolate.name()
240          + "', but was " + msg, msg.contains(policyToViolate.name()));
241      }
242    }
243    assertTrue("Expected to see an exception writing data to a table exceeding its quota",
244      sawError);
245  }
246
247  /**
248   * Verifies that no policy has been violated on the given table
249   */
250  void verifyNoViolation(TableName tn, Mutation m) throws Exception {
251    // But let's try a few times to write data before failing
252    boolean sawSuccess = false;
253    for (int i = 0; i < NUM_RETRIES && !sawSuccess; i++) {
254      try (Table table = testUtil.getConnection().getTable(tn)) {
255        if (m instanceof Put) {
256          table.put((Put) m);
257        } else if (m instanceof Delete) {
258          table.delete((Delete) m);
259        } else if (m instanceof Append) {
260          table.append((Append) m);
261        } else if (m instanceof Increment) {
262          table.increment((Increment) m);
263        } else {
264          fail("Failed to apply " + m.getClass().getSimpleName() + " to the table."
265            + " Programming error");
266        }
267        sawSuccess = true;
268      } catch (Exception e) {
269        LOG.info("Rejected the " + m.getClass().getSimpleName() + ", will sleep and retry");
270        Thread.sleep(2000);
271      }
272    }
273    if (!sawSuccess) {
274      try (Table quotaTable = testUtil.getConnection().getTable(QuotaUtil.QUOTA_TABLE_NAME)) {
275        ResultScanner scanner = quotaTable.getScanner(new Scan());
276        Result result = null;
277        LOG.info("Dumping contents of hbase:quota table");
278        while ((result = scanner.next()) != null) {
279          LOG.info(Bytes.toString(result.getRow()) + " => " + result.toString());
280        }
281        scanner.close();
282      }
283    }
284    assertTrue("Expected to succeed in writing data to a table not having quota ", sawSuccess);
285  }
286
287  /**
288   * Verifies that table usage snapshot exists for the table
289   */
290  void verifyTableUsageSnapshotForSpaceQuotaExist(TableName tn) throws Exception {
291    boolean sawUsageSnapshot = false;
292    try (Table quotaTable = testUtil.getConnection().getTable(QuotaTableUtil.QUOTA_TABLE_NAME)) {
293      Scan s = QuotaTableUtil.makeQuotaSnapshotScanForTable(tn);
294      ResultScanner rs = quotaTable.getScanner(s);
295      sawUsageSnapshot = (rs.next() != null);
296    }
297    assertTrue("Expected to succeed in getting table usage snapshots for space quota",
298      sawUsageSnapshot);
299  }
300
301  /**
302   * Sets the given quota (policy & limit) on the passed table.
303   */
304  void setQuotaLimit(final TableName tn, SpaceViolationPolicy policy, long sizeInMBs)
305    throws Exception {
306    final long sizeLimit = sizeInMBs * SpaceQuotaHelperForTests.ONE_MEGABYTE;
307    QuotaSettings settings = QuotaSettingsFactory.limitTableSpace(tn, sizeLimit, policy);
308    testUtil.getAdmin().setQuota(settings);
309    LOG.debug("Quota limit set for table = {}, limit = {}", tn, sizeLimit);
310  }
311
312  /**
313   * Sets the given quota (policy & limit) on the passed namespace.
314   */
315  void setQuotaLimit(String ns, SpaceViolationPolicy policy, long sizeInMBs) throws Exception {
316    final long sizeLimit = sizeInMBs * SpaceQuotaHelperForTests.ONE_MEGABYTE;
317    QuotaSettings settings = QuotaSettingsFactory.limitNamespaceSpace(ns, sizeLimit, policy);
318    testUtil.getAdmin().setQuota(settings);
319    LOG.debug("Quota limit set for namespace = {}, limit = {}", ns, sizeLimit);
320  }
321
322  /**
323   * Removes the space quota from the given table
324   */
325  void removeQuotaFromtable(final TableName tn) throws Exception {
326    QuotaSettings removeQuota = QuotaSettingsFactory.removeTableSpaceLimit(tn);
327    testUtil.getAdmin().setQuota(removeQuota);
328    LOG.debug("Space quota settings removed from the table ", tn);
329  }
330
331  /**
332   * @param tn             the tablename
333   * @param numFiles       number of files
334   * @param numRowsPerFile number of rows per file
335   * @return a clientServiceCallable which can be used with the Caller factory for bulk load
336   * @throws Exception when failed to get connection, table or preparation of the bulk load
337   */
338  ClientServiceCallable<Void> generateFileToLoad(TableName tn, int numFiles, int numRowsPerFile)
339    throws Exception {
340    Connection conn = testUtil.getConnection();
341    FileSystem fs = testUtil.getTestFileSystem();
342    Configuration conf = testUtil.getConfiguration();
343    Path baseDir = new Path(fs.getHomeDirectory(), testName.getMethodName() + "_files");
344    fs.mkdirs(baseDir);
345    final List<Pair<byte[], String>> famPaths = new ArrayList<Pair<byte[], String>>();
346    for (int i = 1; i <= numFiles; i++) {
347      Path hfile = new Path(baseDir, "file" + i);
348      TestHRegionServerBulkLoad.createHFile(fs, hfile, Bytes.toBytes(SpaceQuotaHelperForTests.F1),
349        Bytes.toBytes("to"), Bytes.toBytes("reject"), numRowsPerFile);
350      famPaths.add(new Pair<>(Bytes.toBytes(SpaceQuotaHelperForTests.F1), hfile.toString()));
351    }
352
353    // bulk load HFiles
354    Table table = conn.getTable(tn);
355    final String bulkToken = new SecureBulkLoadClient(conf, table).prepareBulkLoad(conn);
356    return new ClientServiceCallable<Void>(conn, tn, Bytes.toBytes("row"),
357      new RpcControllerFactory(conf).newController(), HConstants.PRIORITY_UNSET) {
358      @Override
359      public Void rpcCall() throws Exception {
360        SecureBulkLoadClient secureClient = null;
361        byte[] regionName = getLocation().getRegionInfo().getRegionName();
362        try (Table table = conn.getTable(getTableName())) {
363          secureClient = new SecureBulkLoadClient(conf, table);
364          secureClient.secureBulkLoadHFiles(getStub(), famPaths, regionName, true, null, bulkToken);
365        }
366        return null;
367      }
368    };
369  }
370
371  /**
372   * Bulk-loads a number of files with a number of rows to the given table.
373   */
374  // ClientServiceCallable<Boolean> generateFileToLoad(
375  // TableName tn, int numFiles, int numRowsPerFile) throws Exception {
376  // Connection conn = testUtil.getConnection();
377  // FileSystem fs = testUtil.getTestFileSystem();
378  // Configuration conf = testUtil.getConfiguration();
379  // Path baseDir = new Path(fs.getHomeDirectory(), testName.getMethodName() + "_files");
380  // fs.mkdirs(baseDir);
381  // final List<Pair<byte[], String>> famPaths = new ArrayList<>();
382  // for (int i = 1; i <= numFiles; i++) {
383  // Path hfile = new Path(baseDir, "file" + i);
384  // TestHRegionServerBulkLoad.createHFile(
385  // fs, hfile, Bytes.toBytes(SpaceQuotaHelperForTests.F1), Bytes.toBytes("my"),
386  // Bytes.toBytes("file"), numRowsPerFile);
387  // famPaths.add(new Pair<>(Bytes.toBytes(SpaceQuotaHelperForTests.F1), hfile.toString()));
388  // }
389  //
390  // // bulk load HFiles
391  // Table table = conn.getTable(tn);
392  // final String bulkToken = new SecureBulkLoadClient(conf, table).prepareBulkLoad(conn);
393  // return new ClientServiceCallable<Boolean>(
394  // conn, tn, Bytes.toBytes("row"), new RpcControllerFactory(conf).newController(),
395  // HConstants.PRIORITY_UNSET) {
396  // @Override
397  // public Boolean rpcCall() throws Exception {
398  // SecureBulkLoadClient secureClient = null;
399  // byte[] regionName = getLocation().getRegion().getRegionName();
400  // try (Table table = conn.getTable(getTableName())) {
401  // secureClient = new SecureBulkLoadClient(conf, table);
402  // return secureClient.secureBulkLoadHFiles(getStub(), famPaths, regionName,
403  // true, null, bulkToken);
404  // }
405  // }
406  // };
407  // }
408
409  /**
410   * Removes the space quota from the given namespace
411   */
412  void removeQuotaFromNamespace(String ns) throws Exception {
413    QuotaSettings removeQuota = QuotaSettingsFactory.removeNamespaceSpaceLimit(ns);
414    Admin admin = testUtil.getAdmin();
415    admin.setQuota(removeQuota);
416    LOG.debug("Space quota settings removed from the namespace ", ns);
417  }
418
419  /**
420   * Removes all quotas defined in the HBase quota table.
421   */
422  void removeAllQuotas() throws Exception {
423    final Connection conn = testUtil.getConnection();
424    removeAllQuotas(conn);
425    assertEquals(0, listNumDefinedQuotas(conn));
426  }
427
428  /**
429   * Removes all quotas defined in the HBase quota table.
430   */
431  void removeAllQuotas(Connection conn) throws IOException {
432    // Wait for the quota table to be created
433    if (!conn.getAdmin().tableExists(QuotaUtil.QUOTA_TABLE_NAME)) {
434      waitForQuotaTable(conn);
435    } else {
436      // Or, clean up any quotas from previous test runs.
437      QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration());
438      try {
439        for (QuotaSettings quotaSettings : scanner) {
440          final String namespace = quotaSettings.getNamespace();
441          final TableName tableName = quotaSettings.getTableName();
442          final String userName = quotaSettings.getUserName();
443          if (namespace != null) {
444            LOG.debug("Deleting quota for namespace: " + namespace);
445            QuotaUtil.deleteNamespaceQuota(conn, namespace);
446          } else if (tableName != null) {
447            LOG.debug("Deleting quota for table: " + tableName);
448            QuotaUtil.deleteTableQuota(conn, tableName);
449          } else if (userName != null) {
450            LOG.debug("Deleting quota for user: " + userName);
451            QuotaUtil.deleteUserQuota(conn, userName);
452          }
453        }
454      } finally {
455        if (scanner != null) {
456          scanner.close();
457        }
458      }
459    }
460  }
461
462  QuotaSettings getTableSpaceQuota(Connection conn, TableName tn) throws IOException {
463    try (QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration(),
464      new QuotaFilter().setTableFilter(tn.getNameAsString()))) {
465      for (QuotaSettings setting : scanner) {
466        if (setting.getTableName().equals(tn) && setting.getQuotaType() == QuotaType.SPACE) {
467          return setting;
468        }
469      }
470      return null;
471    }
472  }
473
474  /**
475   * Waits 30seconds for the HBase quota table to exist.
476   */
477  public void waitForQuotaTable(Connection conn) throws IOException {
478    waitForQuotaTable(conn, 30_000);
479  }
480
481  /**
482   * Waits {@code timeout} milliseconds for the HBase quota table to exist.
483   */
484  public void waitForQuotaTable(Connection conn, long timeout) throws IOException {
485    testUtil.waitFor(timeout, 1000, new Predicate<IOException>() {
486      @Override
487      public boolean evaluate() throws IOException {
488        return conn.getAdmin().tableExists(QuotaUtil.QUOTA_TABLE_NAME);
489      }
490    });
491  }
492
493  void writeData(TableName tn, long sizeInBytes) throws IOException {
494    writeData(testUtil.getConnection(), tn, sizeInBytes);
495  }
496
497  void writeData(Connection conn, TableName tn, long sizeInBytes) throws IOException {
498    writeData(tn, sizeInBytes, Bytes.toBytes("q1"));
499  }
500
501  void writeData(TableName tn, long sizeInBytes, String qual) throws IOException {
502    writeData(tn, sizeInBytes, Bytes.toBytes(qual));
503  }
504
505  void writeData(TableName tn, long sizeInBytes, byte[] qual) throws IOException {
506    final Connection conn = testUtil.getConnection();
507    final Table table = conn.getTable(tn);
508    try {
509      List<Put> updates = new ArrayList<>();
510      long bytesToWrite = sizeInBytes;
511      long rowKeyId = 0L;
512      final StringBuilder sb = new StringBuilder();
513      while (bytesToWrite > 0L) {
514        sb.setLength(0);
515        sb.append(Long.toString(rowKeyId));
516        // Use the reverse counter as the rowKey to get even spread across all regions
517        Put p = new Put(Bytes.toBytes(sb.reverse().toString()));
518        byte[] value = new byte[SIZE_PER_VALUE];
519        Bytes.random(value);
520        p.addColumn(Bytes.toBytes(F1), qual, value);
521        updates.add(p);
522
523        // Batch ~13KB worth of updates
524        if (updates.size() > 50) {
525          table.put(updates);
526          updates.clear();
527        }
528
529        // Just count the value size, ignore the size of rowkey + column
530        bytesToWrite -= SIZE_PER_VALUE;
531        rowKeyId++;
532      }
533
534      // Write the final batch
535      if (!updates.isEmpty()) {
536        table.put(updates);
537      }
538
539      LOG.debug("Data was written to HBase");
540      // Push the data to disk.
541      testUtil.getAdmin().flush(tn);
542      LOG.debug("Data flushed to disk");
543    } finally {
544      table.close();
545    }
546  }
547
548  NamespaceDescriptor createNamespace() throws Exception {
549    return createNamespace(null);
550  }
551
552  NamespaceDescriptor createNamespace(String namespace) throws Exception {
553    if (namespace == null || namespace.trim().isEmpty())
554      namespace = "ns" + counter.getAndIncrement();
555    NamespaceDescriptor nd = NamespaceDescriptor.create(namespace).build();
556    testUtil.getAdmin().createNamespace(nd);
557    return nd;
558  }
559
560  Multimap<TableName, QuotaSettings> createTablesWithSpaceQuotas() throws Exception {
561    final Admin admin = testUtil.getAdmin();
562    final Multimap<TableName, QuotaSettings> tablesWithQuotas = HashMultimap.create();
563
564    final TableName tn1 = createTable();
565    final TableName tn2 = createTable();
566
567    NamespaceDescriptor nd = createNamespace();
568    final TableName tn3 = createTableInNamespace(nd);
569    final TableName tn4 = createTableInNamespace(nd);
570    final TableName tn5 = createTableInNamespace(nd);
571
572    final long sizeLimit1 = 1024L * 1024L * 1024L * 1024L * 5L; // 5TB
573    final SpaceViolationPolicy violationPolicy1 = SpaceViolationPolicy.NO_WRITES;
574    QuotaSettings qs1 = QuotaSettingsFactory.limitTableSpace(tn1, sizeLimit1, violationPolicy1);
575    tablesWithQuotas.put(tn1, qs1);
576    admin.setQuota(qs1);
577
578    final long sizeLimit2 = 1024L * 1024L * 1024L * 200L; // 200GB
579    final SpaceViolationPolicy violationPolicy2 = SpaceViolationPolicy.NO_WRITES_COMPACTIONS;
580    QuotaSettings qs2 = QuotaSettingsFactory.limitTableSpace(tn2, sizeLimit2, violationPolicy2);
581    tablesWithQuotas.put(tn2, qs2);
582    admin.setQuota(qs2);
583
584    final long sizeLimit3 = 1024L * 1024L * 1024L * 1024L * 100L; // 100TB
585    final SpaceViolationPolicy violationPolicy3 = SpaceViolationPolicy.NO_INSERTS;
586    QuotaSettings qs3 =
587      QuotaSettingsFactory.limitNamespaceSpace(nd.getName(), sizeLimit3, violationPolicy3);
588    tablesWithQuotas.put(tn3, qs3);
589    tablesWithQuotas.put(tn4, qs3);
590    tablesWithQuotas.put(tn5, qs3);
591    admin.setQuota(qs3);
592
593    final long sizeLimit4 = 1024L * 1024L * 1024L * 5L; // 5GB
594    final SpaceViolationPolicy violationPolicy4 = SpaceViolationPolicy.NO_INSERTS;
595    QuotaSettings qs4 = QuotaSettingsFactory.limitTableSpace(tn5, sizeLimit4, violationPolicy4);
596    // Override the ns quota for tn5, import edge-case to catch table quota taking
597    // precedence over ns quota.
598    tablesWithQuotas.put(tn5, qs4);
599    admin.setQuota(qs4);
600
601    return tablesWithQuotas;
602  }
603
604  TableName getNextTableName() {
605    return getNextTableName(NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR);
606  }
607
608  TableName getNextTableName(String namespace) {
609    return TableName.valueOf(namespace, testName.getMethodName() + counter.getAndIncrement());
610  }
611
612  TableName createTable() throws Exception {
613    return createTableWithRegions(1);
614  }
615
616  TableName createTableWithRegions(int numRegions) throws Exception {
617    return createTableWithRegions(NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions);
618  }
619
620  TableName createTableWithRegions(Admin admin, int numRegions) throws Exception {
621    return createTableWithRegions(admin, NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions,
622      0);
623  }
624
625  TableName createTableWithRegions(String namespace, int numRegions) throws Exception {
626    return createTableWithRegions(testUtil.getAdmin(), namespace, numRegions, 0);
627  }
628
629  TableName createTableWithRegions(Admin admin, String namespace, int numRegions,
630    int numberOfReplicas) throws Exception {
631    final TableName tn = getNextTableName(namespace);
632
633    // Delete the old table
634    if (admin.tableExists(tn)) {
635      admin.disableTable(tn);
636      admin.deleteTable(tn);
637    }
638
639    // Create the table
640    TableDescriptor tableDesc;
641    if (numberOfReplicas > 0) {
642      tableDesc = TableDescriptorBuilder.newBuilder(tn).setRegionReplication(numberOfReplicas)
643        .setColumnFamily(ColumnFamilyDescriptorBuilder.of(F1)).build();
644    } else {
645      tableDesc = TableDescriptorBuilder.newBuilder(tn)
646        .setColumnFamily(ColumnFamilyDescriptorBuilder.of(F1)).build();
647    }
648    if (numRegions == 1) {
649      admin.createTable(tableDesc);
650    } else {
651      admin.createTable(tableDesc, Bytes.toBytes("0"), Bytes.toBytes("9"), numRegions);
652    }
653    return tn;
654  }
655
656  TableName createTableInNamespace(NamespaceDescriptor nd) throws Exception {
657    final Admin admin = testUtil.getAdmin();
658    final TableName tn =
659      TableName.valueOf(nd.getName(), testName.getMethodName() + counter.getAndIncrement());
660
661    // Delete the old table
662    if (admin.tableExists(tn)) {
663      admin.disableTable(tn);
664      admin.deleteTable(tn);
665    }
666
667    // Create the table
668    TableDescriptor tableDesc = TableDescriptorBuilder.newBuilder(tn)
669      .addColumnFamily(ColumnFamilyDescriptorBuilder.of(F1)).build();
670
671    admin.createTable(tableDesc);
672    return tn;
673  }
674
675  void partitionTablesByQuotaTarget(Multimap<TableName, QuotaSettings> quotas,
676    Set<TableName> tablesWithTableQuota, Set<TableName> tablesWithNamespaceQuota) {
677    // Partition the tables with quotas by table and ns quota
678    for (Entry<TableName, QuotaSettings> entry : quotas.entries()) {
679      SpaceLimitSettings settings = (SpaceLimitSettings) entry.getValue();
680      TableName tn = entry.getKey();
681      if (settings.getTableName() != null) {
682        tablesWithTableQuota.add(tn);
683      }
684      if (settings.getNamespace() != null) {
685        tablesWithNamespaceQuota.add(tn);
686      }
687
688      if (settings.getTableName() == null && settings.getNamespace() == null) {
689        fail("Unexpected table name with null tableName and namespace: " + tn);
690      }
691    }
692  }
693
694  /**
695   * Abstraction to simplify the case where a test needs to verify a certain state on a
696   * {@code SpaceQuotaSnapshot}. This class fails-fast when there is no such snapshot obtained from
697   * the Master. As such, it is not useful to verify the lack of a snapshot.
698   */
699  static abstract class SpaceQuotaSnapshotPredicate implements Predicate<Exception> {
700    private final Connection conn;
701    private final TableName tn;
702    private final String ns;
703
704    SpaceQuotaSnapshotPredicate(Connection conn, TableName tn) {
705      this(Objects.requireNonNull(conn), Objects.requireNonNull(tn), null);
706    }
707
708    SpaceQuotaSnapshotPredicate(Connection conn, String ns) {
709      this(Objects.requireNonNull(conn), null, Objects.requireNonNull(ns));
710    }
711
712    SpaceQuotaSnapshotPredicate(Connection conn, TableName tn, String ns) {
713      if ((null != tn && null != ns) || (null == tn && null == ns)) {
714        throw new IllegalArgumentException(
715          "One of TableName and Namespace must be non-null, and the other null");
716      }
717      this.conn = conn;
718      this.tn = tn;
719      this.ns = ns;
720    }
721
722    @Override
723    public boolean evaluate() throws Exception {
724      SpaceQuotaSnapshot snapshot;
725      if (null == ns) {
726        snapshot = (SpaceQuotaSnapshot) conn.getAdmin().getCurrentSpaceQuotaSnapshot(tn);
727      } else {
728        snapshot = (SpaceQuotaSnapshot) conn.getAdmin().getCurrentSpaceQuotaSnapshot(ns);
729      }
730
731      LOG.debug("Saw quota snapshot for " + (null == tn ? ns : tn) + ": " + snapshot);
732      if (null == snapshot) {
733        return false;
734      }
735      return evaluate(snapshot);
736    }
737
738    /**
739     * Must determine if the given {@code SpaceQuotaSnapshot} meets some criteria.
740     * @param snapshot a non-null snapshot obtained from the HBase Master
741     * @return true if the criteria is met, false otherwise
742     */
743    abstract boolean evaluate(SpaceQuotaSnapshot snapshot) throws Exception;
744  }
745
746  /**
747   * Predicate that waits for all store files in a table to have no compacted files.
748   */
749  static class NoFilesToDischarge implements Predicate<Exception> {
750    private final MiniHBaseCluster cluster;
751    private final TableName tn;
752
753    NoFilesToDischarge(MiniHBaseCluster cluster, TableName tn) {
754      this.cluster = cluster;
755      this.tn = tn;
756    }
757
758    @Override
759    public boolean evaluate() throws Exception {
760      for (HRegion region : cluster.getRegions(tn)) {
761        for (HStore store : region.getStores()) {
762          Collection<HStoreFile> files =
763            store.getStoreEngine().getStoreFileManager().getCompactedfiles();
764          if (null != files && !files.isEmpty()) {
765            LOG.debug(region.getRegionInfo().getEncodedName() + " still has compacted files");
766            return false;
767          }
768        }
769      }
770      return true;
771    }
772  }
773}