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.fail;
020
021import java.io.IOException;
022import java.util.ArrayList;
023import java.util.Collection;
024import java.util.List;
025import java.util.Map.Entry;
026import java.util.Objects;
027import java.util.Random;
028import java.util.Set;
029import java.util.concurrent.atomic.AtomicLong;
030
031import org.apache.hadoop.conf.Configuration;
032import org.apache.hadoop.hbase.HBaseTestingUtility;
033import org.apache.hadoop.hbase.HColumnDescriptor;
034import org.apache.hadoop.hbase.HTableDescriptor;
035import org.apache.hadoop.hbase.MiniHBaseCluster;
036import org.apache.hadoop.hbase.NamespaceDescriptor;
037import org.apache.hadoop.hbase.TableName;
038import org.apache.hadoop.hbase.Waiter.Predicate;
039import org.apache.hadoop.hbase.client.Admin;
040import org.apache.hadoop.hbase.client.Connection;
041import org.apache.hadoop.hbase.client.Put;
042import org.apache.hadoop.hbase.client.Table;
043import org.apache.hadoop.hbase.regionserver.HRegion;
044import org.apache.hadoop.hbase.regionserver.HStore;
045import org.apache.hadoop.hbase.regionserver.HStoreFile;
046import org.apache.hadoop.hbase.util.Bytes;
047import org.apache.yetus.audience.InterfaceAudience;
048import org.junit.rules.TestName;
049import org.slf4j.Logger;
050import org.slf4j.LoggerFactory;
051import org.apache.hbase.thirdparty.com.google.common.collect.HashMultimap;
052import org.apache.hbase.thirdparty.com.google.common.collect.Iterables;
053import org.apache.hbase.thirdparty.com.google.common.collect.Multimap;
054
055@InterfaceAudience.Private
056public class SpaceQuotaHelperForTests {
057  private static final Logger LOG = LoggerFactory.getLogger(SpaceQuotaHelperForTests.class);
058
059  public static final int SIZE_PER_VALUE = 256;
060  public static final String F1 = "f1";
061  public static final long ONE_KILOBYTE = 1024L;
062  public static final long ONE_MEGABYTE = ONE_KILOBYTE * ONE_KILOBYTE;
063  public static final long ONE_GIGABYTE = ONE_MEGABYTE * ONE_KILOBYTE;
064
065  private final HBaseTestingUtility testUtil;
066  private final TestName testName;
067  private final AtomicLong counter;
068
069  public SpaceQuotaHelperForTests(
070      HBaseTestingUtility testUtil, TestName testName, AtomicLong counter) {
071    this.testUtil = Objects.requireNonNull(testUtil);
072    this.testName = Objects.requireNonNull(testName);
073    this.counter = Objects.requireNonNull(counter);
074  }
075
076  //
077  // Static helpers
078  //
079
080  static void updateConfigForQuotas(Configuration conf) {
081    // Increase the frequency of some of the chores for responsiveness of the test
082    conf.setInt(FileSystemUtilizationChore.FS_UTILIZATION_CHORE_DELAY_KEY, 1000);
083    conf.setInt(FileSystemUtilizationChore.FS_UTILIZATION_CHORE_PERIOD_KEY, 1000);
084    conf.setInt(QuotaObserverChore.QUOTA_OBSERVER_CHORE_DELAY_KEY, 1000);
085    conf.setInt(QuotaObserverChore.QUOTA_OBSERVER_CHORE_PERIOD_KEY, 1000);
086    conf.setInt(SpaceQuotaRefresherChore.POLICY_REFRESHER_CHORE_DELAY_KEY, 1000);
087    conf.setInt(SpaceQuotaRefresherChore.POLICY_REFRESHER_CHORE_PERIOD_KEY, 1000);
088    conf.setInt(SnapshotQuotaObserverChore.SNAPSHOT_QUOTA_CHORE_DELAY_KEY, 1000);
089    conf.setInt(SnapshotQuotaObserverChore.SNAPSHOT_QUOTA_CHORE_PERIOD_KEY, 1000);
090    // The period at which we check for compacted files that should be deleted from HDFS
091    conf.setInt("hbase.hfile.compaction.discharger.interval", 5 * 1000);
092    conf.setBoolean(QuotaUtil.QUOTA_CONF_KEY, true);
093  }
094
095  //
096  // Helpers
097  //
098
099  /**
100   * Returns the number of quotas defined in the HBase quota table.
101   */
102  long listNumDefinedQuotas(Connection conn) throws IOException {
103    QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration());
104    try {
105      return Iterables.size(scanner);
106    } finally {
107      if (scanner != null) {
108        scanner.close();
109      }
110    }
111  }
112
113  /**
114   * Removes all quotas defined in the HBase quota table.
115   */
116  void removeAllQuotas(Connection conn) throws IOException, InterruptedException {
117    // Wait for the quota table to be created
118    if (!conn.getAdmin().tableExists(QuotaUtil.QUOTA_TABLE_NAME)) {
119      do {
120        LOG.debug("Quota table does not yet exist");
121        Thread.sleep(1000);
122      } while (!conn.getAdmin().tableExists(QuotaUtil.QUOTA_TABLE_NAME));
123    } else {
124      // Or, clean up any quotas from previous test runs.
125      QuotaRetriever scanner = QuotaRetriever.open(conn.getConfiguration());
126      try {
127        for (QuotaSettings quotaSettings : scanner) {
128          final String namespace = quotaSettings.getNamespace();
129          final TableName tableName = quotaSettings.getTableName();
130          if (namespace != null) {
131            LOG.debug("Deleting quota for namespace: " + namespace);
132            QuotaUtil.deleteNamespaceQuota(conn, namespace);
133          } else {
134            assert tableName != null;
135            LOG.debug("Deleting quota for table: "+ tableName);
136            QuotaUtil.deleteTableQuota(conn, tableName);
137          }
138        }
139      } finally {
140        if (scanner != null) {
141          scanner.close();
142        }
143      }
144    }
145  }
146
147  QuotaSettings getTableSpaceQuota(Connection conn, TableName tn) throws IOException {
148    try (QuotaRetriever scanner = QuotaRetriever.open(
149        conn.getConfiguration(), new QuotaFilter().setTableFilter(tn.getNameAsString()))) {
150      for (QuotaSettings setting : scanner) {
151        if (setting.getTableName().equals(tn) && setting.getQuotaType() == QuotaType.SPACE) {
152          return setting;
153        }
154      }
155      return null;
156    }
157  }
158
159  /**
160   * Waits 30seconds for the HBase quota table to exist.
161   */
162  void waitForQuotaTable(Connection conn) throws IOException {
163    waitForQuotaTable(conn, 30_000);
164  }
165
166  /**
167   * Waits {@code timeout} milliseconds for the HBase quota table to exist.
168   */
169  void waitForQuotaTable(Connection conn, long timeout) throws IOException {
170    testUtil.waitFor(timeout, 1000, new Predicate<IOException>() {
171      @Override
172      public boolean evaluate() throws IOException {
173        return conn.getAdmin().tableExists(QuotaUtil.QUOTA_TABLE_NAME);
174      }
175    });
176  }
177
178  void writeData(TableName tn, long sizeInBytes) throws IOException {
179    writeData(testUtil.getConnection(), tn, sizeInBytes);
180  }
181
182  void writeData(Connection conn, TableName tn, long sizeInBytes) throws IOException {
183    writeData(tn, sizeInBytes, Bytes.toBytes("q1"));
184  }
185
186  void writeData(TableName tn, long sizeInBytes, String qual) throws IOException {
187    writeData(tn, sizeInBytes, Bytes.toBytes(qual));
188  }
189
190  void writeData(TableName tn, long sizeInBytes, byte[] qual) throws IOException {
191    final Connection conn = testUtil.getConnection();
192    final Table table = conn.getTable(tn);
193    try {
194      List<Put> updates = new ArrayList<>();
195      long bytesToWrite = sizeInBytes;
196      long rowKeyId = 0L;
197      final StringBuilder sb = new StringBuilder();
198      final Random r = new Random();
199      while (bytesToWrite > 0L) {
200        sb.setLength(0);
201        sb.append(Long.toString(rowKeyId));
202        // Use the reverse counter as the rowKey to get even spread across all regions
203        Put p = new Put(Bytes.toBytes(sb.reverse().toString()));
204        byte[] value = new byte[SIZE_PER_VALUE];
205        r.nextBytes(value);
206        p.addColumn(Bytes.toBytes(F1), qual, value);
207        updates.add(p);
208
209        // Batch ~13KB worth of updates
210        if (updates.size() > 50) {
211          table.put(updates);
212          updates.clear();
213        }
214
215        // Just count the value size, ignore the size of rowkey + column
216        bytesToWrite -= SIZE_PER_VALUE;
217        rowKeyId++;
218      }
219
220      // Write the final batch
221      if (!updates.isEmpty()) {
222        table.put(updates);
223      }
224
225      LOG.debug("Data was written to HBase");
226      // Push the data to disk.
227      testUtil.getAdmin().flush(tn);
228      LOG.debug("Data flushed to disk");
229    } finally {
230      table.close();
231    }
232  }
233
234  NamespaceDescriptor createNamespace() throws Exception {
235    NamespaceDescriptor nd = NamespaceDescriptor.create("ns" + counter.getAndIncrement()).build();
236    testUtil.getAdmin().createNamespace(nd);
237    return nd;
238  }
239
240  Multimap<TableName, QuotaSettings> createTablesWithSpaceQuotas() throws Exception {
241    final Admin admin = testUtil.getAdmin();
242    final Multimap<TableName, QuotaSettings> tablesWithQuotas = HashMultimap.create();
243
244    final TableName tn1 = createTable();
245    final TableName tn2 = createTable();
246
247    NamespaceDescriptor nd = createNamespace();
248    final TableName tn3 = createTableInNamespace(nd);
249    final TableName tn4 = createTableInNamespace(nd);
250    final TableName tn5 = createTableInNamespace(nd);
251
252    final long sizeLimit1 = 1024L * 1024L * 1024L * 1024L * 5L; // 5TB
253    final SpaceViolationPolicy violationPolicy1 = SpaceViolationPolicy.NO_WRITES;
254    QuotaSettings qs1 = QuotaSettingsFactory.limitTableSpace(tn1, sizeLimit1, violationPolicy1);
255    tablesWithQuotas.put(tn1, qs1);
256    admin.setQuota(qs1);
257
258    final long sizeLimit2 = 1024L * 1024L * 1024L * 200L; // 200GB
259    final SpaceViolationPolicy violationPolicy2 = SpaceViolationPolicy.NO_WRITES_COMPACTIONS;
260    QuotaSettings qs2 = QuotaSettingsFactory.limitTableSpace(tn2, sizeLimit2, violationPolicy2);
261    tablesWithQuotas.put(tn2, qs2);
262    admin.setQuota(qs2);
263
264    final long sizeLimit3 = 1024L * 1024L * 1024L * 1024L * 100L; // 100TB
265    final SpaceViolationPolicy violationPolicy3 = SpaceViolationPolicy.NO_INSERTS;
266    QuotaSettings qs3 = QuotaSettingsFactory.limitNamespaceSpace(
267        nd.getName(), sizeLimit3, violationPolicy3);
268    tablesWithQuotas.put(tn3, qs3);
269    tablesWithQuotas.put(tn4, qs3);
270    tablesWithQuotas.put(tn5, qs3);
271    admin.setQuota(qs3);
272
273    final long sizeLimit4 = 1024L * 1024L * 1024L * 5L; // 5GB
274    final SpaceViolationPolicy violationPolicy4 = SpaceViolationPolicy.NO_INSERTS;
275    QuotaSettings qs4 = QuotaSettingsFactory.limitTableSpace(tn5, sizeLimit4, violationPolicy4);
276    // Override the ns quota for tn5, import edge-case to catch table quota taking
277    // precedence over ns quota.
278    tablesWithQuotas.put(tn5, qs4);
279    admin.setQuota(qs4);
280
281    return tablesWithQuotas;
282  }
283
284  TableName getNextTableName() {
285    return getNextTableName(NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR);
286  }
287
288  TableName getNextTableName(String namespace) {
289    return TableName.valueOf(namespace, testName.getMethodName() + counter.getAndIncrement());
290  }
291
292  TableName createTable() throws Exception {
293    return createTableWithRegions(1);
294  }
295
296  TableName createTableWithRegions(int numRegions) throws Exception {
297    return createTableWithRegions(NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions);
298  }
299
300  TableName createTableWithRegions(Admin admin, int numRegions) throws Exception {
301    return createTableWithRegions(
302        testUtil.getAdmin(), NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions);
303  }
304
305  TableName createTableWithRegions(String namespace, int numRegions) throws Exception {
306    return createTableWithRegions(testUtil.getAdmin(), namespace, numRegions);
307  }
308
309  TableName createTableWithRegions(Admin admin, String namespace, int numRegions) throws Exception {
310    final TableName tn = getNextTableName(namespace);
311
312    // Delete the old table
313    if (admin.tableExists(tn)) {
314      admin.disableTable(tn);
315      admin.deleteTable(tn);
316    }
317
318    // Create the table
319    HTableDescriptor tableDesc = new HTableDescriptor(tn);
320    tableDesc.addFamily(new HColumnDescriptor(F1));
321    if (numRegions == 1) {
322      admin.createTable(tableDesc);
323    } else {
324      admin.createTable(tableDesc, Bytes.toBytes("0"), Bytes.toBytes("9"), numRegions);
325    }
326    return tn;
327  }
328
329  TableName createTableInNamespace(NamespaceDescriptor nd) throws Exception {
330    final Admin admin = testUtil.getAdmin();
331    final TableName tn = TableName.valueOf(nd.getName(),
332        testName.getMethodName() + counter.getAndIncrement());
333
334    // Delete the old table
335    if (admin.tableExists(tn)) {
336      admin.disableTable(tn);
337      admin.deleteTable(tn);
338    }
339
340    // Create the table
341    HTableDescriptor tableDesc = new HTableDescriptor(tn);
342    tableDesc.addFamily(new HColumnDescriptor(F1));
343
344    admin.createTable(tableDesc);
345    return tn;
346  }
347
348  void partitionTablesByQuotaTarget(Multimap<TableName,QuotaSettings> quotas,
349      Set<TableName> tablesWithTableQuota, Set<TableName> tablesWithNamespaceQuota) {
350    // Partition the tables with quotas by table and ns quota
351    for (Entry<TableName, QuotaSettings> entry : quotas.entries()) {
352      SpaceLimitSettings settings = (SpaceLimitSettings) entry.getValue();
353      TableName tn = entry.getKey();
354      if (settings.getTableName() != null) {
355        tablesWithTableQuota.add(tn);
356      }
357      if (settings.getNamespace() != null) {
358        tablesWithNamespaceQuota.add(tn);
359      }
360
361      if (settings.getTableName() == null && settings.getNamespace() == null) {
362        fail("Unexpected table name with null tableName and namespace: " + tn);
363      }
364    }
365  }
366
367  /**
368   * Abstraction to simplify the case where a test needs to verify a certain state
369   * on a {@code SpaceQuotaSnapshot}. This class fails-fast when there is no such
370   * snapshot obtained from the Master. As such, it is not useful to verify the
371   * lack of a snapshot.
372   */
373  static abstract class SpaceQuotaSnapshotPredicate implements Predicate<Exception> {
374    private final Connection conn;
375    private final TableName tn;
376    private final String ns;
377
378    SpaceQuotaSnapshotPredicate(Connection conn, TableName tn) {
379      this(Objects.requireNonNull(conn), Objects.requireNonNull(tn), null);
380    }
381
382    SpaceQuotaSnapshotPredicate(Connection conn, String ns) {
383      this(Objects.requireNonNull(conn), null, Objects.requireNonNull(ns));
384    }
385
386    SpaceQuotaSnapshotPredicate(Connection conn, TableName tn, String ns) {
387      if ((null != tn && null != ns) || (null == tn && null == ns)) {
388        throw new IllegalArgumentException(
389            "One of TableName and Namespace must be non-null, and the other null");
390      }
391      this.conn = conn;
392      this.tn = tn;
393      this.ns = ns;
394    }
395
396    @Override
397    public boolean evaluate() throws Exception {
398      SpaceQuotaSnapshot snapshot;
399      if (null == ns) {
400        snapshot = QuotaTableUtil.getCurrentSnapshot(conn, tn);
401      } else {
402        snapshot = QuotaTableUtil.getCurrentSnapshot(conn, ns);
403      }
404
405      LOG.debug("Saw quota snapshot for " + (null == tn ? ns : tn) + ": " + snapshot);
406      if (null == snapshot) {
407        return false;
408      }
409      return evaluate(snapshot);
410    }
411
412    /**
413     * Must determine if the given {@code SpaceQuotaSnapshot} meets some criteria.
414     *
415     * @param snapshot a non-null snapshot obtained from the HBase Master
416     * @return true if the criteria is met, false otherwise
417     */
418    abstract boolean evaluate(SpaceQuotaSnapshot snapshot) throws Exception;
419  }
420
421  /**
422   * Predicate that waits for all store files in a table to have no compacted files.
423   */
424  static class NoFilesToDischarge implements Predicate<Exception> {
425    private final MiniHBaseCluster cluster;
426    private final TableName tn;
427
428    NoFilesToDischarge(MiniHBaseCluster cluster, TableName tn) {
429      this.cluster = cluster;
430      this.tn = tn;
431    }
432
433    @Override
434    public boolean evaluate() throws Exception {
435      for (HRegion region : cluster.getRegions(tn)) {
436        for (HStore store : region.getStores()) {
437          Collection<HStoreFile> files =
438              store.getStoreEngine().getStoreFileManager().getCompactedfiles();
439          if (null != files && !files.isEmpty()) {
440            LOG.debug(region.getRegionInfo().getEncodedName() + " still has compacted files");
441            return false;
442          }
443        }
444      }
445      return true;
446    }
447  }
448}