HBASE-29281 Atomic request throttles are missing QuotaSettingsFactory support (#6953) (#6956)

Signed-off-by: Nick Dimiduk <[email protected]>
Co-authored-by: Ray Mattingly <[email protected]>

Author: rmdmattingly

hbase-client/src/main/java/org/apache/hadoop/hbase/quotas/QuotaSettingsFactory.java

@@ -168,6 +168,18 @@ protected static List<ThrottleSettings> fromThrottle(final String userName,
       settings.add(ThrottleSettings.fromTimedQuota(userName, tableName, namespace, regionServer,
         ThrottleType.WRITE_CAPACITY_UNIT, throttle.getWriteCapacityUnit()));
     }
+    if (throttle.hasAtomicReadSize()) {
+      settings.add(ThrottleSettings.fromTimedQuota(userName, tableName, namespace, regionServer,
+        ThrottleType.ATOMIC_READ_SIZE, throttle.getAtomicReadSize()));
+    }
+    if (throttle.hasAtomicWriteSize()) {
+      settings.add(ThrottleSettings.fromTimedQuota(userName, tableName, namespace, regionServer,
+        ThrottleType.ATOMIC_WRITE_SIZE, throttle.getAtomicWriteSize()));
+    }
+    if (throttle.hasAtomicReqNum()) {
+      settings.add(ThrottleSettings.fromTimedQuota(userName, tableName, namespace, regionServer,
+        ThrottleType.ATOMIC_REQUEST_NUMBER, throttle.getAtomicReqNum()));
+    }
     return settings;
   }
 

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

@@ -26,7 +26,10 @@
 import static org.apache.hadoop.hbase.quotas.ThrottleQuotaTestUtil.triggerUserCacheRefresh;
 import static org.apache.hadoop.hbase.quotas.ThrottleQuotaTestUtil.waitMinuteQuota;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
+import java.io.IOException;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 import org.apache.hadoop.hbase.HBaseClassTestRule;
 import org.apache.hadoop.hbase.HBaseTestingUtil;
@@ -652,4 +655,39 @@ public void testExceedThrottleQuota() throws Exception {
     admin.setQuota(QuotaSettingsFactory.unthrottleTable(TABLE_NAMES[0]));
     triggerTableCacheRefresh(TEST_UTIL, true, TABLE_NAMES[0]);
   }
+
+  @Test
+  public void testSetAndGetAllThrottleTypes() throws Exception {
+    for (ThrottleType throttleType : ThrottleType.values()) {
+      canSetAndGetUserThrottle(throttleType);
+    }
+  }
+
+  private void canSetAndGetUserThrottle(ThrottleType throttleType) throws IOException {
+    final Admin admin = TEST_UTIL.getAdmin();
+    final String userName = User.getCurrent().getShortName();
+
+    QuotaSettings setQuota =
+      QuotaSettingsFactory.throttleUser(userName, throttleType, 123, TimeUnit.SECONDS);
+    admin.setQuota(setQuota);
+
+    boolean found = false;
+    List<QuotaSettings> quotaSettings = admin.getQuota(new QuotaFilter().setUserFilter(userName));
+    for (QuotaSettings settings : quotaSettings) {
+      if (settings instanceof ThrottleSettings) {
+        ThrottleSettings throttle = (ThrottleSettings) settings;
+        if (
+          userName.equals(throttle.getUserName()) && throttle.getThrottleType() == throttleType
+            && throttle.getSoftLimit() == 123 && throttle.getTimeUnit() == TimeUnit.SECONDS
+        ) {
+          found = true;
+          break;
+        }
+      }
+    }
+
+    assertTrue("Expected to find " + throttleType.name() + " quota for user " + userName, found);
+    admin.setQuota(QuotaSettingsFactory.unthrottleUserByThrottleType(userName, throttleType));
+  }
+
 }