HBASE-13593 Quota support for namespace should take snapshot restore and clone into account (Ashish Singhi)

hbase-server/src/main/java/org/apache/hadoop/hbase/master/snapshot/SnapshotManager.java

@@ -734,7 +734,16 @@ public class SnapshotManager extends MasterProcedureManager implements Stoppable
       if (cpHost != null) {
         cpHost.preRestoreSnapshot(reqSnapshot, snapshotTableDesc);
       }
-      restoreSnapshot(snapshot, snapshotTableDesc);
+      try {
+        // Table already exist. Check and update the region quota for this table namespace
+        checkAndUpdateNamespaceRegionQuota(manifest, tableName);
+        restoreSnapshot(snapshot, snapshotTableDesc);
+      } catch (IOException e) {
+        this.master.getMasterQuotaManager().removeTableFromNamespaceQuota(tableName);
+        LOG.error("Exception occurred while restoring the snapshot " + snapshot.getName()
+            + " as table " + tableName.getNameAsString(), e);
+        throw e;
+      }
       LOG.info("Restore snapshot=" + snapshot.getName() + " as table=" + tableName);
 
       if (cpHost != null) {
@@ -745,7 +754,15 @@ public class SnapshotManager extends MasterProcedureManager implements Stoppable
       if (cpHost != null) {
         cpHost.preCloneSnapshot(reqSnapshot, htd);
       }
-      cloneSnapshot(snapshot, htd);
+      try {
+        checkAndUpdateNamespaceQuota(manifest, tableName);
+        cloneSnapshot(snapshot, htd);
+      } catch (IOException e) {
+        this.master.getMasterQuotaManager().removeTableFromNamespaceQuota(tableName);
+        LOG.error("Exception occurred while cloning the snapshot " + snapshot.getName()
+            + " as table " + tableName.getNameAsString(), e);
+        throw e;
+      }
       LOG.info("Clone snapshot=" + snapshot.getName() + " as table=" + tableName);
 
       if (cpHost != null) {
@@ -754,6 +771,22 @@ public class SnapshotManager extends MasterProcedureManager implements Stoppable
     }
   }
   
