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