Fix misleading shutdown interruption warnings (#18939)

Author: JackieTien97

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java

@@ -1263,6 +1263,7 @@ public final class DataNodeQueryMessages {
       "Interrupted when dispatching read async";
   public static final String INTERRUPTED_WHEN_DISPATCHING_WRITE_ASYNC =
       "Interrupted when dispatching write async";
+  public static final String MESSAGE_INTERRUPTED_ERRORS_ARG_876F1B4B = "Interrupted errors: %s";
   public static final String DESERIALIZE_CONSENSUSGROUPID_FAILED =
       "Deserialize ConsensusGroupId failed. ";
   public static final String CAN_T_CONNECT_TO_NODE =
@@ -3476,6 +3477,7 @@ private DataNodeQueryMessages() {}
       "already being cancelled";
   public static final String DRIVER_TASK_ABORTED_BY_INTERNAL_ERROR_SCHEDULED =
       "internal error scheduled";
+  public static final String EXCEPTION_SCHEDULER_SHUTDOWN_323938F5 = "scheduler shutdown";
   public static final String DRIVER_TASK_ABORTED_BY_MEMORY_NOT_ENOUGH =
       "Memory is not enough to execute the query task.";
   public static final String ROOT_FRAGMENT_INSTANCE_PLACEMENT_ERROR_FMT =

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java

@@ -1243,6 +1243,7 @@ public final class DataNodeQueryMessages {
       "异步分发读取操作时被中断";
   public static final String INTERRUPTED_WHEN_DISPATCHING_WRITE_ASYNC =
       "异步分发写入操作时被中断";
+  public static final String MESSAGE_INTERRUPTED_ERRORS_ARG_876F1B4B = "中断错误：%s";
   public static final String DESERIALIZE_CONSENSUSGROUPID_FAILED =
       "反序列化 ConsensusGroupId 失败。";
   public static final String CAN_T_CONNECT_TO_NODE =
@@ -4127,6 +4128,7 @@ private DataNodeQueryMessages() {}
       "已经在取消中";
   public static final String DRIVER_TASK_ABORTED_BY_INTERNAL_ERROR_SCHEDULED =
       "内部错误调度";
+  public static final String EXCEPTION_SCHEDULER_SHUTDOWN_323938F5 = "调度器关闭";
   public static final String DRIVER_TASK_ABORTED_BY_MEMORY_NOT_ENOUGH =
       "内存不足，无法执行查询任务。";
   public static final String ROOT_FRAGMENT_INSTANCE_PLACEMENT_ERROR_FMT =

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

@@ -54,7 +54,7 @@ public class QueryStateMachine {
   private final StateMachine<QueryState> queryState;
 
   private final AtomicReference<Throwable> failureException = new AtomicReference<>();
-  private final AtomicReference<TSStatus> failureStatus = new AtomicReference<>();
+  private final AtomicReference<FailureStatus> failureStatus = new AtomicReference<>();
 
   public QueryStateMachine(QueryId queryId, ExecutorService executor) {
     this.queryState =
@@ -90,7 +90,7 @@ public void transitionToDispatching() {
   }
 
   public void transitionToPendingRetry(TSStatus failureStatus) {
-    this.failureStatus.compareAndSet(null, failureStatus);
+    recordFailureStatus(failureStatus, false);
     queryState.setIf(PENDING_RETRY, currentState -> currentState == DISPATCHING);
   }
 
@@ -110,7 +110,7 @@ public void transitionToCanceled() {
   }
 
   public boolean transitionToCanceled(Throwable throwable, TSStatus failureStatus) {
-    this.failureStatus.compareAndSet(null, failureStatus);
+    recordFailureStatus(failureStatus, false);
     this.failureException.compareAndSet(null, throwable);
     return transitionToDoneState(CANCELED);
   }
@@ -129,10 +129,34 @@ public void transitionToFailed(Throwable throwable) {
   }
 
   public void transitionToFailed(TSStatus failureStatus) {
-    this.failureStatus.compareAndSet(null, failureStatus);
+    transitionToFailed(failureStatus, false);
+  }
+
+  public void transitionToFailed(TSStatus failureStatus, boolean failureCausedByInterruption) {
+    recordFailureStatus(failureStatus, failureCausedByInterruption);
     transitionToDoneState(FAILED);
   }
 
+  private void recordFailureStatus(TSStatus status, boolean failureCausedByInterruption) {
+    if (status != null) {
+      // Publish the status and its origin together. A later interruption must not reclassify
+      // the first recorded failure.
+      failureStatus.compareAndSet(null, new FailureStatus(status, failureCausedByInterruption));
+    }
+  }
+
+  /** Only the exact recorded status carries the interruption provenance stored with it. */
+  public boolean isFailureCausedByInterruption(TSStatus status) {
+    FailureStatus failure = failureStatus.get();
+    if (failure == null || failure.status != status || !failure.failureCausedByInterruption) {
+      return false;
+    }
+    // Statuses and exceptions can be recorded independently. Keep an unrelated exception
+    // visible even if a subsequent dispatch interruption supplied the status.
+    Throwable throwable = failureException.get();
+    return throwable == null || getRootCause(throwable) instanceof InterruptedException;
+  }
+
   private boolean transitionToDoneState(QueryState doneState) {
     requireNonNull(doneState, DataNodeQueryMessages.EXCEPTION_DONESTATE_IS_NULL_D88F77E5);
     checkArgument(
@@ -161,9 +185,9 @@ public Throwable getFailureException() {
   }
 
   public TSStatus getFailureStatus() {
-    TSStatus status = failureStatus.get();
+    FailureStatus status = failureStatus.get();
     if (status != null) {
-      return status;
+      return status.status;
     } else {
       Throwable throwable = failureException.get();
       if (throwable != null) {
@@ -174,7 +198,18 @@ public TSStatus getFailureStatus() {
           return RpcUtils.getStatus(((IoTDBException) t).getErrorCode(), t.getMessage());
         }
       }
-      return failureStatus.get();
+      status = failureStatus.get();
+      return status == null ? null : status.status;
+    }
+  }
+
+  private static class FailureStatus {
+    private final TSStatus status;
+    private final boolean failureCausedByInterruption;
+
+    private FailureStatus(TSStatus status, boolean failureCausedByInterruption) {
+      this.status = status;
+      this.failureCausedByInterruption = failureCausedByInterruption;
     }
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/AbstractDriverThread.java

@@ -67,8 +67,11 @@ public void run() {
         try {
           next = queue.poll();
         } catch (InterruptedException e) {
-          logger.warn(
-              DataNodeQueryMessages.EXECUTOR_FAILED_TO_POLL_DRIVER_TASK_FROM_QUEUE, this.getName());
+          if (!closed) {
+            logger.warn(
+                DataNodeQueryMessages.EXECUTOR_FAILED_TO_POLL_DRIVER_TASK_FROM_QUEUE,
+                this.getName());
+          }
           Thread.currentThread().interrupt();
           break;
         }
@@ -86,7 +89,14 @@ public void run() {
           try (SetThreadName driverTaskName =
               new SetThreadName(next.getDriver().getDriverTaskId().getFullId())) {
             Throwable rootCause = ErrorHandlingCommonUtils.getRootCause(e);
-            if (rootCause instanceof CorruptedTsFileException) {
+            if (closed && rootCause instanceof InterruptedException) {
+              Thread.currentThread().interrupt();
+              // Still abort the task below so that shutdown releases its query resources.
+              next.setAbortCause(
+                  new DriverTaskAbortedException(
+                      next.getDriverTaskId().getFullId(),
+                      DriverTaskAbortedException.BY_SCHEDULER_SHUTDOWN));
+            } else if (rootCause instanceof CorruptedTsFileException) {
               // CorruptedTsFileException no longer chains the original IOException as its
               // cause (it uses addSuppressed instead), so getRootCause returns the exception
               // itself and we can match it here.

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/schedule/DriverTaskAbortedException.java

@@ -34,6 +34,8 @@ public class DriverTaskAbortedException extends Exception {
       DataNodeQueryMessages.DRIVER_TASK_ABORTED_BY_ALREADY_BEING_CANCELLED;
   public static final String BY_INTERNAL_ERROR_SCHEDULED =
       DataNodeQueryMessages.DRIVER_TASK_ABORTED_BY_INTERNAL_ERROR_SCHEDULED;
+  public static final String BY_SCHEDULER_SHUTDOWN =
+      DataNodeQueryMessages.EXCEPTION_SCHEDULER_SHUTDOWN_323938F5;
 
   public static final String BY_MEMORY_NOT_ENOUGH =
       DataNodeQueryMessages.DRIVER_TASK_ABORTED_BY_MEMORY_NOT_ENOUGH;

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/ExecutionResult.java

@@ -26,8 +26,16 @@ public class ExecutionResult {
   public final QueryId queryId;
   public final TSStatus status;
 
+  /** Whether this failure originated from an interrupted wait rather than another query error. */
+  public final boolean failureCausedByInterruption;
+
   public ExecutionResult(QueryId queryId, TSStatus status) {
+    this(queryId, status, false);
+  }
+
+  ExecutionResult(QueryId queryId, TSStatus status, boolean failureCausedByInterruption) {
     this.queryId = queryId;
     this.status = status;
+    this.failureCausedByInterruption = failureCausedByInterruption;
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/execution/QueryExecution.java

@@ -644,19 +644,29 @@ public ExecutionResult getStatus() {
     } catch (InterruptedException e) {
       // TODO: (xingtanzjr) use more accurate error handling
       Thread.currentThread().interrupt();
+      // Dispatch may already have converted an interruption to a failure status. Preserve its
+      // recorded origin without reclassifying an unrelated failure.
+      TSStatus failureStatus = stateMachine.getFailureStatus();
+      if (failureStatus != null) {
+        return new ExecutionResult(
+            context.getQueryId(),
+            failureStatus,
+            stateMachine.isFailureCausedByInterruption(failureStatus));
+      }
       return new ExecutionResult(
           context.getQueryId(),
-          stateMachine.getFailureStatus() == null
-              ? RpcUtils.getStatus(
-                  TSStatusCode.INTERNAL_SERVER_ERROR, stateMachine.getFailureMessage())
-              : stateMachine.getFailureStatus());
+          RpcUtils.getStatus(TSStatusCode.INTERNAL_SERVER_ERROR, stateMachine.getFailureMessage()),
+          // A terminal query may have failed with an exception that has no dedicated status code.
+          !stateMachine.getState().isDone());
     } catch (ExecutionException e) {
+      TSStatus failureStatus = stateMachine.getFailureStatus();
       return new ExecutionResult(
           context.getQueryId(),
-          stateMachine.getFailureStatus() == null
+          failureStatus == null
               ? RpcUtils.getStatus(
                   TSStatusCode.INTERNAL_SERVER_ERROR, stateMachine.getFailureMessage())
-              : stateMachine.getFailureStatus());
+              : failureStatus,
+          stateMachine.isFailureCausedByInterruption(failureStatus));
     }
   }
 
@@ -727,7 +737,8 @@ private ExecutionResult getExecutionResult(QueryState state) {
         || IoTDBConstant.ClientVersion.V_1_0.equals(context.getSession().getVersion())) {
       planner.setRedirectInfo(analysis, CONFIG.getAddressAndPort(), tsstatus);
     }
-    return new ExecutionResult(context.getQueryId(), tsstatus);
+    return new ExecutionResult(
+        context.getQueryId(), tsstatus, stateMachine.isFailureCausedByInterruption(tsstatus));
   }
 
   public DistributedQueryPlan getDistributedPlan() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/ClusterScheduler.java

@@ -125,7 +125,8 @@ public void start() {
         if (needRetry(result.getFailureStatus())) {
           stateMachine.transitionToPendingRetry(result.getFailureStatus());
         } else {
-          stateMachine.transitionToFailed(result.getFailureStatus());
+          stateMachine.transitionToFailed(
+              result.getFailureStatus(), result.isFailureCausedByInterruption());
         }
         return;
       }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragInstanceDispatchResult.java

@@ -26,13 +26,21 @@ public class FragInstanceDispatchResult {
 
   private TSStatus failureStatus;
 
+  private final boolean failureCausedByInterruption;
+
   public FragInstanceDispatchResult(boolean successful) {
     this.successful = successful;
+    this.failureCausedByInterruption = false;
   }
 
   public FragInstanceDispatchResult(TSStatus failureStatus) {
+    this(failureStatus, false);
+  }
+
+  public FragInstanceDispatchResult(TSStatus failureStatus, boolean failureCausedByInterruption) {
     this.successful = false;
     this.failureStatus = failureStatus;
+    this.failureCausedByInterruption = failureCausedByInterruption;
   }
 
   public boolean isSuccessful() {
@@ -42,4 +50,8 @@ public boolean isSuccessful() {
   public TSStatus getFailureStatus() {
     return failureStatus;
   }
+
+  public boolean isFailureCausedByInterruption() {
+    return failureCausedByInterruption;
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/FragmentInstanceDispatcherImpl.java

@@ -330,7 +330,11 @@ private Future<FragInstanceDispatchResult> dispatchWrite(List<FragmentInstance>
       return immediateFuture(
           new FragInstanceDispatchResult(
               RpcUtils.getStatus(
-                  TSStatusCode.INTERNAL_SERVER_ERROR, "Interrupted errors: " + e.getMessage())));
+                  TSStatusCode.INTERNAL_SERVER_ERROR,
+                  String.format(
+                      DataNodeQueryMessages.MESSAGE_INTERRUPTED_ERRORS_ARG_876F1B4B,
+                      e.getMessage())),
+              true));
     }
 
     if (dispatchFailures.isEmpty()) {