+  private void checkAndUpdateNamespaceQuota(SnapshotManifest manifest, TableName tableName)
+      throws IOException {
+    if (this.master.getMasterQuotaManager().isQuotaEnabled()) {
+      this.master.getMasterQuotaManager().checkNamespaceTableAndRegionQuota(tableName,
+        manifest.getRegionManifestsMap().size());
+    }
+  }
+
+  private void checkAndUpdateNamespaceRegionQuota(SnapshotManifest manifest, TableName tableName)
+      throws IOException {
+    if (this.master.getMasterQuotaManager().isQuotaEnabled()) {
+      this.master.getMasterQuotaManager().checkAndUpdateNamespaceRegionQuota(tableName,
+        manifest.getRegionManifestsMap().size());
+    }
+  }
+
   /**
    * Restore the specified snapshot.
    * The restore will fail if the destination table has a snapshot or restore in progress.

hbase-server/src/main/java/org/apache/hadoop/hbase/namespace/NamespaceAuditor.java

@@ -69,6 +69,20 @@ public class NamespaceAuditor {
     }
   }
   
+  /**
+   * Check and update region count quota for an existing table.
+   * @param tName - table name for which region count to be updated.
+   * @param regions - Number of regions that will be added.
+   * @throws IOException Signals that an I/O exception has occurred.
+   */
+  public void checkQuotaToUpdateRegion(TableName tName, int regions) throws IOException {
+    if (stateManager.isInitialized()) {
+      stateManager.checkAndUpdateNamespaceRegionCount(tName, regions);
+    } else {
+      checkTableTypeAndThrowException(tName);
+    }
+  }
+
   private void checkTableTypeAndThrowException(TableName name) throws IOException {
     if (name.isSystemTable()) {
       LOG.debug("Namespace auditor checks not performed for table " + name.getNameAsString());

hbase-server/src/main/java/org/apache/hadoop/hbase/namespace/NamespaceStateManager.java

@@ -109,6 +109,33 @@ class NamespaceStateManager extends ZooKeeperListener {
     return true;
   }
   
+  /**
+   * Check and update region count for an existing table. To handle scenarios like restore snapshot
+   * @param TableName name of the table for region count needs to be checked and updated
+   * @param incr count of regions
+   * @throws QuotaExceededException if quota exceeds for the number of regions allowed in a
+   *           namespace
+   * @throws IOException Signals that an I/O exception has occurred.
+   */
+  synchronized void checkAndUpdateNamespaceRegionCount(TableName name, int incr) 
+      throws IOException {
+    String namespace = name.getNamespaceAsString();
+    NamespaceDescriptor nspdesc = getNamespaceDescriptor(namespace);
+    if (nspdesc != null) {
+      NamespaceTableAndRegionInfo currentStatus = getState(namespace);
+      int regionCountOfTable = currentStatus.getRegionCountOfTable(name);
+      if ((currentStatus.getRegionCount() - regionCountOfTable + incr) > TableNamespaceManager
+          .getMaxRegions(nspdesc)) {
+        throw new QuotaExceededException("The table " + name.getNameAsString()
+            + " region count cannot be updated as it would exceed maximum number "
+            + "of regions allowed in the namespace.  The total number of regions permitted is "
+            + TableNamespaceManager.getMaxRegions(nspdesc));
+      }
+      currentStatus.removeTable(name);
+      currentStatus.addTable(name, incr);
+    }
+  }
+
   private NamespaceDescriptor getNamespaceDescriptor(String namespaceAsString) {
     try {
       return this.master.getNamespaceDescriptor(namespaceAsString);

hbase-server/src/main/java/org/apache/hadoop/hbase/quotas/MasterQuotaManager.java

@@ -326,6 +326,12 @@ public class MasterQuotaManager implements RegionStateListener {
     }
   }
   
+  public void checkAndUpdateNamespaceRegionQuota(TableName tName, int regions) throws IOException {
+    if (enabled) {
+      namespaceQuotaManager.checkQuotaToUpdateRegion(tName, regions);
+    }
+  }
+
   public void onRegionMerged(HRegionInfo hri) throws IOException {
     if (enabled) {
       namespaceQuotaManager.updateQuotaForRegionMerge(hri);

hbase-server/src/test/java/org/apache/hadoop/hbase/namespace/TestNamespaceAuditor.java

@@ -46,6 +46,7 @@ import org.apache.hadoop.hbase.client.ConnectionFactory;
 import org.apache.hadoop.hbase.client.HBaseAdmin;
 import org.apache.hadoop.hbase.client.HTable;
 import org.apache.hadoop.hbase.client.Mutation;
+import org.apache.hadoop.hbase.client.RegionLocator;
 import org.apache.hadoop.hbase.client.Table;
 import org.apache.hadoop.hbase.coprocessor.BaseMasterObserver;
 import org.apache.hadoop.hbase.coprocessor.BaseRegionObserver;
@@ -68,6 +69,7 @@ import org.apache.hadoop.hbase.regionserver.HRegion;
 import org.apache.hadoop.hbase.regionserver.HRegionServer;
 import org.apache.hadoop.hbase.regionserver.Region;
 import org.apache.hadoop.hbase.regionserver.RegionServerCoprocessorHost;
+import org.apache.hadoop.hbase.snapshot.RestoreSnapshotException;
 import org.apache.hadoop.hbase.testclassification.MediumTests;
 import org.apache.hadoop.hbase.util.Bytes;
 import org.apache.hadoop.hbase.util.FSUtils;
@@ -679,4 +681,131 @@ public class TestNamespaceAuditor {
     ADMIN.createTable(tableDescOne);
     ADMIN.createTable(tableDescTwo, Bytes.toBytes("AAA"), Bytes.toBytes("ZZZ"), 4);
   }
+  
+  @Test(expected = QuotaExceededException.class, timeout = 30000)
+  public void testCloneSnapshotQuotaExceed() throws Exception {
+    String nsp = prefix + "_testTableQuotaExceedWithCloneSnapshot";
+    NamespaceDescriptor nspDesc =
+        NamespaceDescriptor.create(nsp).addConfiguration(TableNamespaceManager.KEY_MAX_TABLES, "1")
+            .build();
+    ADMIN.createNamespace(nspDesc);
+    assertNotNull("Namespace descriptor found null.", ADMIN.getNamespaceDescriptor(nsp));
+    TableName tableName = TableName.valueOf(nsp + TableName.NAMESPACE_DELIM + "table1");
+    TableName cloneTableName = TableName.valueOf(nsp + TableName.NAMESPACE_DELIM + "table2");
+    HTableDescriptor tableDescOne = new HTableDescriptor(tableName);
+    ADMIN.createTable(tableDescOne);
+    String snapshot = "snapshot_testTableQuotaExceedWithCloneSnapshot";
+    ADMIN.snapshot(snapshot, tableName);
+    ADMIN.cloneSnapshot(snapshot, cloneTableName);
+    ADMIN.deleteSnapshot(snapshot);
+  }
+
+  @Test(timeout = 180000)
+  public void testCloneSnapshot() throws Exception {
+    String nsp = prefix + "_testCloneSnapshot";
+    NamespaceDescriptor nspDesc =
+        NamespaceDescriptor.create(nsp).addConfiguration(TableNamespaceManager.KEY_MAX_TABLES, "2")
+            .addConfiguration(TableNamespaceManager.KEY_MAX_REGIONS, "20").build();
+    ADMIN.createNamespace(nspDesc);
+    assertNotNull("Namespace descriptor found null.", ADMIN.getNamespaceDescriptor(nsp));
+    TableName tableName = TableName.valueOf(nsp + TableName.NAMESPACE_DELIM + "table1");
+    TableName cloneTableName = TableName.valueOf(nsp + TableName.NAMESPACE_DELIM + "table2");
+    HTableDescriptor tableDescOne = new HTableDescriptor(tableName);
+
+    ADMIN.createTable(tableDescOne, Bytes.toBytes("AAA"), Bytes.toBytes("ZZZ"), 4);
+    String snapshot = "snapshot_testCloneSnapshot";
+    ADMIN.snapshot(snapshot, tableName);
+    ADMIN.cloneSnapshot(snapshot, cloneTableName);
+
+    int tableLength;
+    try (RegionLocator locator = ADMIN.getConnection().getRegionLocator(tableName)) {
+      tableLength = locator.getStartKeys().length;
+    }
+    assertEquals(tableName.getNameAsString() + " should have four regions.", 4, tableLength);
+
+    try (RegionLocator locator = ADMIN.getConnection().getRegionLocator(cloneTableName)) {
+      tableLength = locator.getStartKeys().length;
+    }
+    assertEquals(cloneTableName.getNameAsString() + " should have four regions.", 4, tableLength);
+
+    NamespaceTableAndRegionInfo nstate = getNamespaceState(nsp);
+    assertEquals("Total tables count should be 2.", 2, nstate.getTables().size());
+    assertEquals("Total regions count should be.", 8, nstate.getRegionCount());
+
+    ADMIN.deleteSnapshot(snapshot);
+  }
+
+  @Test(timeout = 180000)
+  public void testRestoreSnapshot() throws Exception {
+    String nsp = prefix + "_testRestoreSnapshot";
+    NamespaceDescriptor nspDesc =
+        NamespaceDescriptor.create(nsp)
+            .addConfiguration(TableNamespaceManager.KEY_MAX_REGIONS, "10").build();
+    ADMIN.createNamespace(nspDesc);
+    assertNotNull("Namespace descriptor found null.", ADMIN.getNamespaceDescriptor(nsp));
+    TableName tableName1 = TableName.valueOf(nsp + TableName.NAMESPACE_DELIM + "table1");
+    HTableDescriptor tableDescOne = new HTableDescriptor(tableName1);
+    ADMIN.createTable(tableDescOne, Bytes.toBytes("AAA"), Bytes.toBytes("ZZZ"), 4);
+
+    NamespaceTableAndRegionInfo nstate = getNamespaceState(nsp);
+    assertEquals("Intial region count should be 4.", 4, nstate.getRegionCount());
+
+    String snapshot = "snapshot_testRestoreSnapshot";
+    ADMIN.snapshot(snapshot, tableName1);
+
+    List<HRegionInfo> regions = ADMIN.getTableRegions(tableName1);
+    Collections.sort(regions);
+
+    ADMIN.split(tableName1, Bytes.toBytes("JJJ"));
+    Thread.sleep(2000);
+    assertEquals("Total regions count should be 5.", 5, nstate.getRegionCount());
+
+    ADMIN.disableTable(tableName1);
+    ADMIN.restoreSnapshot(snapshot);
+
+    assertEquals("Total regions count should be 4 after restore.", 4, nstate.getRegionCount());
+
+    ADMIN.enableTable(tableName1);
+    ADMIN.deleteSnapshot(snapshot);
+  }
+
+  @Test(timeout = 180000)
+  public void testRestoreSnapshotQuotaExceed() throws Exception {
+    String nsp = prefix + "_testRestoreSnapshotQuotaExceed";
+    NamespaceDescriptor nspDesc =
+        NamespaceDescriptor.create(nsp)
+            .addConfiguration(TableNamespaceManager.KEY_MAX_REGIONS, "10").build();
+    ADMIN.createNamespace(nspDesc);
+    NamespaceDescriptor ndesc = ADMIN.getNamespaceDescriptor(nsp);
+    assertNotNull("Namespace descriptor found null.", ndesc);
+    TableName tableName1 = TableName.valueOf(nsp + TableName.NAMESPACE_DELIM + "table1");
+    HTableDescriptor tableDescOne = new HTableDescriptor(tableName1);
+    ADMIN.createTable(tableDescOne, Bytes.toBytes("AAA"), Bytes.toBytes("ZZZ"), 4);
+
+    NamespaceTableAndRegionInfo nstate = getNamespaceState(nsp);
+    assertEquals("Intial region count should be 4.", 4, nstate.getRegionCount());
+
+    String snapshot = "snapshot_testRestoreSnapshotQuotaExceed";
+    ADMIN.snapshot(snapshot, tableName1);
+
+    List<HRegionInfo> regions = ADMIN.getTableRegions(tableName1);
+    Collections.sort(regions);
+
+    ADMIN.split(tableName1, Bytes.toBytes("JJJ"));
+    Thread.sleep(2000);
+    assertEquals("Total regions count should be 5.", 5, nstate.getRegionCount());
+
+    ndesc.setConfiguration(TableNamespaceManager.KEY_MAX_REGIONS, "2");
+    ADMIN.modifyNamespace(ndesc);
+
+    ADMIN.disableTable(tableName1);
+    try {
+      ADMIN.restoreSnapshot(snapshot);
+      fail("Region quota is exceeded so QuotaExceededException should be thrown but HBaseAdmin"
+          + " wraps IOException into RestoreSnapshotException");
+    } catch (RestoreSnapshotException ignore) {
+    }
+    ADMIN.enableTable(tableName1);
+    ADMIN.deleteSnapshot(snapshot);
+  }
 }