Stop data region write retry promptly (#18730)

Author: Caideyipi

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java

@@ -265,11 +265,13 @@ protected TSStatus write(PlanNode planNode) {
                   .PIPE_LOG_WRITE_OPERATION_STILL_FAILED_AFTER_RETRY_TIMES_BECAUSE_15EEA702,
               MAX_WRITE_RETRY_TIMES,
               result.getCode());
+          break;
         }
         try {
-          Thread.sleep(WRITE_RETRY_WAIT_TIME_IN_MS);
+          waitBeforeNextWriteRetry();
         } catch (InterruptedException e) {
           Thread.currentThread().interrupt();
+          break;
         }
       } else {
         if (TSStatusCode.TABLE_NOT_EXISTS.getStatusCode() == result.getCode()
@@ -282,6 +284,10 @@ protected TSStatus write(PlanNode planNode) {
     return result;
   }
 
+  protected void waitBeforeNextWriteRetry() throws InterruptedException {
+    Thread.sleep(WRITE_RETRY_WAIT_TIME_IN_MS);
+  }
+
   @Override
   public DataSet read(IConsensusRequest request) {
     if (region == null) {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachineTest.java

@@ -141,13 +141,14 @@ public void testPipeGeneratedWriteProcessRejectSkipsStateMachineRetry() {
 
   @Test
   public void testNonPipeWriteProcessRejectCanStillRetry() {
-    final DataRegionStateMachine stateMachine = new DataRegionStateMachine(null);
+    final TestingDataRegionStateMachine stateMachine = new TestingDataRegionStateMachine(false);
     final RetryControlledPlanNode planNode = new RetryControlledPlanNode(false);
 
     final TSStatus status = stateMachine.write(planNode);
 
     Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode());
     Assert.assertEquals(2, planNode.getAcceptCount());
+    Assert.assertEquals(1, stateMachine.getWaitCount());
   }
 
   @Test
@@ -164,7 +165,7 @@ public void testMetadataLeaseFencedNoRetry() {
 
   @Test
   public void testMetadataLeaseFencedRetryRequiredRetriesAndFails() {
-    final DataRegionStateMachine stateMachine = new DataRegionStateMachine(null);
+    final TestingDataRegionStateMachine stateMachine = new TestingDataRegionStateMachine(false);
     final FixedStatusPlanNode planNode =
         new FixedStatusPlanNode(TSStatusCode.METADATA_LEASE_FENCED_RETRY_REQUIRED.getStatusCode());
 
@@ -173,11 +174,12 @@ public void testMetadataLeaseFencedRetryRequiredRetriesAndFails() {
     Assert.assertEquals(
         TSStatusCode.METADATA_LEASE_FENCED_RETRY_REQUIRED.getStatusCode(), status.getCode());
     Assert.assertEquals(5, planNode.getAcceptCount());
+    Assert.assertEquals(4, stateMachine.getWaitCount());
   }
 
   @Test
   public void testMetadataLeaseFencedRetryRequiredRetriesAndSucceeds() {
-    final DataRegionStateMachine stateMachine = new DataRegionStateMachine(null);
+    final TestingDataRegionStateMachine stateMachine = new TestingDataRegionStateMachine(false);
     final SequenceStatusPlanNode planNode =
         new SequenceStatusPlanNode(
             TSStatusCode.METADATA_LEASE_FENCED_RETRY_REQUIRED.getStatusCode(),
@@ -188,6 +190,51 @@ public void testMetadataLeaseFencedRetryRequiredRetriesAndSucceeds() {
 
     Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), status.getCode());
     Assert.assertEquals(3, planNode.getAcceptCount());
+    Assert.assertEquals(2, stateMachine.getWaitCount());
+  }
+
+  @Test
+  public void testInterruptedWriteRetryStopsImmediately() {
+    Thread.interrupted();
+    try {
+      final TestingDataRegionStateMachine stateMachine = new TestingDataRegionStateMachine(true);
+      final FixedStatusPlanNode planNode =
+          new FixedStatusPlanNode(
+              TSStatusCode.METADATA_LEASE_FENCED_RETRY_REQUIRED.getStatusCode());
+
+      final TSStatus status = stateMachine.write(planNode);
+
+      Assert.assertEquals(
+          TSStatusCode.METADATA_LEASE_FENCED_RETRY_REQUIRED.getStatusCode(), status.getCode());
+      Assert.assertEquals(1, planNode.getAcceptCount());
+      Assert.assertEquals(1, stateMachine.getWaitCount());
+      Assert.assertTrue(Thread.currentThread().isInterrupted());
+    } finally {
+      Thread.interrupted();
+    }
+  }
+
+  private static class TestingDataRegionStateMachine extends DataRegionStateMachine {
+
+    private final boolean interruptOnWait;
+    private int waitCount;
+
+    private TestingDataRegionStateMachine(final boolean interruptOnWait) {
+      super(null);
+      this.interruptOnWait = interruptOnWait;
+    }
+
+    private int getWaitCount() {
+      return waitCount;
+    }
+
+    @Override
+    protected void waitBeforeNextWriteRetry() throws InterruptedException {
+      ++waitCount;
+      if (interruptOnWait) {
+        throw new InterruptedException();
+      }
+    }
   }
 
   private static class FixedStatusPlanNode extends PlanNode {