HBASE-22013 SpaceQuotas - getNumRegions() returning wrong number of regions due to region replicas

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

@@ -32,6 +32,7 @@
 import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.client.Connection;
 import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.RegionReplicaUtil;
 import org.apache.hadoop.hbase.client.Scan;
 import org.apache.hadoop.hbase.master.HMaster;
 import org.apache.hadoop.hbase.master.MetricsMaster;
@@ -765,6 +766,8 @@ int getNumRegions(TableName table) throws IOException {
       if (regions == null) {
         return 0;
       }
+      // Filter the region replicas if any and return the original number of regions for a table.
+      RegionReplicaUtil.removeNonDefaultRegions(regions);
       return regions.size();
     }
 

hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/SpaceQuotaHelperForTests.java

@@ -477,15 +477,16 @@ TableName createTableWithRegions(int numRegions) throws Exception {
   }
 
   TableName createTableWithRegions(Admin admin, int numRegions) throws Exception {
-    return createTableWithRegions(
-        testUtil.getAdmin(), NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions);
+    return createTableWithRegions(admin, NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, numRegions,
+        0);
   }
 
   TableName createTableWithRegions(String namespace, int numRegions) throws Exception {
-    return createTableWithRegions(testUtil.getAdmin(), namespace, numRegions);
+    return createTableWithRegions(testUtil.getAdmin(), namespace, numRegions, 0);
   }
 
-  TableName createTableWithRegions(Admin admin, String namespace, int numRegions) throws Exception {
+  TableName createTableWithRegions(Admin admin, String namespace, int numRegions,
+      int numberOfReplicas) throws Exception {
     final TableName tn = getNextTableName(namespace);
 
     // Delete the old table
@@ -495,8 +496,14 @@ TableName createTableWithRegions(Admin admin, String namespace, int numRegions)
     }
 
     // Create the table
-    TableDescriptor tableDesc = TableDescriptorBuilder.newBuilder(tn)
-        .setColumnFamily(ColumnFamilyDescriptorBuilder.of(F1)).build();
+    TableDescriptor tableDesc;
+    if (numberOfReplicas > 0) {
+      tableDesc = TableDescriptorBuilder.newBuilder(tn).setRegionReplication(numberOfReplicas)
+          .setColumnFamily(ColumnFamilyDescriptorBuilder.of(F1)).build();
+    } else {
+      tableDesc = TableDescriptorBuilder.newBuilder(tn)
+          .setColumnFamily(ColumnFamilyDescriptorBuilder.of(F1)).build();
+    }
     if (numRegions == 1) {
       admin.createTable(tableDesc);
     } else {

hbase-server/src/test/java/org/apache/hadoop/hbase/quotas/TestSpaceQuotasWithRegionReplicas.java

@@ -0,0 +1,112 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hbase.quotas;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.HBaseClassTestRule;
+import org.apache.hadoop.hbase.HBaseTestingUtility;
+import org.apache.hadoop.hbase.NamespaceDescriptor;
+import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.client.Put;
+import org.apache.hadoop.hbase.testclassification.MediumTests;
+import org.apache.hadoop.hbase.util.Bytes;
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.rules.TestName;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+@Category(MediumTests.class)
+public class TestSpaceQuotasWithRegionReplicas {
+
+  @ClassRule
+  public static final HBaseClassTestRule CLASS_RULE =
+      HBaseClassTestRule.forClass(TestSpaceQuotasWithRegionReplicas.class);
+
+  private static final Logger LOG =
+      LoggerFactory.getLogger(TestSpaceQuotasWithRegionReplicas.class);
+  private static final HBaseTestingUtility TEST_UTIL = new HBaseTestingUtility();
+
+  @Rule
+  public TestName testName = new TestName();
+  private SpaceQuotaHelperForTests helper;
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    Configuration conf = TEST_UTIL.getConfiguration();
+    SpaceQuotaHelperForTests.updateConfigForQuotas(conf);
+    TEST_UTIL.startMiniCluster(1);
+  }
+
+  @AfterClass
+  public static void tearDown() throws Exception {
+    TEST_UTIL.shutdownMiniCluster();
+  }
+
+  @Before
+  public void removeAllQuotas() throws Exception {
+    helper = new SpaceQuotaHelperForTests(TEST_UTIL, testName, new AtomicLong(0));
+    helper.removeAllQuotas();
+  }
+
+  @Test
+  public void testSetQuotaWithRegionReplicaSingleRegion() throws Exception {
+    for (SpaceViolationPolicy policy : SpaceViolationPolicy.values()) {
+      setQuotaAndVerifyForRegionReplication(1, 2, policy);
+    }
+  }
+
+  @Test
+  public void testSetQuotaWithRegionReplicaMultipleRegion() throws Exception {
+    for (SpaceViolationPolicy policy : SpaceViolationPolicy.values()) {
+      setQuotaAndVerifyForRegionReplication(6, 3, policy);
+    }
+  }
+
+  @Test
+  public void testSetQuotaWithSingleRegionZeroRegionReplica() throws Exception {
+    for (SpaceViolationPolicy policy : SpaceViolationPolicy.values()) {
+      setQuotaAndVerifyForRegionReplication(1, 0, policy);
+    }
+  }
+
+  @Test
+  public void testSetQuotaWithMultipleRegionZeroRegionReplicas() throws Exception {
+    for (SpaceViolationPolicy policy : SpaceViolationPolicy.values()) {
+      setQuotaAndVerifyForRegionReplication(6, 0, policy);
+    }
+  }
+
+  private void setQuotaAndVerifyForRegionReplication(int region, int replicatedRegion,
+      SpaceViolationPolicy policy) throws Exception {
+    TableName tn = helper.createTableWithRegions(TEST_UTIL.getAdmin(),
+        NamespaceDescriptor.DEFAULT_NAMESPACE_NAME_STR, region, replicatedRegion);
+    helper.setQuotaLimit(tn, policy, 5L);
+    helper.writeData(tn, 5L * SpaceQuotaHelperForTests.ONE_MEGABYTE);
+    Put p = new Put(Bytes.toBytes("to_reject"));
+    p.addColumn(Bytes.toBytes(SpaceQuotaHelperForTests.F1), Bytes.toBytes("to"),
+        Bytes.toBytes("reject"));
+    helper.verifyViolation(policy, tn, p);
+  }
+}