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