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