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