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