[fix](arrow-flight) Keep coordinator alive across GetFlightInfo/DoGet for external table scan

fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java

@@ -934,6 +934,40 @@ public void clear() {
         statementContext = null;
     }
 
+    // Arrow Flight SQL only.
+    // Executors of already-planned queries whose results are produced on the BE and pulled later
+    // during the DoGet phase. Their coordinators must stay alive until the BE finishes scanning:
+    // an external-table scan in batch mode lazily fetches splits from the FE (a batch SplitSource
+    // held by the coordinator's scan nodes), so closing the coordinator at the end of
+    // GetFlightInfo would release the SplitSource too early and make the BE's fetchSplitBatch fail
+    // with "Split source X is released". These executors are finalized when the next query starts
+    // on this connection, or when the connection is torn down. See #62259.
+    private final List<StmtExecutor> flightSqlDeferredExecutors = new ArrayList<>();
+
+    public void addFlightSqlDeferredExecutor(StmtExecutor executor) {
+        synchronized (flightSqlDeferredExecutors) {
+            flightSqlDeferredExecutors.add(executor);
+        }
+    }
+
+    public void closeFlightSqlDeferredExecutors() {
+        List<StmtExecutor> toClose;
+        synchronized (flightSqlDeferredExecutors) {
+            if (flightSqlDeferredExecutors.isEmpty()) {
+                return;
+            }
+            toClose = new ArrayList<>(flightSqlDeferredExecutors);
+            flightSqlDeferredExecutors.clear();
+        }
+        for (StmtExecutor deferredExecutor : toClose) {
+            try {
+                deferredExecutor.finalizeArrowFlightQuery();
+            } catch (Throwable t) {
+                LOG.warn("failed to finalize deferred arrow flight executor", t);
+            }
+        }
+    }
+
     /**
      * This method is idempotent.
      */

fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java

