Retry peer-list reset during AddRegionPeer rollback

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/region/AddRegionPeerProcedure.java

@@ -55,6 +55,8 @@
 
 public class AddRegionPeerProcedure extends RegionOperationProcedure<AddRegionPeerState> {
   private static final Logger LOGGER = LoggerFactory.getLogger(AddRegionPeerProcedure.class);
+  private static final int MAX_RESET_PEER_LIST_RETRY = 3;
+  private static final long RESET_PEER_LIST_RETRY_INTERVAL_MS = 5_000;
 
   private TDataNodeLocation coordinator;
 
@@ -252,6 +254,21 @@ private Flow warnAndRollBackAndNoMoreState(
 
     Map<Integer, TSStatus> resultMap =
         handler.resetPeerList(regionId, correctDataNodeLocations, relatedDataNodeLocationMap);
+    for (int retry = 0; retry < MAX_RESET_PEER_LIST_RETRY; retry++) {
+      if (resultMap.keySet().containsAll(relatedDataNodeLocationMap.keySet())
+          && resultMap.values().stream()
+              .allMatch(status -> status.getCode() == SUCCESS_STATUS.getStatusCode())) {
+        break;
+      }
+      try {
+        Thread.sleep(RESET_PEER_LIST_RETRY_INTERVAL_MS);
+      } catch (InterruptedException interruptedException) {
+        Thread.currentThread().interrupt();
+        break;
+      }
+      resultMap =
+          handler.resetPeerList(regionId, correctDataNodeLocations, relatedDataNodeLocationMap);
+    }
 
     resultMap.forEach(
         (dataNodeId, resetResult) -> {

iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/impl/region/AddRegionPeerProcedureTest.java

@@ -23,16 +23,66 @@
 import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
 import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
 import org.apache.iotdb.common.rpc.thrift.TEndPoint;
+import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.confignode.manager.ConfigManager;
+import org.apache.iotdb.confignode.manager.partition.PartitionManager;
+import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv;
+import org.apache.iotdb.confignode.procedure.env.RegionMaintainHandler;
+import org.apache.iotdb.confignode.procedure.impl.StateMachineProcedure.Flow;
+import org.apache.iotdb.confignode.procedure.state.AddRegionPeerState;
 import org.apache.iotdb.confignode.procedure.store.ProcedureFactory;
+import org.apache.iotdb.rpc.TSStatusCode;
 
 import org.apache.tsfile.utils.PublicBAOS;
 import org.junit.Assert;
+import org.junit.Before;
 import org.junit.Test;
 
 import java.io.DataOutputStream;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.ArgumentMatchers.anyMap;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class AddRegionPeerProcedureTest {
+  private final TConsensusGroupId regionId =
+      new TConsensusGroupId(TConsensusGroupType.DataRegion, 10);
+  private final TDataNodeLocation coordinator = location(1);
+  private final TDataNodeLocation target = location(2);
+  private RegionMaintainHandler handler;
+  private ConfigNodeProcedureEnv env;
+
+  @Before
+  public void setUp() throws Exception {
+    ConfigManager configManager = mock(ConfigManager.class);
+    PartitionManager partitionManager = mock(PartitionManager.class);
+    handler = mock(RegionMaintainHandler.class);
+    env = mock(ConfigNodeProcedureEnv.class);
+    when(env.getConfigManager()).thenReturn(configManager);
+    when(env.getRegionMaintainHandler()).thenReturn(handler);
+    when(configManager.getPartitionManager()).thenReturn(partitionManager);
+    when(partitionManager.getAllReplicaSets())
+        .thenReturn(
+            Collections.singletonList(
+                new TRegionReplicaSet(
+                    regionId, new ArrayList<>(Collections.singletonList(coordinator)))));
+    when(handler.createNewRegionPeer(regionId, target))
+        .thenReturn(new TSStatus(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode()));
+    when(handler.removeRegionLocation(regionId, target))
+        .thenReturn(new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
+  }
+
   @Test
   public void serDeTest() throws Exception {
     AddRegionPeerProcedure procedure =
@@ -60,4 +110,76 @@ public void serDeTest() throws Exception {
       Assert.assertEquals(procedure, ProcedureFactory.getInstance().create(buffer));
     }
   }
+
+  @Test
+  public void rollbackRetriesUnavailablePeerAfterDelay() throws Exception {
+    when(handler.resetPeerList(eq(regionId), anyList(), anyMap()))
+        .thenReturn(
+            responses(TSStatusCode.CAN_NOT_CONNECT_DATANODE),
+            responses(TSStatusCode.SUCCESS_STATUS));
+
+    long start = System.nanoTime();
+    executeFailedCreate();
+    Assert.assertTrue(System.nanoTime() - start >= TimeUnit.SECONDS.toNanos(5));
+    verify(handler, times(2)).resetPeerList(eq(regionId), anyList(), anyMap());
+  }
+
+  @Test
+  public void rollbackStopsAfterFourFailedRpcBatches() throws Exception {
+    when(handler.resetPeerList(eq(regionId), anyList(), anyMap()))
+        .thenReturn(responses(TSStatusCode.CAN_NOT_CONNECT_DATANODE));
+
+    long start = System.nanoTime();
+    executeFailedCreate();
+    Assert.assertTrue(System.nanoTime() - start >= TimeUnit.SECONDS.toNanos(15));
+    verify(handler, times(4)).resetPeerList(eq(regionId), anyList(), anyMap());
+  }
+
+  @Test
+  public void rollbackDoesNotRetrySuccessfulRpcBatch() throws Exception {
+    when(handler.resetPeerList(eq(regionId), anyList(), anyMap()))
+        .thenReturn(responses(TSStatusCode.SUCCESS_STATUS));
+
+    executeFailedCreate();
+    verify(handler, times(1)).resetPeerList(eq(regionId), anyList(), anyMap());
+  }
+
+  @Test
+  public void interruptedRollbackWaitDoesNotRetry() throws Exception {
+    when(handler.resetPeerList(eq(regionId), anyList(), anyMap()))
+        .thenReturn(responses(TSStatusCode.CAN_NOT_CONNECT_DATANODE));
+
+    try {
+      Thread.currentThread().interrupt();
+      executeFailedCreate();
+      Assert.assertTrue(Thread.currentThread().isInterrupted());
+      verify(handler, times(1)).resetPeerList(eq(regionId), anyList(), anyMap());
+    } finally {
+      Thread.interrupted();
+    }
+  }
+
+  private void executeFailedCreate() throws Exception {
+    AddRegionPeerProcedure procedure = new AddRegionPeerProcedure(regionId, coordinator, target);
+    Assert.assertEquals(
+        Flow.NO_MORE_STATE,
+        procedure.executeFromState(env, AddRegionPeerState.CREATE_NEW_REGION_PEER));
+    Assert.assertFalse(procedure.isFailed());
+    verify(handler, times(1)).createNewRegionPeer(regionId, target);
+    verify(handler, times(1)).removeRegionLocation(regionId, target);
+  }
+
+  private Map<Integer, TSStatus> responses(TSStatusCode targetStatus) {
+    Map<Integer, TSStatus> responses = new HashMap<>();
+    responses.put(
+        coordinator.getDataNodeId(), new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
+    responses.put(target.getDataNodeId(), new TSStatus(targetStatus.getStatusCode()));
+    return responses;
+  }
+
+  private static TDataNodeLocation location(int id) {
+    return new TDataNodeLocation()
+        .setDataNodeId(id)
+        .setInternalEndPoint(new TEndPoint("127.0.0.1", 10730 + id));
+  }
 }