[To dev/1.3] Fix timeout didn't pass to client bug (#15743)

(cherry picked from commit 232e0c9)

Author: JackieTien97

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/exception/query/QueryTimeoutRuntimeException.java

@@ -29,4 +29,8 @@ public QueryTimeoutRuntimeException(long startTime, long currentTime, long timeo
         String.format(
             QUERY_TIMEOUT_EXCEPTION_MESSAGE, startTime, startTime + timeout, currentTime));
   }
+
+  public QueryTimeoutRuntimeException(String message) {
+    super(message);
+  }
 }

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

@@ -37,7 +37,9 @@
 import javax.annotation.concurrent.NotThreadSafe;
 
 import java.util.LinkedList;
+import java.util.Optional;
 import java.util.Queue;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 
 import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
@@ -81,6 +83,8 @@ public class SharedTsBlockQueue {
   private long maxBytesCanReserve =
       IoTDBDescriptor.getInstance().getConfig().getMaxBytesPerFragmentInstance();
 
+  private volatile Throwable abortedCause = null;
+
   // used for SharedTsBlockQueue listener
   private final ExecutorService executorService;
 
@@ -177,6 +181,18 @@ public void setNoMoreTsBlocks(boolean noMoreTsBlocks) {
    */
   public TsBlock remove() {
     if (closed) {
+      // try throw underlying exception instead of "Source handle is aborted."
+      if (abortedCause != null) {
+        throw new IllegalStateException(abortedCause);
+      }
+      try {
+        blocked.get();
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new IllegalStateException(e);
+      } catch (ExecutionException e) {
+        throw new IllegalStateException(e.getCause() == null ? e : e.getCause());
+      }
       throw new IllegalStateException("queue has been destroyed");
     }
     TsBlock tsBlock = queue.remove();
@@ -332,6 +348,7 @@ public void abort(Throwable t) {
     if (closed) {
       return;
     }
+    abortedCause = t;
     closed = true;
     if (!blocked.isDone()) {
       blocked.setException(t);
@@ -354,4 +371,8 @@ public void abort(Throwable t) {
       bufferRetainedSizeInBytes = 0;
     }
   }
+
+  public Optional<Throwable> getAbortedCause() {
+    return Optional.ofNullable(abortedCause);
+  }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/LocalSourceHandle.java

@@ -36,6 +36,8 @@
 import org.slf4j.LoggerFactory;
 
 import java.nio.ByteBuffer;
+import java.util.Optional;
+import java.util.concurrent.ExecutionException;
 
 import static com.google.common.util.concurrent.Futures.nonCancellationPropagating;
 import static org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.createFullIdFrom;
@@ -156,6 +158,7 @@ public ByteBuffer getSerializedTsBlock() throws IoTDBException {
   @Override
   public boolean isFinished() {
     synchronized (queue) {
+      checkSharedQueueIfAborted();
       return queue.hasNoMoreTsBlocks() && queue.isEmpty();
     }
   }
@@ -251,10 +254,28 @@ public void close() {
   }
 
   private void checkState() {
-    if (aborted) {
-      throw new IllegalStateException("Source handle is aborted.");
-    } else if (closed) {
-      throw new IllegalStateException("Source Handle is closed.");
+    if (aborted || closed) {
+      checkSharedQueueIfAborted();
+      if (queue.isBlocked().isDone()) {
+        // try throw underlying exception instead of "Source handle is aborted."
+        try {
+          queue.isBlocked().get();
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          throw new IllegalStateException(e);
+        } catch (ExecutionException e) {
+          throw new IllegalStateException(e.getCause() == null ? e : e.getCause());
+        }
+      }
+      throw new IllegalStateException(
+          "LocalSinkChannel state is ." + (aborted ? "ABORTED" : "CLOSED"));
+    }
+  }
+
+  private void checkSharedQueueIfAborted() {
+    Optional<Throwable> abortedCause = queue.getAbortedCause();
+    if (abortedCause.isPresent()) {
+      throw new IllegalStateException(abortedCause.get());
     }
   }
 

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

@@ -26,6 +26,7 @@
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.exception.IoTDBException;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.exception.query.QueryTimeoutRuntimeException;
 import org.apache.iotdb.db.queryengine.common.FragmentInstanceId;
 import org.apache.iotdb.db.queryengine.common.QueryId;
 import org.apache.iotdb.db.queryengine.execution.driver.IDriver;
@@ -56,7 +57,6 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicLong;
 
 import static java.util.Objects.requireNonNull;
@@ -421,8 +421,10 @@ private void cancelTimeoutFlushingInstances() {
             execution
                 .getStateMachine()
                 .failed(
-                    new TimeoutException(
-                        "Query has executed more than " + execution.getTimeoutInMs() + "ms"));
+                    new QueryTimeoutRuntimeException(
+                        "Query has executed more than "
+                            + execution.getTimeoutInMs()
+                            + "ms, and now is in flushing state"));
           }
         });
   }