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 */
018
019package org.apache.hadoop.hbase.quotas;
020
021import java.io.ByteArrayInputStream;
022import java.io.ByteArrayOutputStream;
023import java.io.IOException;
024import java.util.ArrayList;
025import java.util.Collection;
026import java.util.HashMap;
027import java.util.HashSet;
028import java.util.List;
029import java.util.Map;
030import java.util.Objects;
031import java.util.Set;
032import java.util.regex.Pattern;
033
034import org.apache.commons.lang3.StringUtils;
035import org.apache.hadoop.hbase.Cell;
036import org.apache.hadoop.hbase.CellScanner;
037import org.apache.hadoop.hbase.CompareOperator;
038import org.apache.hadoop.hbase.NamespaceDescriptor;
039import org.apache.hadoop.hbase.ServerName;
040import org.apache.hadoop.hbase.TableName;
041import org.apache.yetus.audience.InterfaceAudience;
042import org.apache.yetus.audience.InterfaceStability;
043import org.slf4j.Logger;
044import org.slf4j.LoggerFactory;
045import org.apache.hadoop.hbase.client.ClusterConnection;
046import org.apache.hadoop.hbase.client.Connection;
047import org.apache.hadoop.hbase.client.Delete;
048import org.apache.hadoop.hbase.client.Get;
049import org.apache.hadoop.hbase.client.Put;
050import org.apache.hadoop.hbase.client.QuotaStatusCalls;
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.filter.ColumnPrefixFilter;
056import org.apache.hadoop.hbase.filter.Filter;
057import org.apache.hadoop.hbase.filter.FilterList;
058import org.apache.hadoop.hbase.filter.QualifierFilter;
059import org.apache.hadoop.hbase.filter.RegexStringComparator;
060import org.apache.hadoop.hbase.filter.RowFilter;
061import org.apache.hadoop.hbase.protobuf.ProtobufMagic;
062import org.apache.hadoop.hbase.util.Bytes;
063import org.apache.yetus.audience.InterfaceAudience;
064import org.apache.yetus.audience.InterfaceStability;
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.Multimap;
070import org.apache.hbase.thirdparty.com.google.protobuf.ByteString;
071import org.apache.hbase.thirdparty.com.google.protobuf.InvalidProtocolBufferException;
072import org.apache.hbase.thirdparty.com.google.protobuf.UnsafeByteOperations;
073import org.apache.hadoop.hbase.shaded.protobuf.ProtobufUtil;
074import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos;
075import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos;
076import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.GetQuotaStatesResponse;
077import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.GetSpaceQuotaRegionSizesResponse;
078import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.GetSpaceQuotaSnapshotsResponse;
079import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.GetSpaceQuotaSnapshotsResponse.TableQuotaSnapshot;
080import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.GetSpaceQuotaRegionSizesResponse.RegionSizes;
081import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.Quotas;
082import org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.SpaceQuota;
083import org.apache.hadoop.hbase.util.Bytes;
084
085/**
086 * Helper class to interact with the quota table.
087 * <table>
088 *   <tr><th>ROW-KEY</th><th>FAM/QUAL</th><th>DATA</th></tr>
089 *   <tr><td>n.&lt;namespace&gt;</td><td>q:s</td><td>&lt;global-quotas&gt;</td></tr>
090 *   <tr><td>n.&lt;namespace&gt;</td><td>u:p</td><td>&lt;namespace-quota policy&gt;</td></tr>
091 *   <tr><td>n.&lt;namespace&gt;</td><td>u:s</td><td>&lt;SpaceQuotaSnapshot&gt;</td></tr>
092 *   <tr><td>t.&lt;table&gt;</td><td>q:s</td><td>&lt;global-quotas&gt;</td></tr>
093 *   <tr><td>t.&lt;table&gt;</td><td>u:p</td><td>&lt;table-quota policy&gt;</td></tr>
094 *   <tr><td>t.&lt;table&gt;</td><td>u:ss.&lt;snapshot name&gt;</td><td>&lt;SpaceQuotaSnapshot&gt;</td></tr>
095 *   <tr><td>u.&lt;user&gt;</td><td>q:s</td><td>&lt;global-quotas&gt;</td></tr>
096 *   <tr><td>u.&lt;user&gt;</td><td>q:s.&lt;table&gt;</td><td>&lt;table-quotas&gt;</td></tr>
097 *   <tr><td>u.&lt;user&gt;</td><td>q:s.&lt;ns&gt;</td><td>&lt;namespace-quotas&gt;</td></tr>
098 * </table>
099 */
100@InterfaceAudience.Private
101@InterfaceStability.Evolving
102public class QuotaTableUtil {
103  private static final Logger LOG = LoggerFactory.getLogger(QuotaTableUtil.class);
104
105  /** System table for quotas */
106  public static final TableName QUOTA_TABLE_NAME =
107      TableName.valueOf(NamespaceDescriptor.SYSTEM_NAMESPACE_NAME_STR, "quota");
108
109  protected static final byte[] QUOTA_FAMILY_INFO = Bytes.toBytes("q");
110  protected static final byte[] QUOTA_FAMILY_USAGE = Bytes.toBytes("u");
111  protected static final byte[] QUOTA_QUALIFIER_SETTINGS = Bytes.toBytes("s");
112  protected static final byte[] QUOTA_QUALIFIER_SETTINGS_PREFIX = Bytes.toBytes("s.");
113  protected static final byte[] QUOTA_QUALIFIER_POLICY = Bytes.toBytes("p");
114  protected static final byte[] QUOTA_SNAPSHOT_SIZE_QUALIFIER = Bytes.toBytes("ss");
115  protected static final String QUOTA_POLICY_COLUMN =
116      Bytes.toString(QUOTA_FAMILY_USAGE) + ":" + Bytes.toString(QUOTA_QUALIFIER_POLICY);
117  protected static final byte[] QUOTA_USER_ROW_KEY_PREFIX = Bytes.toBytes("u.");
118  protected static final byte[] QUOTA_TABLE_ROW_KEY_PREFIX = Bytes.toBytes("t.");
119  protected static final byte[] QUOTA_NAMESPACE_ROW_KEY_PREFIX = Bytes.toBytes("n.");
120
121  /* =========================================================================
122   *  Quota "settings" helpers
123   */
124  public static Quotas getTableQuota(final Connection connection, final TableName table)
125      throws IOException {
126    return getQuotas(connection, getTableRowKey(table));
127  }
128
129  public static Quotas getNamespaceQuota(final Connection connection, final String namespace)
130      throws IOException {
131    return getQuotas(connection, getNamespaceRowKey(namespace));
132  }
133
134  public static Quotas getUserQuota(final Connection connection, final String user)
135      throws IOException {
136    return getQuotas(connection, getUserRowKey(user));
137  }
138
139  public static Quotas getUserQuota(final Connection connection, final String user,
140      final TableName table) throws IOException {
141    return getQuotas(connection, getUserRowKey(user), getSettingsQualifierForUserTable(table));
142  }
143
144  public static Quotas getUserQuota(final Connection connection, final String user,
145      final String namespace) throws IOException {
146    return getQuotas(connection, getUserRowKey(user),
147      getSettingsQualifierForUserNamespace(namespace));
148  }
149
150  private static Quotas getQuotas(final Connection connection, final byte[] rowKey)
151      throws IOException {
152    return getQuotas(connection, rowKey, QUOTA_QUALIFIER_SETTINGS);
153  }
154
155  private static Quotas getQuotas(final Connection connection, final byte[] rowKey,
156      final byte[] qualifier) throws IOException {
157    Get get = new Get(rowKey);
158    get.addColumn(QUOTA_FAMILY_INFO, qualifier);
159    Result result = doGet(connection, get);
160    if (result.isEmpty()) {
161      return null;
162    }
163    return quotasFromData(result.getValue(QUOTA_FAMILY_INFO, qualifier));
164  }
165
166  public static Get makeGetForTableQuotas(final TableName table) {
167    Get get = new Get(getTableRowKey(table));
168    get.addFamily(QUOTA_FAMILY_INFO);
169    return get;
170  }
171
172  public static Get makeGetForNamespaceQuotas(final String namespace) {
173    Get get = new Get(getNamespaceRowKey(namespace));
174    get.addFamily(QUOTA_FAMILY_INFO);
175    return get;
176  }
177
178  public static Get makeGetForUserQuotas(final String user, final Iterable<TableName> tables,
179      final Iterable<String> namespaces) {
180    Get get = new Get(getUserRowKey(user));
181    get.addColumn(QUOTA_FAMILY_INFO, QUOTA_QUALIFIER_SETTINGS);
182    for (final TableName table: tables) {
183      get.addColumn(QUOTA_FAMILY_INFO, getSettingsQualifierForUserTable(table));
184    }
185    for (final String ns: namespaces) {
186      get.addColumn(QUOTA_FAMILY_INFO, getSettingsQualifierForUserNamespace(ns));
187    }
188    return get;
189  }
190
191  public static Scan makeScan(final QuotaFilter filter) {
192    Scan scan = new Scan();
193    scan.addFamily(QUOTA_FAMILY_INFO);
194    if (filter != null && !filter.isNull()) {
195      scan.setFilter(makeFilter(filter));
196    }
197    return scan;
198  }
199
200  /**
201   * converts quotafilter to serializeable filterlists.
202   */
203  public static Filter makeFilter(final QuotaFilter filter) {
204    FilterList filterList = new FilterList(FilterList.Operator.MUST_PASS_ALL);
205    if (StringUtils.isNotEmpty(filter.getUserFilter())) {
206      FilterList userFilters = new FilterList(FilterList.Operator.MUST_PASS_ONE);
207      boolean hasFilter = false;
208
209      if (StringUtils.isNotEmpty(filter.getNamespaceFilter())) {
210        FilterList nsFilters = new FilterList(FilterList.Operator.MUST_PASS_ALL);
211        nsFilters.addFilter(new RowFilter(CompareOperator.EQUAL,
212            new RegexStringComparator(getUserRowKeyRegex(filter.getUserFilter()), 0)));
213        nsFilters.addFilter(new QualifierFilter(CompareOperator.EQUAL,
214            new RegexStringComparator(
215              getSettingsQualifierRegexForUserNamespace(filter.getNamespaceFilter()), 0)));
216        userFilters.addFilter(nsFilters);
217        hasFilter = true;
218      }
219      if (StringUtils.isNotEmpty(filter.getTableFilter())) {
220        FilterList tableFilters = new FilterList(FilterList.Operator.MUST_PASS_ALL);
221        tableFilters.addFilter(new RowFilter(CompareOperator.EQUAL,
222            new RegexStringComparator(getUserRowKeyRegex(filter.getUserFilter()), 0)));
223        tableFilters.addFilter(new QualifierFilter(CompareOperator.EQUAL,
224            new RegexStringComparator(
225              getSettingsQualifierRegexForUserTable(filter.getTableFilter()), 0)));
226        userFilters.addFilter(tableFilters);
227        hasFilter = true;
228      }
229      if (!hasFilter) {
230        userFilters.addFilter(new RowFilter(CompareOperator.EQUAL,
231            new RegexStringComparator(getUserRowKeyRegex(filter.getUserFilter()), 0)));
232      }
233
234      filterList.addFilter(userFilters);
235    } else if (StringUtils.isNotEmpty(filter.getTableFilter())) {
236      filterList.addFilter(new RowFilter(CompareOperator.EQUAL,
237          new RegexStringComparator(getTableRowKeyRegex(filter.getTableFilter()), 0)));
238    } else if (StringUtils.isNotEmpty(filter.getNamespaceFilter())) {
239      filterList.addFilter(new RowFilter(CompareOperator.EQUAL,
240          new RegexStringComparator(getNamespaceRowKeyRegex(filter.getNamespaceFilter()), 0)));
241    }
242    return filterList;
243  }
244
245  /**
246   * Creates a {@link Scan} which returns only quota snapshots from the quota table.
247   */
248  public static Scan makeQuotaSnapshotScan() {
249    return makeQuotaSnapshotScanForTable(null);
250  }
251
252  /**
253   * Fetches all {@link SpaceQuotaSnapshot} objects from the {@code hbase:quota} table.
254   *
255   * @param conn The HBase connection
256   * @return A map of table names and their computed snapshot.
257   */
258  public static Map<TableName,SpaceQuotaSnapshot> getSnapshots(Connection conn) throws IOException {
259    Map<TableName,SpaceQuotaSnapshot> snapshots = new HashMap<>();
260    try (Table quotaTable = conn.getTable(QUOTA_TABLE_NAME);
261        ResultScanner rs = quotaTable.getScanner(makeQuotaSnapshotScan())) {
262      for (Result r : rs) {
263        extractQuotaSnapshot(r, snapshots);
264      }
265    }
266    return snapshots;
267  }
268
269  /**
270   * Creates a {@link Scan} which returns only {@link SpaceQuotaSnapshot} from the quota table for a
271   * specific table.
272   * @param tn Optionally, a table name to limit the scan's rowkey space. Can be null.
273   */
274  public static Scan makeQuotaSnapshotScanForTable(TableName tn) {
275    Scan s = new Scan();
276    // Limit to "u:v" column
277    s.addColumn(QUOTA_FAMILY_USAGE, QUOTA_QUALIFIER_POLICY);
278    if (null == tn) {
279      s.setRowPrefixFilter(QUOTA_TABLE_ROW_KEY_PREFIX);
280    } else {
281      byte[] row = getTableRowKey(tn);
282      // Limit rowspace to the "t:" prefix
283      s.withStartRow(row, true).withStopRow(row, true);
284    }
285    return s;
286  }
287
288  /**
289   * Creates a {@link Get} which returns only {@link SpaceQuotaSnapshot} from the quota table for a
290   * specific table.
291   * @param tn table name to get from. Can't be null.
292   */
293  public static Get makeQuotaSnapshotGetForTable(TableName tn) {
294    Get g = new Get(getTableRowKey(tn));
295    // Limit to "u:v" column
296    g.addColumn(QUOTA_FAMILY_USAGE, QUOTA_QUALIFIER_POLICY);
297    return g;
298  }
299
300  /**
301   * Extracts the {@link SpaceViolationPolicy} and {@link TableName} from the provided
302   * {@link Result} and adds them to the given {@link Map}. If the result does not contain
303   * the expected information or the serialized policy in the value is invalid, this method
304   * will throw an {@link IllegalArgumentException}.
305   *
306   * @param result A row from the quota table.
307   * @param snapshots A map of snapshots to add the result of this method into.
308   */
309  public static void extractQuotaSnapshot(
310      Result result, Map<TableName,SpaceQuotaSnapshot> snapshots) {
311    byte[] row = Objects.requireNonNull(result).getRow();
312    if (row == null || row.length == 0) {
313      throw new IllegalArgumentException("Provided result had a null row");
314    }
315    final TableName targetTableName = getTableFromRowKey(row);
316    Cell c = result.getColumnLatestCell(QUOTA_FAMILY_USAGE, QUOTA_QUALIFIER_POLICY);
317    if (c == null) {
318      throw new IllegalArgumentException("Result did not contain the expected column "
319          + QUOTA_POLICY_COLUMN + ", " + result.toString());
320    }
321    ByteString buffer = UnsafeByteOperations.unsafeWrap(
322        c.getValueArray(), c.getValueOffset(), c.getValueLength());
323    try {
324      QuotaProtos.SpaceQuotaSnapshot snapshot = QuotaProtos.SpaceQuotaSnapshot.parseFrom(buffer);
325      snapshots.put(targetTableName, SpaceQuotaSnapshot.toSpaceQuotaSnapshot(snapshot));
326    } catch (InvalidProtocolBufferException e) {
327      throw new IllegalArgumentException(
328          "Result did not contain a valid SpaceQuota protocol buffer message", e);
329    }
330  }
331
332  public static interface UserQuotasVisitor {
333    void visitUserQuotas(final String userName, final Quotas quotas)
334      throws IOException;
335    void visitUserQuotas(final String userName, final TableName table, final Quotas quotas)
336      throws IOException;
337    void visitUserQuotas(final String userName, final String namespace, final Quotas quotas)
338      throws IOException;
339  }
340
341  public static interface TableQuotasVisitor {
342    void visitTableQuotas(final TableName tableName, final Quotas quotas)
343      throws IOException;
344  }
345
346  public static interface NamespaceQuotasVisitor {
347    void visitNamespaceQuotas(final String namespace, final Quotas quotas)
348      throws IOException;
349  }
350
351  public static interface QuotasVisitor extends UserQuotasVisitor,
352      TableQuotasVisitor, NamespaceQuotasVisitor {
353  }
354
355  public static void parseResult(final Result result, final QuotasVisitor visitor)
356      throws IOException {
357    byte[] row = result.getRow();
358    if (isNamespaceRowKey(row)) {
359      parseNamespaceResult(result, visitor);
360    } else if (isTableRowKey(row)) {
361      parseTableResult(result, visitor);
362    } else if (isUserRowKey(row)) {
363      parseUserResult(result, visitor);
364    } else {
365      LOG.warn("unexpected row-key: " + Bytes.toString(row));
366    }
367  }
368
369  public static void parseResultToCollection(final Result result,
370      Collection<QuotaSettings> quotaSettings) throws IOException {
371
372    QuotaTableUtil.parseResult(result, new QuotaTableUtil.QuotasVisitor() {
373      @Override
374      public void visitUserQuotas(String userName, Quotas quotas) {
375        quotaSettings.addAll(QuotaSettingsFactory.fromUserQuotas(userName, quotas));
376      }
377
378      @Override
379      public void visitUserQuotas(String userName, TableName table, Quotas quotas) {
380        quotaSettings.addAll(QuotaSettingsFactory.fromUserQuotas(userName, table, quotas));
381      }
382
383      @Override
384      public void visitUserQuotas(String userName, String namespace, Quotas quotas) {
385        quotaSettings.addAll(QuotaSettingsFactory.fromUserQuotas(userName, namespace, quotas));
386      }
387
388      @Override
389      public void visitTableQuotas(TableName tableName, Quotas quotas) {
390        quotaSettings.addAll(QuotaSettingsFactory.fromTableQuotas(tableName, quotas));
391      }
392
393      @Override
394      public void visitNamespaceQuotas(String namespace, Quotas quotas) {
395        quotaSettings.addAll(QuotaSettingsFactory.fromNamespaceQuotas(namespace, quotas));
396      }
397    });
398  }
399
400  public static void parseNamespaceResult(final Result result,
401      final NamespaceQuotasVisitor visitor) throws IOException {
402    String namespace = getNamespaceFromRowKey(result.getRow());
403    parseNamespaceResult(namespace, result, visitor);
404  }
405
406  protected static void parseNamespaceResult(final String namespace, final Result result,
407      final NamespaceQuotasVisitor visitor) throws IOException {
408    byte[] data = result.getValue(QUOTA_FAMILY_INFO, QUOTA_QUALIFIER_SETTINGS);
409    if (data != null) {
410      Quotas quotas = quotasFromData(data);
411      visitor.visitNamespaceQuotas(namespace, quotas);
412    }
413  }
414
415  public static void parseTableResult(final Result result, final TableQuotasVisitor visitor)
416      throws IOException {
417    TableName table = getTableFromRowKey(result.getRow());
418    parseTableResult(table, result, visitor);
419  }
420
421  protected static void parseTableResult(final TableName table, final Result result,
422      final TableQuotasVisitor visitor) throws IOException {
423    byte[] data = result.getValue(QUOTA_FAMILY_INFO, QUOTA_QUALIFIER_SETTINGS);
424    if (data != null) {
425      Quotas quotas = quotasFromData(data);
426      visitor.visitTableQuotas(table, quotas);
427    }
428  }
429
430  public static void parseUserResult(final Result result, final UserQuotasVisitor visitor)
431      throws IOException {
432    String userName = getUserFromRowKey(result.getRow());
433    parseUserResult(userName, result, visitor);
434  }
435
436  protected static void parseUserResult(final String userName, final Result result,
437      final UserQuotasVisitor visitor) throws IOException {
438    Map<byte[], byte[]> familyMap = result.getFamilyMap(QUOTA_FAMILY_INFO);
439    if (familyMap == null || familyMap.isEmpty()) return;
440
441    for (Map.Entry<byte[], byte[]> entry: familyMap.entrySet()) {
442      Quotas quotas = quotasFromData(entry.getValue());
443      if (Bytes.startsWith(entry.getKey(), QUOTA_QUALIFIER_SETTINGS_PREFIX)) {
444        String name = Bytes.toString(entry.getKey(), QUOTA_QUALIFIER_SETTINGS_PREFIX.length);
445        if (name.charAt(name.length() - 1) == TableName.NAMESPACE_DELIM) {
446          String namespace = name.substring(0, name.length() - 1);
447          visitor.visitUserQuotas(userName, namespace, quotas);
448        } else {
449          TableName table = TableName.valueOf(name);
450          visitor.visitUserQuotas(userName, table, quotas);
451        }
452      } else if (Bytes.equals(entry.getKey(), QUOTA_QUALIFIER_SETTINGS)) {
453        visitor.visitUserQuotas(userName, quotas);
454      }
455    }
456  }
457
458  /**
459   * Creates a {@link Put} to store the given {@code snapshot} for the given {@code tableName} in
460   * the quota table.
461   */
462  static Put createPutForSpaceSnapshot(TableName tableName, SpaceQuotaSnapshot snapshot) {
463    Put p = new Put(getTableRowKey(tableName));
464    p.addColumn(
465        QUOTA_FAMILY_USAGE, QUOTA_QUALIFIER_POLICY,
466        SpaceQuotaSnapshot.toProtoSnapshot(snapshot).toByteArray());
467    return p;
468  }
469
470  /**
471   * Creates a {@link Get} for the HBase snapshot's size against the given table.
472   */
473  static Get makeGetForSnapshotSize(TableName tn, String snapshot) {
474    Get g = new Get(Bytes.add(QUOTA_TABLE_ROW_KEY_PREFIX, Bytes.toBytes(tn.toString())));
475    g.addColumn(
476        QUOTA_FAMILY_USAGE,
477        Bytes.add(QUOTA_SNAPSHOT_SIZE_QUALIFIER, Bytes.toBytes(snapshot)));
478    return g;
479  }
480
481  /**
482   * Creates a {@link Put} to persist the current size of the {@code snapshot} with respect to
483   * the given {@code table}.
484   */
485  static Put createPutForSnapshotSize(TableName tableName, String snapshot, long size) {
486    // We just need a pb message with some `long usage`, so we can just reuse the
487    // SpaceQuotaSnapshot message instead of creating a new one.
488    Put p = new Put(getTableRowKey(tableName));
489    p.addColumn(QUOTA_FAMILY_USAGE, getSnapshotSizeQualifier(snapshot),
490        org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.SpaceQuotaSnapshot
491            .newBuilder().setQuotaUsage(size).build().toByteArray());
492    return p;
493  }
494
495  /**
496   * Creates a {@code Put} for the namespace's total snapshot size.
497   */
498  static Put createPutForNamespaceSnapshotSize(String namespace, long size) {
499    Put p = new Put(getNamespaceRowKey(namespace));
500    p.addColumn(QUOTA_FAMILY_USAGE, QUOTA_SNAPSHOT_SIZE_QUALIFIER,
501        org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.SpaceQuotaSnapshot
502            .newBuilder().setQuotaUsage(size).build().toByteArray());
503    return p;
504  }
505
506  /**
507   * Returns a list of {@code Delete} to remove given table snapshot
508   * entries to remove from quota table
509   * @param snapshotEntriesToRemove the entries to remove
510   */
511  static List<Delete> createDeletesForExistingTableSnapshotSizes(
512      Multimap<TableName, String> snapshotEntriesToRemove) {
513    List<Delete> deletes = new ArrayList<>();
514    for (Map.Entry<TableName, Collection<String>> entry : snapshotEntriesToRemove.asMap()
515        .entrySet()) {
516      for (String snapshot : entry.getValue()) {
517        Delete d = new Delete(getTableRowKey(entry.getKey()));
518        d.addColumns(QUOTA_FAMILY_USAGE,
519            Bytes.add(QUOTA_SNAPSHOT_SIZE_QUALIFIER, Bytes.toBytes(snapshot)));
520        deletes.add(d);
521      }
522    }
523    return deletes;
524  }
525
526  /**
527   * Returns a list of {@code Delete} to remove all table snapshot entries from quota table.
528   * @param connection connection to re-use
529   */
530  static List<Delete> createDeletesForExistingTableSnapshotSizes(Connection connection)
531      throws IOException {
532    return createDeletesForExistingSnapshotsFromScan(connection, createScanForSpaceSnapshotSizes());
533  }
534
535  /**
536   * Returns a list of {@code Delete} to remove given namespace snapshot
537   * entries to removefrom quota table
538   * @param snapshotEntriesToRemove the entries to remove
539   */
540  static List<Delete> createDeletesForExistingNamespaceSnapshotSizes(
541      Set<String> snapshotEntriesToRemove) {
542    List<Delete> deletes = new ArrayList<>();
543    for (String snapshot : snapshotEntriesToRemove) {
544      Delete d = new Delete(getNamespaceRowKey(snapshot));
545      d.addColumns(QUOTA_FAMILY_USAGE, QUOTA_SNAPSHOT_SIZE_QUALIFIER);
546      deletes.add(d);
547    }
548    return deletes;
549  }
550
551  /**
552   * Returns a list of {@code Delete} to remove all namespace snapshot entries from quota table.
553   * @param connection connection to re-use
554   */
555  static List<Delete> createDeletesForExistingNamespaceSnapshotSizes(Connection connection)
556      throws IOException {
557    return createDeletesForExistingSnapshotsFromScan(connection,
558        createScanForNamespaceSnapshotSizes());
559  }
560
561  /**
562   * Returns a list of {@code Delete} to remove all entries returned by the passed scanner.
563   * @param connection connection to re-use
564   * @param scan the scanner to use to generate the list of deletes
565   */
566  static List<Delete> createDeletesForExistingSnapshotsFromScan(Connection connection, Scan scan)
567      throws IOException {
568    List<Delete> deletes = new ArrayList<>();
569    try (Table quotaTable = connection.getTable(QUOTA_TABLE_NAME);
570        ResultScanner rs = quotaTable.getScanner(scan)) {
571      for (Result r : rs) {
572        CellScanner cs = r.cellScanner();
573        while (cs.advance()) {
574          Cell c = cs.current();
575          byte[] family = Bytes.copy(c.getFamilyArray(), c.getFamilyOffset(), c.getFamilyLength());
576          byte[] qual =
577              Bytes.copy(c.getQualifierArray(), c.getQualifierOffset(), c.getQualifierLength());
578          Delete d = new Delete(r.getRow());
579          d.addColumns(family, qual);
580          deletes.add(d);
581        }
582      }
583      return deletes;
584    }
585  }
586
587  /**
588   * Fetches the computed size of all snapshots against tables in a namespace for space quotas.
589   */
590  static long getNamespaceSnapshotSize(
591      Connection conn, String namespace) throws IOException {
592    try (Table quotaTable = conn.getTable(QuotaTableUtil.QUOTA_TABLE_NAME)) {
593      Result r = quotaTable.get(createGetNamespaceSnapshotSize(namespace));
594      if (r.isEmpty()) {
595        return 0L;
596      }
597      r.advance();
598      return parseSnapshotSize(r.current());
599    } catch (InvalidProtocolBufferException e) {
600      throw new IOException("Could not parse snapshot size value for namespace " + namespace, e);
601    }
602  }
603
604  /**
605   * Creates a {@code Get} to fetch the namespace's total snapshot size.
606   */
607  static Get createGetNamespaceSnapshotSize(String namespace) {
608    Get g = new Get(getNamespaceRowKey(namespace));
609    g.addColumn(QUOTA_FAMILY_USAGE, QUOTA_SNAPSHOT_SIZE_QUALIFIER);
610    return g;
611  }
612
613  /**
614   * Parses the snapshot size from the given Cell's value.
615   */
616  static long parseSnapshotSize(Cell c) throws InvalidProtocolBufferException {
617    ByteString bs = UnsafeByteOperations.unsafeWrap(
618        c.getValueArray(), c.getValueOffset(), c.getValueLength());
619    return QuotaProtos.SpaceQuotaSnapshot.parseFrom(bs).getQuotaUsage();
620  }
621
622  /**
623   * Returns a scanner for all existing namespace snapshot entries.
624   */
625  static Scan createScanForNamespaceSnapshotSizes() {
626    return createScanForNamespaceSnapshotSizes(null);
627  }
628
629  /**
630   * Returns a scanner for all namespace snapshot entries of the given namespace
631   * @param namespace name of the namespace whose snapshot entries are to be scanned
632   */
633  static Scan createScanForNamespaceSnapshotSizes(String namespace) {
634    Scan s = new Scan();
635    if (namespace == null || namespace.isEmpty()) {
636      // Read all namespaces, just look at the row prefix
637      s.setRowPrefixFilter(QUOTA_NAMESPACE_ROW_KEY_PREFIX);
638    } else {
639      // Fetch the exact row for the table
640      byte[] rowkey = getNamespaceRowKey(namespace);
641      // Fetch just this one row
642      s.withStartRow(rowkey).withStopRow(rowkey, true);
643    }
644
645    // Just the usage family and only the snapshot size qualifiers
646    return s.addFamily(QUOTA_FAMILY_USAGE)
647        .setFilter(new ColumnPrefixFilter(QUOTA_SNAPSHOT_SIZE_QUALIFIER));
648  }
649
650  static Scan createScanForSpaceSnapshotSizes() {
651    return createScanForSpaceSnapshotSizes(null);
652  }
653
654  static Scan createScanForSpaceSnapshotSizes(TableName table) {
655    Scan s = new Scan();
656    if (null == table) {
657      // Read all tables, just look at the row prefix
658      s.setRowPrefixFilter(QUOTA_TABLE_ROW_KEY_PREFIX);
659    } else {
660      // Fetch the exact row for the table
661      byte[] rowkey = getTableRowKey(table);
662      // Fetch just this one row
663      s.withStartRow(rowkey).withStopRow(rowkey, true);
664    }
665
666    // Just the usage family and only the snapshot size qualifiers
667    return s.addFamily(QUOTA_FAMILY_USAGE).setFilter(
668        new ColumnPrefixFilter(QUOTA_SNAPSHOT_SIZE_QUALIFIER));
669  }
670
671  /**
672   * Fetches any persisted HBase snapshot sizes stored in the quota table. The sizes here are
673   * computed relative to the table which the snapshot was created from. A snapshot's size will
674   * not include the size of files which the table still refers. These sizes, in bytes, are what
675   * is used internally to compute quota violation for tables and namespaces.
676   *
677   * @return A map of snapshot name to size in bytes per space quota computations
678   */
679  public static Map<String,Long> getObservedSnapshotSizes(Connection conn) throws IOException {
680    try (Table quotaTable = conn.getTable(QUOTA_TABLE_NAME);
681        ResultScanner rs = quotaTable.getScanner(createScanForSpaceSnapshotSizes())) {
682      final Map<String,Long> snapshotSizes = new HashMap<>();
683      for (Result r : rs) {
684        CellScanner cs = r.cellScanner();
685        while (cs.advance()) {
686          Cell c = cs.current();
687          final String snapshot = extractSnapshotNameFromSizeCell(c);
688          final long size = parseSnapshotSize(c);
689          snapshotSizes.put(snapshot, size);
690        }
691      }
692      return snapshotSizes;
693    }
694  }
695
696  /**
697   * Returns a multimap for all existing table snapshot entries.
698   * @param conn connection to re-use
699   */
700  public static Multimap<TableName, String> getTableSnapshots(Connection conn) throws IOException {
701    try (Table quotaTable = conn.getTable(QUOTA_TABLE_NAME);
702        ResultScanner rs = quotaTable.getScanner(createScanForSpaceSnapshotSizes())) {
703      Multimap<TableName, String> snapshots = HashMultimap.create();
704      for (Result r : rs) {
705        CellScanner cs = r.cellScanner();
706        while (cs.advance()) {
707          Cell c = cs.current();
708
709          final String snapshot = extractSnapshotNameFromSizeCell(c);
710          snapshots.put(getTableFromRowKey(r.getRow()), snapshot);
711        }
712      }
713      return snapshots;
714    }
715  }
716
717  /**
718   * Returns a set of the names of all namespaces containing snapshot entries.
719   * @param conn connection to re-use
720   */
721  public static Set<String> getNamespaceSnapshots(Connection conn) throws IOException {
722    try (Table quotaTable = conn.getTable(QUOTA_TABLE_NAME);
723        ResultScanner rs = quotaTable.getScanner(createScanForNamespaceSnapshotSizes())) {
724      Set<String> snapshots = new HashSet<>();
725      for (Result r : rs) {
726        CellScanner cs = r.cellScanner();
727        while (cs.advance()) {
728          cs.current();
729          snapshots.add(getNamespaceFromRowKey(r.getRow()));
730        }
731      }
732      return snapshots;
733    }
734  }
735
736  /* =========================================================================
737   *  Space quota status RPC helpers
738   */
739  /**
740   * Fetches the table sizes on the filesystem as tracked by the HBase Master.
741   */
742  public static Map<TableName,Long> getMasterReportedTableSizes(
743      Connection conn) throws IOException {
744    if (!(conn instanceof ClusterConnection)) {
745      throw new IllegalArgumentException("Expected a ClusterConnection");
746    }
747    ClusterConnection clusterConn = (ClusterConnection) conn;
748    GetSpaceQuotaRegionSizesResponse response = QuotaStatusCalls.getMasterRegionSizes(
749        clusterConn, 0);
750    Map<TableName,Long> tableSizes = new HashMap<>();
751    for (RegionSizes sizes : response.getSizesList()) {
752      TableName tn = ProtobufUtil.toTableName(sizes.getTableName());
753      tableSizes.put(tn, sizes.getSize());
754    }
755    return tableSizes;
756  }
757
758  /**
759   * Fetches the observed {@link SpaceQuotaSnapshot}s observed by a RegionServer.
760   */
761  public static Map<TableName,SpaceQuotaSnapshot> getRegionServerQuotaSnapshots(
762      Connection conn, ServerName regionServer) throws IOException {
763    if (!(conn instanceof ClusterConnection)) {
764      throw new IllegalArgumentException("Expected a ClusterConnection");
765    }
766    ClusterConnection clusterConn = (ClusterConnection) conn;
767    GetSpaceQuotaSnapshotsResponse response = QuotaStatusCalls.getRegionServerQuotaSnapshot(
768        clusterConn, 0, regionServer);
769    Map<TableName,SpaceQuotaSnapshot> snapshots = new HashMap<>();
770    for (TableQuotaSnapshot snapshot : response.getSnapshotsList()) {
771      snapshots.put(
772          ProtobufUtil.toTableName(snapshot.getTableName()),
773          SpaceQuotaSnapshot.toSpaceQuotaSnapshot(snapshot.getSnapshot()));
774    }
775    return snapshots;
776  }
777
778  /**
779   * Returns the Master's view of a quota on the given {@code tableName} or null if the
780   * Master has no quota information on that table.
781   */
782  public static SpaceQuotaSnapshot getCurrentSnapshot(
783      Connection conn, TableName tn) throws IOException {
784    if (!(conn instanceof ClusterConnection)) {
785      throw new IllegalArgumentException("Expected a ClusterConnection");
786    }
787    ClusterConnection clusterConn = (ClusterConnection) conn;
788    GetQuotaStatesResponse resp = QuotaStatusCalls.getMasterQuotaStates(clusterConn, 0);
789    HBaseProtos.TableName protoTableName = ProtobufUtil.toProtoTableName(tn);
790    for (GetQuotaStatesResponse.TableQuotaSnapshot tableSnapshot : resp.getTableSnapshotsList()) {
791      if (protoTableName.equals(tableSnapshot.getTableName())) {
792        return SpaceQuotaSnapshot.toSpaceQuotaSnapshot(tableSnapshot.getSnapshot());
793      }
794    }
795    return null;
796  }
797
798  /**
799   * Returns the Master's view of a quota on the given {@code namespace} or null if the
800   * Master has no quota information on that namespace.
801   */
802  public static SpaceQuotaSnapshot getCurrentSnapshot(
803      Connection conn, String namespace) throws IOException {
804    if (!(conn instanceof ClusterConnection)) {
805      throw new IllegalArgumentException("Expected a ClusterConnection");
806    }
807    ClusterConnection clusterConn = (ClusterConnection) conn;
808    GetQuotaStatesResponse resp = QuotaStatusCalls.getMasterQuotaStates(clusterConn, 0);
809    for (GetQuotaStatesResponse.NamespaceQuotaSnapshot nsSnapshot : resp.getNsSnapshotsList()) {
810      if (namespace.equals(nsSnapshot.getNamespace())) {
811        return SpaceQuotaSnapshot.toSpaceQuotaSnapshot(nsSnapshot.getSnapshot());
812      }
813    }
814    return null;
815  }
816
817  /**
818   * Returns the current space quota snapshot of the given {@code tableName} from
819   * {@code QuotaTableUtil.QUOTA_TABLE_NAME} or null if the no quota information is available for
820   * that tableName.
821   * @param conn connection to re-use
822   * @param tableName name of the table whose current snapshot is to be retreived
823   */
824  public static SpaceQuotaSnapshot getCurrentSnapshotFromQuotaTable(Connection conn,
825      TableName tableName) throws IOException {
826    try (Table quotaTable = conn.getTable(QuotaTableUtil.QUOTA_TABLE_NAME)) {
827      Map<TableName, SpaceQuotaSnapshot> snapshots = new HashMap<>(1);
828      Result result = quotaTable.get(makeQuotaSnapshotGetForTable(tableName));
829      // if we don't have any row corresponding to this get, return null
830      if (result.isEmpty()) {
831        return null;
832      }
833      // otherwise, extract quota snapshot in snapshots object
834      extractQuotaSnapshot(result, snapshots);
835      return snapshots.get(tableName);
836    }
837  }
838
839  /* =========================================================================
840   *  Quotas protobuf helpers
841   */
842  protected static Quotas quotasFromData(final byte[] data) throws IOException {
843    return quotasFromData(data, 0, data.length);
844  }
845
846  protected static Quotas quotasFromData(
847      final byte[] data, int offset, int length) throws IOException {
848    int magicLen = ProtobufMagic.lengthOfPBMagic();
849    if (!ProtobufMagic.isPBMagicPrefix(data, offset, magicLen)) {
850      throw new IOException("Missing pb magic prefix");
851    }
852    return Quotas.parseFrom(new ByteArrayInputStream(data, offset + magicLen, length - magicLen));
853  }
854
855  protected static byte[] quotasToData(final Quotas data) throws IOException {
856    ByteArrayOutputStream stream = new ByteArrayOutputStream();
857    stream.write(ProtobufMagic.PB_MAGIC);
858    data.writeTo(stream);
859    return stream.toByteArray();
860  }
861
862  public static boolean isEmptyQuota(final Quotas quotas) {
863    boolean hasSettings = false;
864    hasSettings |= quotas.hasThrottle();
865    hasSettings |= quotas.hasBypassGlobals();
866    // Only when there is a space quota, make sure there's actually both fields provided
867    // Otherwise, it's a noop.
868    if (quotas.hasSpace()) {
869      hasSettings |= (quotas.getSpace().hasSoftLimit() && quotas.getSpace().hasViolationPolicy());
870    }
871    return !hasSettings;
872  }
873
874  /* =========================================================================
875   *  HTable helpers
876   */
877  protected static Result doGet(final Connection connection, final Get get)
878      throws IOException {
879    try (Table table = connection.getTable(QUOTA_TABLE_NAME)) {
880      return table.get(get);
881    }
882  }
883
884  protected static Result[] doGet(final Connection connection, final List<Get> gets)
885      throws IOException {
886    try (Table table = connection.getTable(QUOTA_TABLE_NAME)) {
887      return table.get(gets);
888    }
889  }
890
891  /* =========================================================================
892   *  Quota table row key helpers
893   */
894  protected static byte[] getUserRowKey(final String user) {
895    return Bytes.add(QUOTA_USER_ROW_KEY_PREFIX, Bytes.toBytes(user));
896  }
897
898  protected static byte[] getTableRowKey(final TableName table) {
899    return Bytes.add(QUOTA_TABLE_ROW_KEY_PREFIX, table.getName());
900  }
901
902  protected static byte[] getNamespaceRowKey(final String namespace) {
903    return Bytes.add(QUOTA_NAMESPACE_ROW_KEY_PREFIX, Bytes.toBytes(namespace));
904  }
905
906  protected static byte[] getSettingsQualifierForUserTable(final TableName tableName) {
907    return Bytes.add(QUOTA_QUALIFIER_SETTINGS_PREFIX, tableName.getName());
908  }
909
910  protected static byte[] getSettingsQualifierForUserNamespace(final String namespace) {
911    return Bytes.add(QUOTA_QUALIFIER_SETTINGS_PREFIX,
912        Bytes.toBytes(namespace + TableName.NAMESPACE_DELIM));
913  }
914
915  protected static String getUserRowKeyRegex(final String user) {
916    return getRowKeyRegEx(QUOTA_USER_ROW_KEY_PREFIX, user);
917  }
918
919  protected static String getTableRowKeyRegex(final String table) {
920    return getRowKeyRegEx(QUOTA_TABLE_ROW_KEY_PREFIX, table);
921  }
922
923  protected static String getNamespaceRowKeyRegex(final String namespace) {
924    return getRowKeyRegEx(QUOTA_NAMESPACE_ROW_KEY_PREFIX, namespace);
925  }
926
927  private static String getRowKeyRegEx(final byte[] prefix, final String regex) {
928    return '^' + Pattern.quote(Bytes.toString(prefix)) + regex + '$';
929  }
930
931  protected static String getSettingsQualifierRegexForUserTable(final String table) {
932    return '^' + Pattern.quote(Bytes.toString(QUOTA_QUALIFIER_SETTINGS_PREFIX)) +
933          table + "(?<!" + Pattern.quote(Character.toString(TableName.NAMESPACE_DELIM)) + ")$";
934  }
935
936  protected static String getSettingsQualifierRegexForUserNamespace(final String namespace) {
937    return '^' + Pattern.quote(Bytes.toString(QUOTA_QUALIFIER_SETTINGS_PREFIX)) +
938                  namespace + Pattern.quote(Character.toString(TableName.NAMESPACE_DELIM)) + '$';
939  }
940
941  protected static boolean isNamespaceRowKey(final byte[] key) {
942    return Bytes.startsWith(key, QUOTA_NAMESPACE_ROW_KEY_PREFIX);
943  }
944
945  protected static String getNamespaceFromRowKey(final byte[] key) {
946    return Bytes.toString(key, QUOTA_NAMESPACE_ROW_KEY_PREFIX.length);
947  }
948
949  protected static boolean isTableRowKey(final byte[] key) {
950    return Bytes.startsWith(key, QUOTA_TABLE_ROW_KEY_PREFIX);
951  }
952
953  protected static TableName getTableFromRowKey(final byte[] key) {
954    return TableName.valueOf(Bytes.toString(key, QUOTA_TABLE_ROW_KEY_PREFIX.length));
955  }
956
957  protected static boolean isUserRowKey(final byte[] key) {
958    return Bytes.startsWith(key, QUOTA_USER_ROW_KEY_PREFIX);
959  }
960
961  protected static String getUserFromRowKey(final byte[] key) {
962    return Bytes.toString(key, QUOTA_USER_ROW_KEY_PREFIX.length);
963  }
964
965  protected static SpaceQuota getProtoViolationPolicy(SpaceViolationPolicy policy) {
966    return SpaceQuota.newBuilder()
967          .setViolationPolicy(ProtobufUtil.toProtoViolationPolicy(policy))
968          .build();
969  }
970
971  protected static SpaceViolationPolicy getViolationPolicy(SpaceQuota proto) {
972    if (!proto.hasViolationPolicy()) {
973      throw new IllegalArgumentException("Protobuf SpaceQuota does not have violation policy.");
974    }
975    return ProtobufUtil.toViolationPolicy(proto.getViolationPolicy());
976  }
977
978  protected static byte[] getSnapshotSizeQualifier(String snapshotName) {
979    return Bytes.add(QUOTA_SNAPSHOT_SIZE_QUALIFIER, Bytes.toBytes(snapshotName));
980  }
981
982  protected static String extractSnapshotNameFromSizeCell(Cell c) {
983    return Bytes.toString(
984        c.getQualifierArray(), c.getQualifierOffset() + QUOTA_SNAPSHOT_SIZE_QUALIFIER.length,
985        c.getQualifierLength() - QUOTA_SNAPSHOT_SIZE_QUALIFIER.length);
986  }
987
988  protected static long extractSnapshotSize(
989      byte[] data, int offset, int length) throws InvalidProtocolBufferException {
990    ByteString byteStr = UnsafeByteOperations.unsafeWrap(data, offset, length);
991    return org.apache.hadoop.hbase.shaded.protobuf.generated.QuotaProtos.SpaceQuotaSnapshot
992        .parseFrom(byteStr).getQuotaUsage();
993  }
994}