Source code

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