Fix delayed operator memory release after fragment termination

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 =

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 =

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java

@@ -514,11 +514,11 @@ protected PatternTreeMap<Modification, PatternTreeMapFactory.ModsSerializer> 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) {

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 @@
   private final List<FragmentInstanceFailureListener> sourceInstanceFailureListeners =
       new ArrayList<>();
 
+  @GuardedBy("finalStateChangeListeners")
+  private final List<StateChangeListener<FragmentInstanceState>> 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 @@
     requireNonNull(doneState, "doneState is null");
     checkArgument(doneState.isDone(), "doneState %s is not a done state", doneState);
 
-    instanceState.setIf(doneState, currentState -> !currentState.isDone());
+    List<StateChangeListener<FragmentInstanceState>> 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 @@
     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.
+   *
+   * <p>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<FragmentInstanceState> 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<FragmentInstanceState> 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<FragmentInstanceId, Throwable> failures;
     synchronized (this) {

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.
    *

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 @@
       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<FragmentInstanceState> 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<FragmentInstanceState> 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());
+  }
 }

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 @@
     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());