@@ -199,6 +199,10 @@ public class StmtExecutor {
 
     @Setter
     private volatile Coordinator coord = null;
+    // Arrow Flight SQL: when true, this query's coordinator is kept alive past GetFlightInfo and
+    // is finalized later by ConnectContext (see #62259), so the eager close in executeAndSendResult
+    // is skipped.
+    private volatile boolean deferredForArrowFlight = false;
     private MasterOpExecutor masterOpExecutor = null;
     private RedirectStatus redirectStatus = null;
     private Planner planner;
@@ -1035,6 +1039,23 @@ public void finalizeQuery() {
         QeProcessorImpl.INSTANCE.unregisterQuery(context.queryId());
     }
 
+    public boolean isDeferredForArrowFlight() {
+        return deferredForArrowFlight;
+    }
+
+    // Finalize an Arrow Flight query whose coordinator was kept alive across the
+    // GetFlightInfo -> DoGet phases: close the coordinator (releasing external-table batch
+    // SplitSources and the query queue slot) and then unregister the query. See #62259.
+    public void finalizeArrowFlightQuery() {
+        try {
+            if (coord != null) {
+                coord.close();
+            }
+        } finally {
+            finalizeQuery();
+        }
+    }
+
     private void handleQueryWithRetry(TUniqueId queryId) throws Exception {
         // queue query here
         int retryTime = Config.max_query_retry_time;
@@ -1470,6 +1491,22 @@ public void executeAndSendResult(boolean isOutfileQuery, boolean isSendFields,
             if (context.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)) {
                 Preconditions.checkState(!context.isReturnResultFromLocal());
                 profile.getSummaryProfile().setTempStartTime();
+                // Defer closing the coordinator to ConnectContext (closed on the next query or
+                // connection teardown) instead of in the finally block below. This gate covers
+                // every Arrow Flight query whose results are produced on the BE (coordBase ==
+                // coord) -- internal-table and external, batch or not. It is REQUIRED only for an
+                // external-table scan in batch mode, where the BE lazily fetches splits from the FE
+                // during the later DoGet phase, so closing the coordinator here would release its
+                // batch SplitSource too early and break DoGet. Other remote-result queries do not
+                // need deferral (the BE buffers their result independently) but are captured by the
+                // same gate; the trade-off is their coordinator, query queue slot and query
+                // registration stay held until the next query / teardown instead of being released
+                // at the end of GetFlightInfo. Point queries use a different coordBase (not
+                // deferred). See #62259.
+                if (coordBase == coord) {
+                    deferredForArrowFlight = true;
+                    context.addFlightSqlDeferredExecutor(this);
+                }
                 return;
             }
 
@@ -1579,7 +1616,12 @@ public void executeAndSendResult(boolean isOutfileQuery, boolean isSendFields,
             this.coord = null;
             throw e;
         } finally {
-            coordBase.close();
+            // For deferred Arrow Flight queries the coordinator is closed later by ConnectContext
+            // (next query / connection teardown), so the BE can still fetch splits during DoGet.
+            // See #62259.
+            if (!deferredForArrowFlight) {
+                coordBase.close();
+            }
         }
     }
 

fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java

@@ -189,6 +189,10 @@ private FlightInfo executeQueryStatement(String peerIdentity, ConnectContext con
         try {
             Preconditions.checkState(null != connectContext);
             Preconditions.checkState(!query.isEmpty());
+            // Finalize the previous query's coordinator on this connection whose close was
+            // deferred (Arrow Flight keeps it alive across GetFlightInfo -> DoGet so the BE can
+            // fetch external-table splits during DoGet). By now the previous DoGet is done. #62259
+            connectContext.closeFlightSqlDeferredExecutors();
             // After the previous query was executed, there was no getStreamStatement to take away the result.
             connectContext.getFlightSqlChannel().reset();
             connectContext.clearFlightSqlEndpointsLocations();
@@ -280,6 +284,14 @@ private FlightInfo executeQueryStatement(String peerIdentity, ConnectContext con
                 }
             }
         } catch (Throwable e) {
+            // GetFlightInfo failed (e.g. the BE Arrow schema fetch above timed out or returned an
+            // error) after this query's coordinator may already have been deferred during planning.
+            // No FlightInfo is returned, so no DoGet will ever pull this query's results; finalize
+            // the deferred coordinator now (releasing its external-table batch SplitSource, query
+            // queue slot and query registration) instead of leaking it until the next query starts
+            // or the connection is torn down. The previous query's deferred coordinator was already
+            // finalized at the top of this method, so this only closes this failed query. See #62259.
+            connectContext.closeFlightSqlDeferredExecutors();
             String errMsg = "get flight info statement failed, " + e.getMessage() + ", " + Util.getRootCauseMessage(e)
                     + ", error code: " + connectContext.getState().getErrorCode() + ", error msg: "
                     + connectContext.getState().getErrorMessage();

fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java

@@ -196,11 +196,20 @@ public void fetchArrowFlightSchema(int timeoutMs) {
     @Override
     public void close() throws Exception {
         ctx.setCommand(MysqlCommand.COM_SLEEP);
+        // Executors whose results are pulled from the BE keep their coordinator alive past
+        // GetFlightInfo (registered as deferred executors on the ConnectContext) so the BE can
+        // still fetch external-table splits during DoGet. Do NOT finalize those here; they are
+        // finalized when the next query starts or the connection is torn down. Executors that are
+        // not deferred (local results, or a query that already failed) are finalized now. See #62259.
         for (StmtExecutor asynExecutor : returnResultFromRemoteExecutor) {
-            asynExecutor.finalizeQuery();
+            if (!asynExecutor.isDeferredForArrowFlight()) {
+                asynExecutor.finalizeQuery();
+            }
         }
         returnResultFromRemoteExecutor.clear();
-        executor.finalizeQuery();
+        if (executor != null && !executor.isDeferredForArrowFlight()) {
+            executor.finalizeQuery();
+        }
         ctx.clear();
         ConnectContext.remove();
     }

fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgr.java

@@ -55,6 +55,11 @@ public int registerConnection(ConnectContext ctx) {
 
     @Override
     public void unregisterConnection(ConnectContext ctx) {
+        // Finalize any Arrow Flight query whose coordinator was kept alive across the
+        // GetFlightInfo -> DoGet phases (see #62259), releasing its resources (e.g. external-table
+        // batch SplitSources and the query queue slot) when the connection is torn down: idle/query
+        // timeout, bearer token expiry, or explicit CloseSession all reach here.
+        ctx.closeFlightSqlDeferredExecutors();
         ctx.closeTxn();
         if (connectionMap.remove(ctx.getConnectionId()) != null) {
             numberConnection.decrementAndGet();

fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java

@@ -753,4 +753,62 @@ public void testConnectAttributesSetNull() {
         Assert.assertNotNull(ctx.getConnectAttributes());
         Assert.assertTrue(ctx.getConnectAttributes().isEmpty());
     }
+
+    // Arrow Flight SQL keeps a query's coordinator alive across GetFlightInfo -> DoGet (see #62259).
+    // closeFlightSqlDeferredExecutors() is the single place that releases those deferred coordinators
+    // (and with them the external-table batch SplitSource and the query queue slot). The following
+    // tests pin the leak-prevention contract of that method: every deferred executor is finalized,
+    // the list is cleared so nothing is finalized twice or retained, and one failing executor does
+    // not strand the others' resources.
+
+    @Test
+    public void testCloseFlightSqlDeferredExecutorsFinalizesEachExecutor() {
+        ConnectContext ctx = new ConnectContext();
+        StmtExecutor deferred1 = Mockito.mock(StmtExecutor.class);
+        StmtExecutor deferred2 = Mockito.mock(StmtExecutor.class);
+        ctx.addFlightSqlDeferredExecutor(deferred1);
+        ctx.addFlightSqlDeferredExecutor(deferred2);
+
+        ctx.closeFlightSqlDeferredExecutors();
+
+        // Both deferred coordinators must be finalized, otherwise their SplitSource and query queue
+        // slot leak after the DoGet phase.
+        Mockito.verify(deferred1).finalizeArrowFlightQuery();
+        Mockito.verify(deferred2).finalizeArrowFlightQuery();
+    }
+
+    @Test
+    public void testCloseFlightSqlDeferredExecutorsClearsListSoSecondCallIsNoOp() {
+        ConnectContext ctx = new ConnectContext();
+        StmtExecutor deferred = Mockito.mock(StmtExecutor.class);
+        ctx.addFlightSqlDeferredExecutor(deferred);
+
+        // More than one teardown path can fire for the same connection (e.g. the next query cleans
+        // up, then the connection is later torn down). The list must be cleared after the first
+        // call so the executor is finalized exactly once and is not retained (leaked) afterwards.
+        ctx.closeFlightSqlDeferredExecutors();
+        ctx.closeFlightSqlDeferredExecutors();
+
+        Mockito.verify(deferred, Mockito.times(1)).finalizeArrowFlightQuery();
+    }
+
+    @Test
+    public void testCloseFlightSqlDeferredExecutorsFinalizesRemainingWhenOneFails() {
+        ConnectContext ctx = new ConnectContext();
+        StmtExecutor failing = Mockito.mock(StmtExecutor.class);
+        StmtExecutor healthy = Mockito.mock(StmtExecutor.class);
+        Mockito.doThrow(new RuntimeException("finalize failed")).when(failing).finalizeArrowFlightQuery();
+        ctx.addFlightSqlDeferredExecutor(failing);
+        ctx.addFlightSqlDeferredExecutor(healthy);
+
+        // A single bad coordinator must not abort the cleanup: the call must not throw, and the
+        // healthy executor must still be finalized so its resources are released.
+        ctx.closeFlightSqlDeferredExecutors();
+        Mockito.verify(healthy).finalizeArrowFlightQuery();
+
+        // The list is cleared up front, so neither executor is reprocessed on a later teardown.
+        ctx.closeFlightSqlDeferredExecutors();
+        Mockito.verify(failing, Mockito.times(1)).finalizeArrowFlightQuery();
+        Mockito.verify(healthy, Mockito.times(1)).finalizeArrowFlightQuery();
+    }
 }

fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java

@@ -29,6 +29,8 @@
 import org.apache.doris.planner.ResultFileSink;
 import org.apache.doris.qe.CommonResultSet.CommonResultSetMetaData;
 import org.apache.doris.qe.ConnectContext.ConnectType;
+import org.apache.doris.thrift.TQueryOptions;
+import org.apache.doris.thrift.TUniqueId;
 import org.apache.doris.utframe.TestWithFeService;
 
 import com.google.common.collect.Lists;
@@ -71,6 +73,40 @@ public void testShowNull() throws Exception {
         Assert.assertEquals(QueryState.MysqlStateType.OK, connectContext.getState().getStateType());
     }
 
+    // Arrow Flight SQL keeps a query's coordinator alive across GetFlightInfo -> DoGet (see #62259);
+    // it is released later by finalizeArrowFlightQuery(), which closes the coordinator and then
+    // unregisters the query. The close and the unregister must be independent: if coord.close()
+    // throws, the query registration must still be released (the try/finally), otherwise the query
+    // leaks in QeProcessorImpl forever. The thrown error is expected to propagate to the caller
+    // (ConnectContext.closeFlightSqlDeferredExecutors), which catches and logs it.
+    @Test
+    public void testFinalizeArrowFlightQueryUnregistersQueryEvenIfCoordCloseThrows() throws Exception {
+        StmtExecutor stmtExecutor = new StmtExecutor(connectContext, "");
+        TUniqueId queryId = new TUniqueId(0x6226259L, 0x62259L);
+        connectContext.setQueryId(queryId);
+
+        Coordinator coord = Mockito.mock(Coordinator.class);
+        Mockito.when(coord.getQueryOptions()).thenReturn(new TQueryOptions());
+        Mockito.doThrow(new RuntimeException("coord close failed")).when(coord).close();
+        stmtExecutor.setCoord(coord);
+
+        // Simulate the in-flight query whose results DoGet is still pulling.
+        QeProcessorImpl.INSTANCE.registerQuery(queryId, new QeProcessorImpl.QueryInfo(coord));
+        Assert.assertNotNull(QeProcessorImpl.INSTANCE.getCoordinator(queryId));
+
+        try {
+            stmtExecutor.finalizeArrowFlightQuery();
+            Assert.fail("expected coord.close() failure to propagate after the query is unregistered");
+        } catch (RuntimeException e) {
+            Assert.assertEquals("coord close failed", e.getMessage());
+        }
+
+        // The coordinator close was attempted (releases SplitSource + query queue slot) ...
+        Mockito.verify(coord).close();
+        // ... and despite it failing, the query registration was still released (no leak).
+        Assert.assertNull(QeProcessorImpl.INSTANCE.getCoordinator(queryId));
+    }
+
     @Test
     public void testKill() throws Exception {
         StmtExecutor stmtExecutor = new StmtExecutor(connectContext, "");