From 1acec07de9aab4d11d33363a2d5b7075df89de69 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:19:09 +0800 Subject: [PATCH] Fix delayed fragment memory release --- .../fragment/FragmentInstanceContext.java | 6 +-- .../FragmentInstanceStateMachine.java | 48 ++++++++++++++++++- .../plan/planner/LocalExecutionPlanner.java | 42 ++++++++++------ .../FragmentInstanceStateMachineTest.java | 31 ++++++++++++ ...alExecutionPlannerOperatorsMemoryTest.java | 34 +++++++++++++ 5 files changed, 142 insertions(+), 19 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index e29e2a2141aa1..7b7cb7ff5ca40 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -459,11 +459,11 @@ private QueryModificationLoader getQueryModificationLoader( modsTreeMatcher); } - // the state change listener is added here in a separate initialize() method + // the final-state listener is added here in a separate initialize() method // instead of the constructor to prevent leaking the "this" reference to - // another thread, which will cause unsafe publication of this instance. + // listener code before this instance has been fully constructed. private void initialize() { - stateMachine.addStateChangeListener(this::updateStatsIfDone); + stateMachine.addFinalStateChangeListener(this::updateStatsIfDone); } private void updateStatsIfDone(FragmentInstanceState newState) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceStateMachine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceStateMachine.java index 771d96693d48f..e0f637d90fa2b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceStateMachine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceStateMachine.java @@ -21,6 +21,7 @@ import org.apache.iotdb.calc.execution.StateMachine; import org.apache.iotdb.calc.execution.StateMachine.StateChangeListener; +import org.apache.iotdb.calc.i18n.CalcMessages; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.utils.SetThreadName; @@ -71,6 +72,10 @@ public class FragmentInstanceStateMachine { private final List sourceInstanceFailureListeners = new ArrayList<>(); + @GuardedBy("finalStateChangeListeners") + private final List> finalStateChangeListeners = + new ArrayList<>(); + public FragmentInstanceStateMachine(FragmentInstanceId fragmentInstanceId, Executor executor) { this.instanceId = requireNonNull( @@ -150,7 +155,15 @@ private void transitionToDoneState(FragmentInstanceState doneState) { DataNodeQueryMessages.EXCEPTION_DONESTATE_ARG_IS_NOT_A_DONE_STATE_8724C618, doneState); - instanceState.setIf(doneState, currentState -> !currentState.isDone()); + List> listeners; + synchronized (finalStateChangeListeners) { + if (!instanceState.setIf(doneState, currentState -> !currentState.isDone())) { + return; + } + listeners = ImmutableList.copyOf(finalStateChangeListeners); + finalStateChangeListeners.clear(); + } + listeners.forEach(listener -> fireFinalStateChangeListener(doneState, listener)); } /** @@ -164,6 +177,39 @@ public void addStateChangeListener( instanceState.addStateChangeListener(stateChangeListener); } + /** + * Adds a listener that is notified synchronously when this state machine first reaches a terminal + * state. If the state is already terminal, the listener is notified before this method returns. + * + *

