From 85cd9f19ebf2986e801ea6bcd99330bcca3a91cc Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 6 Nov 2026 15:54:26 +0800 Subject: [PATCH] Fix delayed operator memory release after fragment termination (#18387) * Fix delayed fragment memory release * Fix final listener delivery when notifications are rejected (cherry picked from commit 23e1740ee07e77b4ac2cae54def7bbaac98bdfb3) --- .../iotdb/db/i18n/DataNodeQueryMessages.java | 7 ++ .../iotdb/db/i18n/DataNodeQueryMessages.java | 7 ++ .../fragment/FragmentInstanceContext.java | 6 +- .../FragmentInstanceStateMachine.java | 77 ++++++++++++++++++- .../plan/planner/LocalExecutionPlanner.java | 47 +++++++---- .../FragmentInstanceStateMachineTest.java | 56 ++++++++++++++ ...alExecutionPlannerOperatorsMemoryTest.java | 34 ++++++++ 7 files changed, 213 insertions(+), 21 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 4d04dbaf1d2e8..1f8300c2807ed 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -23,6 +23,13 @@ public final class DataNodeQueryMessages { private DataNodeQueryMessages() {} + public static final String EXCEPTION_STATECHANGELISTENER_IS_NULL_B030E8A0 = + "stateChangeListener is null"; + public static final String LOG_ERROR_NOTIFYING_STATE_CHANGE_LISTENER_FOR_ARG_00D668A9 = + "Error notifying state change listener for {}"; + public static final String LOG_RELEASEMEMORY_RELEASE_ARG_CURRENT_REMAINING_MEMORY_ARG_4E7B461E = + "[ReleaseMemory] release: {}, current remaining memory: {}"; + public static final String EXCEPTION_FAILED_TO_CONVERT_TSFILE_ARG_TO_TABLETS_STATUS_ARG_E7C9F477 = "Failed to convert TsFile %s to tablets. Status: %s"; public static final String EXCEPTION_FAILED_TO_CONVERT_TSFILE_ARG_TO_TABLETS_ARG_96358FDD = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index e84e99433f5f3..32e74a0e5b94f 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -23,6 +23,13 @@ public final class DataNodeQueryMessages { private DataNodeQueryMessages() {} + public static final String EXCEPTION_STATECHANGELISTENER_IS_NULL_B030E8A0 = + "stateChangeListener 为空"; + public static final String LOG_ERROR_NOTIFYING_STATE_CHANGE_LISTENER_FOR_ARG_00D668A9 = + "通知 {} 的状态变更监听器时出错"; + public static final String LOG_RELEASEMEMORY_RELEASE_ARG_CURRENT_REMAINING_MEMORY_ARG_4E7B461E = + "[ReleaseMemory] 释放：{}，当前剩余内存：{}"; + public static final String EXCEPTION_FAILED_TO_CONVERT_TSFILE_ARG_TO_TABLETS_STATUS_ARG_E7C9F477 = "将 TsFile %s 转换为 tablet 失败。状态：%s"; public static final String EXCEPTION_FAILED_TO_CONVERT_TSFILE_ARG_TO_TABLETS_ARG_96358FDD = 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 102df82d38440..3898be8302da2 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 @@ -514,11 +514,11 @@ protected PatternTreeMap get return cachedResult == null ? atomicReference.get() : cachedResult; } - // 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 bef977488ec68..02a3e3c351613 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 @@ -19,6 +19,7 @@ package org.apache.iotdb.db.queryengine.execution.fragment; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.common.FragmentInstanceId; import org.apache.iotdb.db.queryengine.execution.StateMachine; import org.apache.iotdb.db.queryengine.execution.StateMachine.StateChangeListener; @@ -70,6 +71,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(fragmentInstanceId, "fragmentInstanceId is null"); this.executor = requireNonNull(executor, "executor is null"); @@ -140,7 +145,41 @@ private void transitionToDoneState(FragmentInstanceState doneState) { requireNonNull(doneState, "doneState is null"); checkArgument(doneState.isDone(), "doneState %s is not a done state", doneState); - instanceState.setIf(doneState, currentState -> !currentState.isDone()); + List> listeners; + RuntimeException stateChangeFailure = null; + Error stateChangeError = null; + synchronized (finalStateChangeListeners) { + boolean stateChanged; + try { + stateChanged = instanceState.setIf(doneState, currentState -> !currentState.isDone()); + } catch (RuntimeException e) { + // StateMachine updates its state before submitting asynchronous notifications. Preserve + // the final-state callback guarantee if notification submission is rejected afterwards. + if (!instanceState.get().isDone()) { + throw e; + } + stateChanged = true; + stateChangeFailure = e; + } catch (Error e) { + if (!instanceState.get().isDone()) { + throw e; + } + stateChanged = true; + stateChangeError = e; + } + if (!stateChanged) { + return; + } + listeners = ImmutableList.copyOf(finalStateChangeListeners); + finalStateChangeListeners.clear(); + } + listeners.forEach(listener -> fireFinalStateChangeListener(doneState, listener)); + if (stateChangeFailure != null) { + throw stateChangeFailure; + } + if (stateChangeError != null) { + throw stateChangeError; + } } /** @@ -154,6 +193,42 @@ 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, DataNodeQueryMessages.EXCEPTION_STATECHANGELISTENER_IS_NULL_B030E8A0); + + 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( + DataNodeQueryMessages.LOG_ERROR_NOTIFYING_STATE_CHANGE_LISTENER_FOR_ARG_00D668A9, + 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 04dc7099539da..5cff59d84315e 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 @@ -167,27 +167,40 @@ private long checkMemory( } FragmentInstanceStateMachine stateMachine = instanceContext.getStateMachine(); if (reservedBytes > 0) { - stateMachine.addStateChangeListener( - newState -> { - if (newState.isDone()) { - try (SetThreadName fragmentInstanceName = - new SetThreadName(stateMachine.getFragmentInstanceId().getFullId())) { - synchronized (this) { - this.freeMemoryForOperators += reservedBytes; - if (LOGGER.isDebugEnabled()) { - LOGGER.debug( - "[ReleaseMemory] release: {}, current remaining memory: {}", - reservedBytes, - freeMemoryForOperators); - } - } - } - } - }); + 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())) { + synchronized (this) { + this.freeMemoryForOperators += reservedBytes; + if (LOGGER.isDebugEnabled()) { + LOGGER.debug( + DataNodeQueryMessages + .LOG_RELEASEMEMORY_RELEASE_ARG_CURRENT_REMAINING_MEMORY_ARG_4E7B461E, + reservedBytes, + freeMemoryForOperators); + } + } + } + }); + } + + @TestOnly + void registerOperatorsMemoryReleaseForTest( + long reservedBytes, FragmentInstanceStateMachine stateMachine) { + registerOperatorsMemoryRelease(reservedBytes, stateMachine); + } + /** * Try to reserve bytes from the operators free-memory pool. * 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..75e0acbe08f8a 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,11 @@ import org.mockito.Mockito; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; +import java.util.concurrent.RejectedExecutionException; +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 +121,56 @@ 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()); + } + + @Test + public void testFinalStateChangeListenerIsCalledWhenNotificationIsRejected() { + Executor rejectingNotificationExecutor = + command -> { + throw new RejectedExecutionException(); + }; + FragmentInstanceId instanceId = + new FragmentInstanceId(new PlanFragmentId(MOCK_QUERY_ID, 0), "0"); + FragmentInstanceStateMachine stateMachine = + new FragmentInstanceStateMachine(instanceId, rejectingNotificationExecutor); + AtomicInteger invocationCount = new AtomicInteger(); + stateMachine.addFinalStateChangeListener(state -> invocationCount.incrementAndGet()); + + try { + stateMachine.failed(new RuntimeException("Unknown")); + fail("Expected notification rejection"); + } catch (RejectedExecutionException e) { + // Preserve the existing state-transition failure while guaranteeing final callbacks run. + } + + assertEquals(FragmentInstanceState.FAILED, stateMachine.getState()); + assertEquals(1, invocationCount.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 9768d1177d5c1..cc763763d51e8 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,14 +19,21 @@ 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.exception.MemoryNotEnoughException; +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(); @@ -78,6 +85,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());