From 08991b2aa3740732a88cdc48facf41210c9251ed Mon Sep 17 00:00:00 2001
From: Ray Mattingly <rmattingly@hubspot.com>
Date: Fri, 1 Dec 2023 17:10:13 -0500
Subject: [PATCH] simplify, protect against overflow

---
 .../ReopenTableRegionsProcedure.java          | 25 +++++++++++++++----
 ...stReopenTableRegionsProcedureBatching.java | 13 +++++++++-
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/procedure/ReopenTableRegionsProcedure.java b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/procedure/ReopenTableRegionsProcedure.java
index 7a9d0f655190..d50fa49506dd 100644
--- a/hbase-server/src/main/java/org/apache/hadoop/hbase/master/procedure/ReopenTableRegionsProcedure.java
+++ b/hbase-server/src/main/java/org/apache/hadoop/hbase/master/procedure/ReopenTableRegionsProcedure.java
@@ -106,10 +106,13 @@ public ReopenTableRegionsProcedure(final TableName tableName, final List<byte[]>
     this.tableName = tableName;
     this.regionNames = regionNames;
     this.reopenBatchBackoffMillis = reopenBatchBackoffMillis;
-    this.reopenBatchSize = reopenBatchSizeMax != PROGRESSIVE_BATCH_SIZE_MAX_DISABLED
-      ? 1
-      : PROGRESSIVE_BATCH_SIZE_MAX_DEFAULT_VALUE;
-    this.reopenBatchSizeMax = Math.max(reopenBatchSizeMax, MINIMUM_BATCH_SIZE_MAX);
+    if (reopenBatchSizeMax == PROGRESSIVE_BATCH_SIZE_MAX_DISABLED) {
+      this.reopenBatchSize = Integer.MAX_VALUE;
+      this.reopenBatchSizeMax = Integer.MAX_VALUE;
+    } else {
+      this.reopenBatchSize = 1;
+      this.reopenBatchSizeMax = Math.max(reopenBatchSizeMax, MINIMUM_BATCH_SIZE_MAX);
+    }
   }
 
   @Override
@@ -134,6 +137,18 @@ public long getBatchesProcessed() {
     return batchesProcessed;
   }
 
+  @RestrictedApi(explanation = "Should only be called in tests", link = "",
+      allowedOnPath = ".*/src/test/.*")
+  protected int progressBatchSize() {
+    int previousBatchSize = reopenBatchSize;
+    reopenBatchSize = Math.min(reopenBatchSizeMax, 2 * reopenBatchSize);
+    if (reopenBatchSize < previousBatchSize) {
+      // the batch size should never decrease. this must be overflow, so just use max
+      reopenBatchSize = reopenBatchSizeMax;
+    }
+    return reopenBatchSize;
+  }
+
   private boolean canSchedule(MasterProcedureEnv env, HRegionLocation loc) {
     if (loc.getSeqNum() < 0) {
       return false;
@@ -224,7 +239,7 @@ private Flow reopenIfSchedulable(MasterProcedureEnv env, List<HRegionLocation> r
       retryCounter = null;
       setNextState(ReopenTableRegionsState.REOPEN_TABLE_REGIONS_REOPEN_REGIONS);
       if (shouldBatchBackoff && reopenBatchBackoffMillis > 0) {
-        reopenBatchSize = Math.min(reopenBatchSizeMax, 2 * reopenBatchSize);
+        progressBatchSize();
         setBackoffState(reopenBatchBackoffMillis);
         throw new ProcedureSuspendedException();
       } else {
diff --git a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/procedure/TestReopenTableRegionsProcedureBatching.java b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/procedure/TestReopenTableRegionsProcedureBatching.java
index d98d2f4fe926..8ea9b3c6a309 100644
--- a/hbase-server/src/test/java/org/apache/hadoop/hbase/master/procedure/TestReopenTableRegionsProcedureBatching.java
+++ b/hbase-server/src/test/java/org/apache/hadoop/hbase/master/procedure/TestReopenTableRegionsProcedureBatching.java
@@ -134,7 +134,7 @@ public void testNegativeBatchSizeDoesNotBreak() throws IOException {
     Set<StuckRegion> stuckRegions =
       regions.stream().map(r -> stickRegion(am, procExec, r)).collect(Collectors.toSet());
     ReopenTableRegionsProcedure proc =
-      new ReopenTableRegionsProcedure(TABLE_NAME, BACKOFF_MILLIS_PER_RS, -1);
+      new ReopenTableRegionsProcedure(TABLE_NAME, BACKOFF_MILLIS_PER_RS, -100);
     procExec.submitProcedure(proc);
     UTIL.waitFor(10000, () -> proc.getState() == ProcedureState.WAITING_TIMEOUT);
 
@@ -149,6 +149,17 @@ public void testNegativeBatchSizeDoesNotBreak() throws IOException {
     assertEquals(proc.getRegionsReopened(), regions.size());
   }
 
+  @Test
+  public void testBatchSizeDoesNotOverflow() {
+    ReopenTableRegionsProcedure proc =
+      new ReopenTableRegionsProcedure(TABLE_NAME, BACKOFF_MILLIS_PER_RS, Integer.MAX_VALUE);
+    int currentBatchSize = 1;
+    while (currentBatchSize < Integer.MAX_VALUE) {
+      currentBatchSize = proc.progressBatchSize();
+      assertTrue(currentBatchSize > 0);
+    }
+  }
+
   private void confirmBatchSize(int expectedBatchSize, Set<StuckRegion> stuckRegions,
     ReopenTableRegionsProcedure proc) {
     while (true) {