From 85f4d0210b62359a6e9e58733f3308b06e511a76 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 12 Oct 2026 12:21:08 +0800 Subject: [PATCH] Retry failed procedure rollback without losing state (#18822) * Retry failed procedure rollback without losing state * Fix missing TSStatusCode import in procedure manager test (cherry picked from commit 2751ee133b46ae5d9a17b27de84eb86fa2ead239) --- .../confignode/manager/ProcedureManager.java | 5 +- .../iotdb/confignode/procedure/Procedure.java | 14 +++ .../procedure/ProcedureExecutor.java | 15 +++ .../procedure/impl/StateMachineProcedure.java | 6 +- .../manager/ProcedureManagerTest.java | 20 ++++ .../procedure/STMProcedureTest.java | 108 ++++++++++++++++++ 6 files changed, 166 insertions(+), 2 deletions(-) diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java index 49b79ddc07975..5dc0fbe8eeac0 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/ProcedureManager.java @@ -1697,7 +1697,10 @@ private TSStatus waitingProcedureFinished(Procedure procedure) { && System.currentTimeMillis() - startTimeForCurrentProcedure < PROCEDURE_WAIT_TIME_OUT) { sleepWithoutInterrupt(PROCEDURE_WAIT_RETRY_TIMEOUT); } - if (!procedure.isFinished()) { + // Conflict checks can fail before a procedure is assigned a persistent id. In that case the + // procedure remains at NO_PROC_ID but already carries the user-facing failure. + if (!procedure.isFinished() + && !(procedure.getProcId() == Procedure.NO_PROC_ID && procedure.isFailed())) { // The procedure is still executing status = RpcUtils.getStatus(TSStatusCode.INTERNAL_REQUEST_TIME_OUT, PROCEDURE_TIMEOUT_MESSAGE); diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/Procedure.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/Procedure.java index 862bc449c7319..57c41bd9adf4f 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/Procedure.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/Procedure.java @@ -265,6 +265,15 @@ protected boolean isYieldAfterExecution(Env env) { return false; } + /** + * Delay in milliseconds before retrying a failed rollback. A negative value keeps the default + * behavior of completing rollback after an exception. Procedures opting in must make rollback + * idempotent and retain the state needed by the next attempt. + */ + protected long getRollbackRetryTimeout() { + return -1; + } + // -------------------------Internal methods - called by the procedureExecutor------------------ final boolean tryAcquireExecution() { return executing.compareAndSet(false, true); @@ -679,6 +688,11 @@ protected synchronized void setFailure(final ProcedureException exception) { */ protected synchronized boolean setTimeoutFailure(Env env) { if (state == ProcedureState.WAITING_TIMEOUT) { + if (exception != null && getRollbackRetryTimeout() >= 0) { + // Resume compensation without replacing the original execution failure. + setState(ProcedureState.FAILED); + return true; + } long timeDiff = System.currentTimeMillis() - lastUpdate; setFailure( "ProcedureExecutor", diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/ProcedureExecutor.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/ProcedureExecutor.java index f79ff888c3825..8165754293e44 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/ProcedureExecutor.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/ProcedureExecutor.java @@ -650,15 +650,30 @@ private ProcedureLockState acquireLock(Procedure proc) { * @return procedure lock state */ private ProcedureLockState executeRollback(Procedure procedure) { + boolean rollbackFailed = false; try { procedure.doRollback(this.environment); } catch (IOException e) { + rollbackFailed = true; LOG.error("Roll back failed for {}", procedure, e); } catch (InterruptedException e) { + rollbackFailed = true; LOG.warn("Interrupted exception occured for {}", procedure, e); } catch (Throwable t) { + rollbackFailed = true; LOG.error("CODE-BUG: runtime exception for {}", procedure, t); } + if (rollbackFailed && procedure.getRollbackRetryTimeout() >= 0) { + procedure.setTimeout(procedure.getRollbackRetryTimeout()); + procedure.setState(ProcedureState.WAITING_TIMEOUT); + try { + store.update(procedure); + } catch (Exception e) { + LOG.warn(procedure.toString(), e); + } + timeoutExecutor.add(procedure); + return ProcedureLockState.LOCK_EVENT_WAIT; + } cleanupAfterRollback(procedure); return ProcedureLockState.LOCK_ACQUIRED; } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/StateMachineProcedure.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/StateMachineProcedure.java index 2f74bddfed853..c87108fdabf57 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/StateMachineProcedure.java +++ b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/procedure/impl/StateMachineProcedure.java @@ -234,11 +234,15 @@ protected void rollback(final Env env) states.removeLast(); } + boolean completed = false; try { updateTimestamp(); rollbackState(env, getCurrentState()); + completed = true; } finally { - states.removeLast(); + if ((completed || getRollbackRetryTimeout() < 0) && !states.isEmpty()) { + states.removeLast(); + } updateTimestamp(); } } diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/ProcedureManagerTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/ProcedureManagerTest.java index 09a8acf16f4f3..3c78d0a41f3de 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/ProcedureManagerTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/ProcedureManagerTest.java @@ -30,15 +30,18 @@ import org.apache.iotdb.confignode.procedure.ProcedureExecutor; import org.apache.iotdb.confignode.procedure.env.ConfigNodeProcedureEnv; import org.apache.iotdb.confignode.procedure.env.RemoveDataNodeHandler; +import org.apache.iotdb.confignode.procedure.exception.ProcedureException; import org.apache.iotdb.confignode.procedure.impl.node.RemoveDataNodesProcedure; import org.apache.iotdb.confignode.procedure.impl.region.RegionMigrateProcedure; import org.apache.iotdb.confignode.procedure.impl.region.RegionMigrationPlan; +import org.apache.iotdb.rpc.TSStatusCode; import org.junit.Assert; import org.junit.BeforeClass; import org.junit.Test; import java.io.IOException; +import java.lang.reflect.Method; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; @@ -49,6 +52,7 @@ import java.util.concurrent.ConcurrentHashMap; import static org.apache.iotdb.db.service.RegionMigrateService.isFailed; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; @@ -141,6 +145,22 @@ public void testCheckRemoveDataNodeWithAnotherRemoveProcedure() { Assert.assertTrue(isFailed(status)); } + @Test + public void testPreSubmissionFailureWithNoProcedureIdIsReportedImmediately() throws Exception { + final Procedure procedure = mock(Procedure.class); + when(procedure.getProcId()).thenReturn(Procedure.NO_PROC_ID); + when(procedure.isFinished()).thenReturn(false); + when(procedure.isFailed()).thenReturn(true); + when(procedure.getException()).thenReturn(new ProcedureException("conflict")); + + final Method waitingProcedureFinished = + ProcedureManager.class.getDeclaredMethod("waitingProcedureFinished", Procedure.class); + waitingProcedureFinished.setAccessible(true); + final TSStatus status = + (TSStatus) waitingProcedureFinished.invoke(PROCEDURE_MANAGER, procedure); + Assert.assertEquals(TSStatusCode.EXECUTE_STATEMENT_ERROR.getStatusCode(), status.getCode()); + } + @Test public void testCheckRemoveDataNodeWithConflictRegionMigrateProcedure() { RegionMigrateProcedure regionMigrateProcedure = diff --git a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/STMProcedureTest.java b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/STMProcedureTest.java index 1f62a5d1aaae3..f18e6882d5636 100644 --- a/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/STMProcedureTest.java +++ b/iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/procedure/STMProcedureTest.java @@ -21,6 +21,7 @@ import org.apache.iotdb.confignode.procedure.entity.SimpleSTMProcedure; import org.apache.iotdb.confignode.procedure.env.TestProcEnv; +import org.apache.iotdb.confignode.procedure.exception.ProcedureException; import org.apache.iotdb.confignode.procedure.impl.StateMachineProcedure; import org.apache.iotdb.confignode.procedure.state.ProcedureState; import org.apache.iotdb.confignode.procedure.util.ProcedureTestUtil; @@ -28,8 +29,17 @@ import org.junit.Assert; import org.junit.Test; +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.io.IOException; import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; import java.util.concurrent.ConcurrentLinkedDeque; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; public class STMProcedureTest extends TestProcedureBase { @@ -60,6 +70,104 @@ public void testRolledBackProcedure() { Assert.assertEquals(1 + success - rolledback, acc.get()); } + @Test + public void testFailedBeforeExecutionCanRollbackWithoutState() throws Exception { + final RetryingRollbackProcedure procedure = new RetryingRollbackProcedure(); + procedure.setFailure(procedure.originalFailure); + procedure.doRollback(env); + Assert.assertEquals(Arrays.asList(0), procedure.attemptedStates); + Assert.assertSame(procedure.originalFailure, procedure.getException()); + } + + @Test + public void testFailedRollbackRetainsStateAfterSerialization() throws Exception { + final RetryingRollbackProcedure procedure = new RetryingRollbackProcedure(); + procedure.setState(ProcedureState.RUNNABLE); + procedure.doExecute(env); + procedure.doExecute(env); + try { + procedure.doRollback(env); + Assert.fail("Compensation should fail while the Region is unavailable"); + } catch (IOException expected) { + Assert.assertEquals(Arrays.asList(1), procedure.attemptedStates); + } + procedure.setTimeout(1000); + procedure.setState(ProcedureState.WAITING_TIMEOUT); + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + procedure.serialize(new DataOutputStream(bytes)); + final RetryingRollbackProcedure restored = new RetryingRollbackProcedure(); + restored.deserialize(ByteBuffer.wrap(bytes.toByteArray())); + Assert.assertTrue(restored.isFailed()); + Assert.assertEquals( + procedure.originalFailure.getMessage(), restored.getException().getMessage()); + restored.attemptedStates.addAll(Arrays.asList(1, 1)); + restored.doRollback(env); + restored.doRollback(env); + Assert.assertEquals(Arrays.asList(1, 1, 1, 0), restored.attemptedStates); + } + + @Test + public void testFailedRollbackRetriesSameStateWithoutOverwritingFailure() throws Exception { + final RetryingRollbackProcedure procedure = new RetryingRollbackProcedure(); + final long procId = procExecutor.submitProcedure(procedure); + Assert.assertTrue(procedure.firstFailure.await(5, TimeUnit.SECONDS)); + Assert.assertTrue(procedure.completed.await(5, TimeUnit.SECONDS)); + ProcedureTestUtil.waitForProcedure(procExecutor, procId); + Assert.assertTrue(procedure.isFinished()); + Assert.assertEquals(Arrays.asList(1, 1, 1, 0), procedure.attemptedStates); + Assert.assertSame(procedure.originalFailure, procedure.getException()); + } + + private static class RetryingRollbackProcedure + extends StateMachineProcedure { + private final CountDownLatch firstFailure = new CountDownLatch(1); + private final CountDownLatch completed = new CountDownLatch(1); + private final List attemptedStates = new ArrayList<>(); + private final ProcedureException originalFailure = new ProcedureException("Execution failed"); + + @Override + protected Flow executeFromState(TestProcEnv env, Integer state) { + if (state == 0) { + setNextState(1); + return Flow.HAS_MORE_STATE; + } + setFailure(originalFailure); + return Flow.NO_MORE_STATE; + } + + @Override + protected void rollbackState(TestProcEnv env, Integer state) throws IOException { + attemptedStates.add(state); + if (state == 1 && attemptedStates.size() < 3) { + firstFailure.countDown(); + throw new IOException("Region temporarily unavailable"); + } + if (state == 0) { + completed.countDown(); + } + } + + @Override + protected long getRollbackRetryTimeout() { + return 50; + } + + @Override + protected Integer getState(int stateId) { + return stateId; + } + + @Override + protected int getStateId(Integer state) { + return state; + } + + @Override + protected Integer getInitialState() { + return 0; + } + } + @Test public void testEofStateReexecutionDoesNotCallExecuteFromState() throws Exception { EofReexecutionProcedure procedure = new EofReexecutionProcedure();