3

xinyiZzz · Dec 16, 2024 · 4b6a18d · 4b6a18d
1 parent 4dd79a2
commit 4b6a18d
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 9 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
@@ -74,6 +74,7 @@
 import org.apache.doris.transaction.TransactionEntry;
 import org.apache.doris.transaction.TransactionStatus;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
@@ -131,7 +132,9 @@ public enum ConnectType {
     protected volatile long loginTime;
     // for arrow flight
     protected volatile String peerIdentity;
-    protected volatile boolean readyFinalizeArrowFlightSqlRequest = false;
+    protected volatile boolean isFlightSqlConnectProcessorClose = false;
+    protected volatile boolean isExecStatusDone = false;
+    protected volatile boolean finalizeArrowFlightSqlRequestFinished = false;
     private final Map<String, String> preparedQuerys = new HashMap<>();
     private String runningQuery;
     private final List<FlightSqlEndpointsLocation> flightSqlEndpointsLocations = Lists.newArrayList();
@@ -732,12 +735,24 @@ public String getPeerIdentity() {
         return peerIdentity;
     }
 
-    public void setReadyFinalizeArrowFlightSqlRequest() {
-        this.readyFinalizeArrowFlightSqlRequest = true;
+    public void setIsFlightSqlConnectProcessorClose() {
+        this.isFlightSqlConnectProcessorClose = true;
     }
 
-    public boolean getReadyFinalizeArrowFlightSqlRequest() {
-        return readyFinalizeArrowFlightSqlRequest;
+    public boolean getIsFlightSqlConnectProcessorClose() {
+        return isFlightSqlConnectProcessorClose;
+    }
+
+    public void setIsExecStatusDone() {
+        this.isExecStatusDone = true;
+    }
+
+    public boolean getIsExecStatusDone() {
+        return isExecStatusDone;
+    }
+
+    public boolean getFinalizeArrowFlightSqlRequestFinished() {
+        return finalizeArrowFlightSqlRequestFinished;
     }
 
     public FlightSqlChannel getFlightSqlChannel() {
@@ -835,6 +850,8 @@ public void addReturnResultFromRemoteExecutor(StmtExecutor executor) {
      * execute FlightSqlConnectProcessorClose.
      */
     public void finalizeArrowFlightSqlRequest() {
+        Preconditions.checkState(!finalizeArrowFlightSqlRequestFinished);
+        finalizeArrowFlightSqlRequestFinished = true;
         boolean setThreadLocal = false;
         if (get() == null) {
             setThreadLocalInfo();

diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java
@@ -272,10 +272,11 @@ public TReportExecStatusResult reportExecStatus(TReportExecStatusParams params,
         }
         ConnectContext ctx = info.getConnectContext();
         if (ctx.getConnectType() == ConnectType.ARROW_FLIGHT_SQL && params.isDone()) {
-            if (ctx.getReadyFinalizeArrowFlightSqlRequest()) {
+            // `params.isDone` may be true in multiple reportExecStatus.
+            if (ctx.getIsFlightSqlConnectProcessorClose() && !ctx.getFinalizeArrowFlightSqlRequestFinished()) {
                 ctx.finalizeArrowFlightSqlRequest();
             }
-            ctx.setReadyFinalizeArrowFlightSqlRequest();
+            ctx.setIsExecStatusDone();
         }
         result.setStatus(new TStatus(TStatusCode.OK));
         return result;

diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java
@@ -193,10 +193,10 @@ public void fetchArrowFlightSchema(int timeoutMs) {
 
     @Override
     public void close() throws Exception {
-        if (ctx.getReadyFinalizeArrowFlightSqlRequest()) {
+        if (ctx.getIsExecStatusDone()) {
             ctx.finalizeArrowFlightSqlRequest();
         }
-        ctx.setReadyFinalizeArrowFlightSqlRequest();
+        ctx.setIsFlightSqlConnectProcessorClose();
         ConnectContext.remove();
     }
 }