[To dev/1.3] Stop data region write retry promptly (#18730)

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

@@ -224,11 +224,13 @@ protected TSStatus write(PlanNode planNode) {
               "write operation still failed after {} retry times, because {}.",
               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 {
         break;
@@ -237,6 +239,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 (request instanceof GetConsensusReqReaderPlan) {

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

@@ -19,18 +19,26 @@
 
 package org.apache.iotdb.db.consensus.statemachine.dataregion;
 
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
 import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.path.PartialPath;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeId;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode;
 import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowsNode;
+import org.apache.iotdb.rpc.TSStatusCode;
 
 import org.apache.tsfile.enums.TSDataType;
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 public class DataRegionStateMachineTest {
@@ -119,4 +127,110 @@ public void testMergeInsertRowsNodes() throws IllegalPathException {
     InsertNode mergedNode = list.get(0).mergeInsertNode(list);
     Assert.assertTrue(mergedNode instanceof InsertRowsNode);
   }
+
+  @Test
+  public void testWriteRetryDoesNotWaitAfterLastAttempt() {
+    final TestingDataRegionStateMachine stateMachine = new TestingDataRegionStateMachine(false);
+    final FixedStatusPlanNode planNode =
+        new FixedStatusPlanNode(TSStatusCode.WRITE_PROCESS_REJECT.getStatusCode());
+
+    final TSStatus status = stateMachine.write(planNode);
+
+    Assert.assertEquals(TSStatusCode.WRITE_PROCESS_REJECT.getStatusCode(), status.getCode());
+    Assert.assertEquals(5, planNode.getAcceptCount());
+    Assert.assertEquals(4, stateMachine.getWaitCount());
+  }
+
+  @Test
+  public void testInterruptedWriteRetryStopsImmediately() {
+    Thread.interrupted();
+    try {
+      final TestingDataRegionStateMachine stateMachine = new TestingDataRegionStateMachine(true);
+      final FixedStatusPlanNode planNode =
+          new FixedStatusPlanNode(TSStatusCode.WRITE_PROCESS_REJECT.getStatusCode());
+
+      final TSStatus status = stateMachine.write(planNode);
+
+      Assert.assertEquals(TSStatusCode.WRITE_PROCESS_REJECT.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 {
+
+    private final int statusCode;
+    private int acceptCount;
+
+    private FixedStatusPlanNode(final int statusCode) {
+      super(new PlanNodeId("test"));
+      this.statusCode = statusCode;
+    }
+
+    private int getAcceptCount() {
+      return acceptCount;
+    }
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public <R, C> R accept(final PlanVisitor<R, C> visitor, final C context) {
+      ++acceptCount;
+      return (R) new TSStatus(statusCode);
+    }
+
+    @Override
+    public List<PlanNode> getChildren() {
+      return Collections.emptyList();
+    }
+
+    @Override
+    public void addChild(final PlanNode child) {}
+
+    @Override
+    public PlanNode clone() {
+      return new FixedStatusPlanNode(statusCode);
+    }
+
+    @Override
+    public int allowedChildCount() {
+      return NO_CHILD_ALLOWED;
+    }
+
+    @Override
+    public List<String> getOutputColumnNames() {
+      return Collections.emptyList();
+    }
+
+    @Override
+    protected void serializeAttributes(final ByteBuffer byteBuffer) {}
+
+    @Override
+    protected void serializeAttributes(final DataOutputStream stream) throws IOException {}
+  }
 }