Fix delayed operator memory release after fragment termination

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) {

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<FragmentInstanceFailureListener> sourceInstanceFailureListeners =
       new ArrayList<>();
 
+  @GuardedBy("finalStateChangeListeners")
+  private final List<StateChangeListener<FragmentInstanceState>> 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<StateChangeListener<FragmentInstanceState>> 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.
+   *
+   * <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, 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<FragmentInstanceState> 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<FragmentInstanceId, Throwable> failures;
     synchronized (this) {

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

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

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