This is intended only for short, non-blocking bookkeeping that must be completed before the + * state transition method returns. Other cleanup should use {@link #addStateChangeListener}. + */ + public void addFinalStateChangeListener( + StateChangeListener stateChangeListener) { + requireNonNull( + stateChangeListener, CalcMessages.EXCEPTION_STATECHANGELISTENER_IS_NULL_635AE7D2); + + FragmentInstanceState currentState; + synchronized (finalStateChangeListeners) { + currentState = instanceState.get(); + if (!currentState.isDone()) { + finalStateChangeListeners.add(stateChangeListener); + return; + } + } + fireFinalStateChangeListener(currentState, stateChangeListener); + } + + @SuppressWarnings("squid:S1181") + private void fireFinalStateChangeListener( + FragmentInstanceState state, StateChangeListener stateChangeListener) { + try { + stateChangeListener.stateChanged(state); + } catch (Throwable t) { + LOGGER.error(CalcMessages.ERROR_NOTIFYING_STATE_CHANGE_LISTENER_FOR, instanceId, t); + } + } + public void addSourceTaskFailureListener(FragmentInstanceFailureListener listener) { Map failures; synchronized (this) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java index a093c7832fdd0..d49753f90c021 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlanner.java @@ -221,25 +221,37 @@ private long checkMemory( } FragmentInstanceStateMachine stateMachine = instanceContext.getStateMachine(); if (reservedBytes > 0) { - stateMachine.addStateChangeListener( - newState -> { - if (newState.isDone()) { - try (SetThreadName fragmentInstanceName = - new SetThreadName(stateMachine.getFragmentInstanceId().getFullId())) { - OPERATORS_MEMORY_BLOCK.release(reservedBytes); - if (LOGGER.isDebugEnabled()) { - LOGGER.debug( - DataNodeQueryMessages.RELEASEMEMORY_RELEASE_ARG_CURRENT_REMAINING_MEMORY_ARG, - reservedBytes, - OPERATORS_MEMORY_BLOCK.getFreeMemoryInBytes()); - } - } - } - }); + registerOperatorsMemoryRelease(reservedBytes, stateMachine); } return reservedBytes; } + private void registerOperatorsMemoryRelease( + long reservedBytes, FragmentInstanceStateMachine stateMachine) { + // Regular state listeners share a bounded executor and can be delayed by driver cleanup. + // Return this admission reservation in the transition thread so consecutive fast failures + // cannot exhaust the operators memory pool. + stateMachine.addFinalStateChangeListener( + ignored -> { + try (SetThreadName fragmentInstanceName = + new SetThreadName(stateMachine.getFragmentInstanceId().getFullId())) { + OPERATORS_MEMORY_BLOCK.release(reservedBytes); + if (LOGGER.isDebugEnabled()) { + LOGGER.debug( + DataNodeQueryMessages.RELEASEMEMORY_RELEASE_ARG_CURRENT_REMAINING_MEMORY_ARG, + reservedBytes, + OPERATORS_MEMORY_BLOCK.getFreeMemoryInBytes()); + } + } + }); + } + + @TestOnly + void registerOperatorsMemoryReleaseForTest( + long reservedBytes, FragmentInstanceStateMachine stateMachine) { + registerOperatorsMemoryRelease(reservedBytes, stateMachine); + } + /** * Try to reserve bytes from the operators memory block. * diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceStateMachineTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceStateMachineTest.java index d0121e320c143..f4a726d30b6aa 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceStateMachineTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceStateMachineTest.java @@ -28,7 +28,10 @@ import org.mockito.Mockito; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import static org.apache.iotdb.db.queryengine.common.QueryId.MOCK_QUERY_ID; import static org.junit.Assert.assertEquals; @@ -117,4 +120,32 @@ public void testToAbort() { instanceNotificationExecutor.shutdown(); } } + + @Test + public void testFinalStateChangeListenerIsSynchronousAndCalledOnce() { + Executor stalledNotificationExecutor = command -> {}; + FragmentInstanceId instanceId = + new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 0), "0"); + FragmentInstanceStateMachine stateMachine = + new FragmentInstanceStateMachine(instanceId, stalledNotificationExecutor); + AtomicInteger invocationCount = new AtomicInteger(); + AtomicReference notifiedState = new AtomicReference<>(); + + stateMachine.addFinalStateChangeListener( + state -> { + invocationCount.incrementAndGet(); + notifiedState.set(state); + }); + stateMachine.failed(new RuntimeException("Unknown")); + + assertEquals(FragmentInstanceState.FAILED, notifiedState.get()); + assertEquals(1, invocationCount.get()); + + stateMachine.abort(); + assertEquals(1, invocationCount.get()); + + AtomicReference lateNotifiedState = new AtomicReference<>(); + stateMachine.addFinalStateChangeListener(lateNotifiedState::set); + assertEquals(FragmentInstanceState.FAILED, lateNotifiedState.get()); + } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java index 6d0cabb044313..d900b5263d654 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/LocalExecutionPlannerOperatorsMemoryTest.java @@ -19,13 +19,20 @@ package org.apache.iotdb.db.queryengine.plan.planner; +import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; import org.apache.iotdb.db.queryengine.common.QueryId; +import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine; import org.apache.iotdb.db.queryengine.plan.planner.memory.NotThreadSafeMemoryReservationManager; import org.junit.After; import org.junit.Assert; import org.junit.Test; +import java.util.concurrent.Executor; + +import static org.apache.iotdb.db.queryengine.common.QueryId.MOCK_QUERY_ID; + public class LocalExecutionPlannerOperatorsMemoryTest { private static final LocalExecutionPlanner PLANNER = LocalExecutionPlanner.getInstance(); @@ -77,6 +84,33 @@ public void testHighestPriorityFallbackWhenPoolInsufficient() { Assert.assertEquals(freeBefore, PLANNER.getFreeMemoryForOperators()); } + @Test + public void testEstimatedMemoryIsReleasedWithoutWaitingForAsyncStateListeners() { + long request = Math.min(1024L, PLANNER.getFreeMemoryForOperators()); + if (request <= 0) { + return; + } + long freeBefore = PLANNER.getFreeMemoryForOperators(); + long reserved = PLANNER.allocateOperatorsMemoryForTest(request, false); + Assert.assertEquals(request, reserved); + bytesHeldByTest = reserved; + + Executor stalledNotificationExecutor = command -> {}; + FragmentInstanceStateMachine stateMachine = + new FragmentInstanceStateMachine( + new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 0), "0"), + stalledNotificationExecutor); + PLANNER.registerOperatorsMemoryReleaseForTest(reserved, stateMachine); + + stateMachine.failed(new RuntimeException("Unknown")); + + long freeAfterFailure = PLANNER.getFreeMemoryForOperators(); + if (freeAfterFailure == freeBefore) { + bytesHeldByTest = 0L; + } + Assert.assertEquals(freeBefore, freeAfterFailure); + } + @Test public void testMemoryReservationManagerHighestPriorityAllocatesWhenPoolHasRoom() { long request = Math.min(1024L, PLANNER.getFreeMemoryForOperators());