diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java index 9a26dba4e664a..a76889d8b63d7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachine.java +++ b/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) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachineTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachineTest.java index 7a77eada418ce..40a769cde057a 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataRegionStateMachineTest.java +++ b/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 accept(final PlanVisitor visitor, final C context) { + ++acceptCount; + return (R) new TSStatus(statusCode); + } + + @Override + public List 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 getOutputColumnNames() { + return Collections.emptyList(); + } + + @Override + protected void serializeAttributes(final ByteBuffer byteBuffer) {} + + @Override + protected void serializeAttributes(final DataOutputStream stream) throws IOException {} + } }