HBASE-15433 SnapshotManager#restoreSnapshot not update table and region count quota correctly when encountering exception (Jianwei Cui)

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

@@ -65,6 +65,7 @@ import org.apache.hadoop.hbase.protobuf.generated.HBaseProtos.NameStringPair;
 import org.apache.hadoop.hbase.protobuf.generated.HBaseProtos.ProcedureDescription;
 import org.apache.hadoop.hbase.protobuf.generated.HBaseProtos.SnapshotDescription;
 import org.apache.hadoop.hbase.protobuf.generated.HBaseProtos.SnapshotDescription.Type;
+import org.apache.hadoop.hbase.quotas.QuotaExceededException;
 import org.apache.hadoop.hbase.security.AccessDeniedException;
 import org.apache.hadoop.hbase.security.User;
 import org.apache.hadoop.hbase.snapshot.ClientSnapshotDescriptionUtils;
@@ -724,12 +725,41 @@ public class SnapshotManager extends MasterProcedureManager implements Stoppable
       if (cpHost != null) {
         cpHost.preRestoreSnapshot(reqSnapshot, snapshotTableDesc);
       }
+
+      int tableRegionCount = -1;
       try {
-        // Table already exist. Check and update the region quota for this table namespace
+        // Table already exist. Check and update the region quota for this table namespace.
-        checkAndUpdateNamespaceRegionQuota(manifest, tableName);
+        // The region quota may not be updated correctly if there are concurrent restore snapshot
+        // requests for the same table
+
+        tableRegionCount = getRegionCountOfTable(tableName);
+        int snapshotRegionCount = manifest.getRegionManifestsMap().size();
+
+        // Update region quota when snapshotRegionCount is larger. If we updated the region count
+        // to a smaller value before retoreSnapshot and the retoreSnapshot fails, we may fail to
+        // reset the region count to its original value if the region quota is consumed by other
+        // tables in the namespace
+        if (tableRegionCount > 0 && tableRegionCount < snapshotRegionCount) {
+          checkAndUpdateNamespaceRegionQuota(snapshotRegionCount, tableName);
+        }
         restoreSnapshot(snapshot, snapshotTableDesc);
+        // Update the region quota if snapshotRegionCount is smaller. This step should not fail
+        // because we have reserved enough region quota before hand
+        if (tableRegionCount > 0 && tableRegionCount > snapshotRegionCount) {
+          checkAndUpdateNamespaceRegionQuota(snapshotRegionCount, tableName);
+        }
+      } catch (QuotaExceededException e) {
+        LOG.error("Region quota exceeded while restoring the snapshot " + snapshot.getName()
+          + " as table " + tableName.getNameAsString(), e);
+        // If QEE is thrown before restoreSnapshot, quota information is not updated, so we
+        // should throw the exception directly. If QEE is thrown after restoreSnapshot, there
+        // must be unexpected reasons, we also throw the exception directly
+        throw e;
       } catch (IOException e) {
-        this.master.getMasterQuotaManager().removeTableFromNamespaceQuota(tableName);
+        if (tableRegionCount > 0) {
+          // reset the region count for table
+          checkAndUpdateNamespaceRegionQuota(tableRegionCount, tableName);
+        }
         LOG.error("Exception occurred while restoring the snapshot " + snapshot.getName()
             + " as table " + tableName.getNameAsString(), e);
         throw e;
@@ -769,14 +799,24 @@ public class SnapshotManager extends MasterProcedureManager implements Stoppable
     }
   }
 
-  private void checkAndUpdateNamespaceRegionQuota(SnapshotManifest manifest, TableName tableName)
+  private void checkAndUpdateNamespaceRegionQuota(int updatedRegionCount, TableName tableName)
       throws IOException {
     if (this.master.getMasterQuotaManager().isQuotaEnabled()) {
       this.master.getMasterQuotaManager().checkAndUpdateNamespaceRegionQuota(tableName,
-        manifest.getRegionManifestsMap().size());
+        updatedRegionCount);
     }
   }
 
+  /**
+   * @return cached region count, or -1 if quota manager is disabled or table status not found
+  */
+  private int getRegionCountOfTable(TableName tableName) throws IOException {
+    if (this.master.getMasterQuotaManager().isQuotaEnabled()) {
+      return this.master.getMasterQuotaManager().getRegionCountOfTable(tableName);
+    }
+    return -1;
+  }
+
   /**
    * 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

@@ -100,6 +100,21 @@ public class NamespaceAuditor {
     }
   }
 
+  /**
+   * Get region count for table
+   * @param tName - table name
+   * @return cached region count, or -1 if table status not found
+   * @throws IOException Signals that the namespace auditor has not been initialized
+   */
+  public int getRegionCountOfTable(TableName tName) throws IOException {
+    if (stateManager.isInitialized()) {
+      NamespaceTableAndRegionInfo state = stateManager.getState(tName.getNamespaceAsString());
+      return state != null ? state.getRegionCountOfTable(tName) : -1;
+    }
+    checkTableTypeAndThrowException(tName);
+    return -1;
+  }
+
   public void checkQuotaToSplitRegion(HRegionInfo hri) throws IOException {
     if (!stateManager.isInitialized()) {
       throw new IOException(

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

@@ -320,6 +320,16 @@ public class MasterQuotaManager implements RegionStateListener {
     }
   }
 
+  /**
+   * @return cached region count, or -1 if quota manager is disabled or table status not found
+  */
+  public int getRegionCountOfTable(TableName tName) throws IOException {
+    if (enabled) {
+      return namespaceQuotaManager.getRegionCountOfTable(tName);
+    }
+    return -1;
+  }
+
   public void onRegionMerged(HRegionInfo hri) throws IOException {
     if (enabled) {
       namespaceQuotaManager.updateQuotaForRegionMerge(hri);

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

@@ -827,16 +827,13 @@ public class TestNamespaceAuditor {
     assertEquals("Intial region count should be 4.", 4, nstate.getRegionCount());
 
     String snapshot = "snapshot_testRestoreSnapshotQuotaExceed";
+    // snapshot has 4 regions
     ADMIN.snapshot(snapshot, tableName1);
-
+    // recreate table with 1 region and set max regions to 3 for namespace
-    List<HRegionInfo> regions = ADMIN.getTableRegions(tableName1);
+    ADMIN.disableTable(tableName1);
-    Collections.sort(regions);
+    ADMIN.deleteTable(tableName1);
-
+    ADMIN.createTable(tableDescOne);
-    ADMIN.split(tableName1, Bytes.toBytes("JJJ"));
+    ndesc.setConfiguration(TableNamespaceManager.KEY_MAX_REGIONS, "3");
-    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);
@@ -845,7 +842,9 @@ public class TestNamespaceAuditor {
       fail("Region quota is exceeded so QuotaExceededException should be thrown but HBaseAdmin"
           + " wraps IOException into RestoreSnapshotException");
     } catch (RestoreSnapshotException ignore) {
+      assertTrue(ignore.getCause() instanceof QuotaExceededException);
     }
+    assertEquals(1, getNamespaceState(nsp).getRegionCount());
     ADMIN.enableTable(tableName1);
     ADMIN.deleteSnapshot(snapshot);
   }