Fix scan termination for finished fragment instances (#18701)

Author: JackieTien97

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/driver/Driver.java

@@ -26,6 +26,7 @@
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
 import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink;
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceFinishedException;
 import org.apache.iotdb.db.queryengine.execution.operator.OperatorContext;
 import org.apache.iotdb.db.queryengine.execution.schedule.task.DriverTaskId;
 import org.apache.iotdb.db.queryengine.metric.QueryMetricsManager;
@@ -248,6 +249,11 @@ private ListenableFuture<?> processInternal() {
         }
       }
       return NOT_BLOCKED;
+    } catch (FragmentInstanceFinishedException e) {
+      // The fragment may finish before its notification thread closes this driver. Release the
+      // driver through the normal cleanup path without aborting other fragments of the query.
+      state.compareAndSet(State.ALIVE, State.NEED_DESTRUCTION);
+      return NOT_BLOCKED;
     } catch (Throwable t) {
       Throwable actualCause = t;
       if (actualCause.getCause() instanceof IoTDBRuntimeException) {

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

@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.execution.fragment;
+
+import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
+import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
+
+/**
+ * Internal control signal for stopping a scan after its fragment has finished. This is unchecked so
+ * scan operators do not wrap it as an I/O failure before the driver can handle normal termination.
+ */
+public class FragmentInstanceFinishedException extends RuntimeException {
+
+  public FragmentInstanceFinishedException(FragmentInstanceId fragmentInstanceId) {
+    super(
+        String.format(
+            DataNodeQueryMessages.EXCEPTION_FRAGMENT_INSTANCE_ARG_IS_ALREADY_ARG_B44984B4,
+            fragmentInstanceId,
+            FragmentInstanceState.FINISHED));
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtil.java

@@ -25,6 +25,7 @@
 import org.apache.iotdb.db.exception.CorruptedTsFileException;
 import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
 import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceFinishedException;
 import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceState;
 import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext;
 import org.apache.iotdb.db.queryengine.metric.SeriesScanCostMetricSet;
@@ -1587,6 +1588,9 @@ private void checkFragmentInstanceState() throws IOException {
       return;
     }
     FragmentInstanceState state = context.getStateMachine().getState();
+    if (state == FragmentInstanceState.FINISHED) {
+      throw new FragmentInstanceFinishedException(context.getId());
+    }
     if (state.isDone()) {
       // A scan over many overlapping files may stay in one operator call long after cancellation.
       // Exit on the driver thread so it can release its lock and finish resource cleanup.

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtilCancellationTest.java

@@ -22,6 +22,7 @@
 import org.apache.iotdb.calc.execution.operator.Operator;
 import org.apache.iotdb.calc.plan.planner.memory.MemoryReservationManager;
 import org.apache.iotdb.commons.exception.QueryTimeoutException;
+import org.apache.iotdb.commons.exception.SemanticException;
 import org.apache.iotdb.commons.path.AlignedFullPath;
 import org.apache.iotdb.commons.path.NonAlignedFullPath;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
@@ -34,6 +35,7 @@
 import org.apache.iotdb.db.queryengine.execution.exchange.sink.ISink;
 import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceContext;
 import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceExecution;
+import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceFinishedException;
 import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceState;
 import org.apache.iotdb.db.queryengine.execution.fragment.FragmentInstanceStateMachine;
 import org.apache.iotdb.db.queryengine.execution.schedule.IDriverScheduler;
@@ -63,6 +65,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.function.IntConsumer;
 
@@ -71,9 +74,11 @@
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyList;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -139,8 +144,12 @@ public void testTerminalStateBeforeScanDoesNotLoadMetadata() {
           throw new AssertionError(state);
       }
 
-      IOException exception = assertThrows(IOException.class, scanner::hasNextFile);
-      assertSame(context.getFailureCause().orElse(null), exception.getCause());
+      if (state == FragmentInstanceState.FINISHED) {
+        assertThrows(FragmentInstanceFinishedException.class, scanner::hasNextFile);
+      } else {
+        IOException exception = assertThrows(IOException.class, scanner::hasNextFile);
+        assertSame(context.getFailureCause().orElse(null), exception.getCause());
+      }
       assertEquals(0, loadedFiles);
       assertEquals(state, context.getStateMachine().getState());
     }
@@ -237,12 +246,22 @@ public void testCompactionContextWithoutStateMachine() throws IOException {
 
   @Test(timeout = 15000)
   public void testTimeoutUnblocksDriverResourceCleanup() throws Exception {
+    assertFailureUnblocksDriverResourceCleanup(new QueryTimeoutException());
+  }
+
+  @Test(timeout = 15000)
+  public void testSemanticFailureUnblocksDriverResourceCleanup() throws Exception {
+    assertFailureUnblocksDriverResourceCleanup(
+        new SemanticException("Scalar sub-query has returned multiple rows."));
+  }
+
+  private void assertFailureUnblocksDriverResourceCleanup(RuntimeException failure)
+      throws Exception {
     ExecutorService notifications = Executors.newSingleThreadExecutor();
     FragmentInstanceContext context = newContext(notifications);
     context.initializeNumOfDrivers(1);
     try {
       CountDownLatch closeRequested = new CountDownLatch(1);
-      QueryTimeoutException timeout = new QueryTimeoutException();
       SeriesScanUtil scanner =
           newScanner(
               context,
@@ -251,7 +270,7 @@ public void testTimeoutUnblocksDriverResourceCleanup() throws Exception {
               4,
               count -> {
                 if (count == 2) {
-                  context.failed(timeout);
+                  context.failed(failure);
                   try {
                     // The notification thread has requested close while this thread owns the
                     // driver lock, just as in a query that is still loading metadata.
@@ -296,10 +315,9 @@ public void close() {
           exchangeManager);
 
       assertSame(
-          timeout,
+          failure,
           assertThrows(
-              QueryTimeoutException.class,
-              () -> driver.processFor(new Duration(1, TimeUnit.SECONDS))));
+              RuntimeException.class, () -> driver.processFor(new Duration(1, TimeUnit.SECONDS))));
       // This can complete only after the preceding cleanup callback gets past allDriversClosed.
       notifications.submit(() -> {}).get(5, TimeUnit.SECONDS);
       assertEquals(2, loadedFiles);
@@ -316,6 +334,146 @@ public void close() {
     }
   }
 
+  @Test(timeout = 60000)
+  public void testFinishedScanUnblocksDriverResourceCleanup() throws Exception {
+    for (boolean sequence : new boolean[] {true, false}) {
+      for (boolean waitForClose : new boolean[] {false, true}) {
+        assertFinishedScanUnblocksDriverResourceCleanup(sequence, waitForClose);
+      }
+    }
+  }
+
+  private void assertFinishedScanUnblocksDriverResourceCleanup(
+      boolean sequence, boolean waitForClose) throws Exception {
+    ExecutorService notifications = Executors.newSingleThreadExecutor();
+    ExecutorService worker = Executors.newSingleThreadExecutor();
+    CountDownLatch allowNotifications = new CountDownLatch(waitForClose ? 0 : 1);
+    CountDownLatch scanEntered = new CountDownLatch(1);
+    CountDownLatch resumeScan = new CountDownLatch(1);
+    CountDownLatch closeRequested = new CountDownLatch(1);
+    // Cover both orderings: the FI is FINISHED before driver.close(), and close is already pending.
+    notifications.submit(() -> await(allowNotifications));
+    FragmentInstanceContext context = newContext(notifications);
+    context.initializeNumOfDrivers(1);
+    Operator operator = mock(Operator.class);
+    doReturn(NOT_BLOCKED).when(operator).isBlocked();
+    when(operator.hasNextWithTimer()).thenReturn(true);
+    SeriesScanUtil scanner =
+        newScanner(
+            context,
+            Ordering.ASC,
+            sequence,
+            4,
+            count -> {
+              if (count == 1) {
+                scanEntered.countDown();
+                await(resumeScan);
+              }
+            });
+    when(operator.nextWithTimer())
+        .thenAnswer(
+            invocation -> {
+              scanner.hasNextFile();
+              return null;
+            });
+    ISink sink = mock(ISink.class);
+    doReturn(NOT_BLOCKED).when(sink).isFull();
+    DataDriverContext driverContext = new DataDriverContext(context, 0);
+    driverContext.setSink(sink);
+    DataDriver driver =
+        new DataDriver(operator, driverContext, 0) {
+          @Override
+          public void close() {
+            super.close();
+            closeRequested.countDown();
+          }
+        };
+    MPPDataExchangeManager exchangeManager = mock(MPPDataExchangeManager.class);
+    IDriverScheduler scheduler = mock(IDriverScheduler.class);
+    FragmentInstanceExecution.createFragmentInstanceExecution(
+        scheduler,
+        context.getId(),
+        context,
+        Collections.singletonList(driver),
+        sink,
+        context.getStateMachine(),
+        1000,
+        false,
+        exchangeManager);
+    try {
+      Future<?> execution =
+          worker.submit(() -> driver.processFor(new Duration(1, TimeUnit.SECONDS)));
+      await(scanEntered);
+      context.finished();
+      if (waitForClose) {
+        await(closeRequested);
+      }
+      resumeScan.countDown();
+      execution.get(5, TimeUnit.SECONDS);
+      assertTrue(driver.isFinished());
+      assertEquals(1, loadedFiles);
+      assertEquals(FragmentInstanceState.FINISHED, context.getStateMachine().getState());
+      assertTrue(context.getStateMachine().getFailureCauses().isEmpty());
+      allowNotifications.countDown();
+      notifications.submit(() -> {}).get(5, TimeUnit.SECONDS);
+      verify(operator).close();
+      verify(context.getMemoryReservationContext()).releaseAllReservedMemory();
+      verify(exchangeManager)
+          .deRegisterFragmentInstanceFromMemoryPool(
+              context.getId().getQueryId().getId(), context.getId().getFragmentInstanceId(), true);
+      verify(scheduler, never()).abortFragmentInstance(any(), any());
+    } finally {
+      resumeScan.countDown();
+      allowNotifications.countDown();
+      worker.shutdownNow();
+      assertTrue(worker.awaitTermination(5, TimeUnit.SECONDS));
+      driver.close();
+      context.decrementNumOfUnClosedDriver();
+      // Let FI cleanup finish after opening its latch, without interrupting its driver-close wait.
+      notifications.shutdown();
+      assertTrue(notifications.awaitTermination(5, TimeUnit.SECONDS));
+    }
+  }
+
+  @Test
+  public void testFinishedFragmentDoesNotSuppressIOException() throws Exception {
+    FragmentInstanceContext context = newContext(Runnable::run);
+    context.initializeNumOfDrivers(1);
+    IOException failure = new IOException("Failed to read file metadata");
+    Operator operator = mock(Operator.class);
+    doReturn(NOT_BLOCKED).when(operator).isBlocked();
+    when(operator.hasNextWithTimer()).thenReturn(true);
+    when(operator.nextWithTimer())
+        .thenAnswer(
+            invocation -> {
+              context.finished();
+              throw failure;
+            });
+    ISink sink = mock(ISink.class);
+    doReturn(NOT_BLOCKED).when(sink).isFull();
+    DataDriverContext driverContext = new DataDriverContext(context, 0);
+    driverContext.setSink(sink);
+    DataDriver driver = new DataDriver(operator, driverContext, 0);
+    try {
+      RuntimeException exception =
+          assertThrows(
+              RuntimeException.class, () -> driver.processFor(new Duration(1, TimeUnit.SECONDS)));
+      assertSame(failure, exception.getCause());
+      assertSame(failure, context.getFailureCause().get());
+    } finally {
+      driver.close();
+    }
+  }
+
+  private static void await(CountDownLatch latch) {
+    try {
+      assertTrue(latch.await(5, TimeUnit.SECONDS));
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new AssertionError(e);
+    }
+  }
+
   private FragmentInstanceContext newContext(Executor executor) {
     FragmentInstanceId id =
         new FragmentInstanceId(new PlanFragmentId(new QueryId("scan_cancellation"), 0), "0");