From ae526d05d12b3d330b071b0e881969b8d493ab57 Mon Sep 17 00:00:00 2001 From: morningman Date: Wed, 24 Jun 2026 22:02:30 +0800 Subject: [PATCH 1/6] [fix](arrow-flight) Keep coordinator alive across GetFlightInfo/DoGet for external table scan Arrow Flight SQL executes a query in two phases: GetFlightInfo (plan + submit to BE) and DoGet (the client pulls results from the BE). For an external table scan in batch split mode, the BE keeps scanning during DoGet and lazily fetches file splits from the FE via the fetchSplitBatch RPC, using a batch SplitSource that the FE's coordinator holds (through its scan nodes). The problem: the FE closed the coordinator at the end of GetFlightInfo (StmtExecutor.executeAndSendResult's finally -> Coordinator.close() -> ScanNode.stop() -> SplitSourceManager.removeSplitSource()) and also unregistered it (FlightSqlConnectProcessor.close() -> StmtExecutor.finalizeQuery()). So by the time the BE called fetchSplitBatch during DoGet, the SplitSource was already gone and the query failed with "Split source X is released". MySQL protocol is unaffected because plan + execute share one request and the coordinator stays alive until all results are consumed. This keeps the coordinator (and its SplitSource) alive across the two phases and cleans it up reliably: - StmtExecutor: for an Arrow Flight query that produces results on the BE (coordBase == coord), mark it deferred, register the executor on the ConnectContext, and skip the eager Coordinator.close() in the finally. A failed query (exec threw) is not deferred and is closed as before. - ConnectContext: hold the deferred executors and provide closeFlightSqlDeferredExecutors() to close their coordinators (releasing the SplitSource and the query queue slot) and unregister the queries. - FlightSqlConnectProcessor.close(): do not finalize deferred executors. - DorisFlightSqlProducer: finalize the previous query's deferred coordinator when the next query starts on the connection. - FlightSqlConnectPoolMgr.unregisterConnection(): finalize deferred coordinators when the connection is torn down. All teardown paths (idle/query timeout, bearer token expiry, explicit CloseSession) reach here, so an abandoned connection cannot leak the coordinator. This is the root-cause fix for #62259. The BE-side error-path hardening (so any fetchSplitBatch failure fails gracefully instead of crashing) is handled separately in #64797. --- .../org/apache/doris/qe/ConnectContext.java | 34 +++++++++++++++++ .../org/apache/doris/qe/StmtExecutor.java | 38 ++++++++++++++++++- .../arrowflight/DorisFlightSqlProducer.java | 4 ++ .../FlightSqlConnectProcessor.java | 13 ++++++- .../sessions/FlightSqlConnectPoolMgr.java | 5 +++ 5 files changed, 91 insertions(+), 3 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 index cc24f38f4d3f27..c79b8d824ddade 100644 --- 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 @@ -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 flightSqlDeferredExecutors = new ArrayList<>(); + + public void addFlightSqlDeferredExecutor(StmtExecutor executor) { + synchronized (flightSqlDeferredExecutors) { + flightSqlDeferredExecutors.add(executor); + } + } + + public void closeFlightSqlDeferredExecutors() { + List 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. */ diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index 05aba741112ff9..159526eb6d8147 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/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,16 @@ public void executeAndSendResult(boolean isOutfileQuery, boolean isSendFields, if (context.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)) { Preconditions.checkState(!context.isReturnResultFromLocal()); profile.getSummaryProfile().setTempStartTime(); + // The BE keeps scanning and pulling results after GetFlightInfo returns; for an + // external-table scan in batch mode it lazily fetches splits from the FE during the + // later DoGet phase. Keep the coordinator (and the batch SplitSource its scan nodes + // hold) alive by deferring its close to ConnectContext, instead of closing it in the + // finally block below which would release the SplitSource too early. See #62259. + // (Point queries use a different coordBase and are not deferred.) + if (coordBase == coord) { + deferredForArrowFlight = true; + context.addFlightSqlDeferredExecutor(this); + } return; } @@ -1579,7 +1610,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(); + } } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java index e2d197cb537ea3..75415327386461 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java +++ b/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(); 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 index a63fab3e9608f7..aa69339fd31a2c 100644 --- 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 @@ -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(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgr.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgr.java index 3002116b386aa7..743f5c49d44a60 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgr.java +++ b/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(); From 0dcc3e7a2ee6525eca0f830cd40c2a68a682202d Mon Sep 17 00:00:00 2001 From: morningman Date: Wed, 24 Jun 2026 22:16:06 +0800 Subject: [PATCH 2/6] [test](arrow-flight) Add regression for Iceberg scan over Arrow Flight in batch split mode Reproduces #62259: scanning an Iceberg external table over Arrow Flight SQL in batch split mode used to fail ("Split source X is released" / BE crash) because the FE released the async SplitSource at the end of GetFlightInfo, before the BE fetched splits during DoGet. The test forces batch split mode on the Arrow Flight session (num_files_in_batch_mode=1), asserts via explain that the scan really uses the batch SplitSource path ("approximate"), then scans format_v2.sample_cow_orc over Arrow Flight and checks all 1000 rows come back. Running several queries on the same connection also exercises the cleanup of the previous query's deferred coordinator. The test is skipped when the Iceberg env or the Arrow Flight endpoint is not configured. --- ...t_iceberg_arrow_flight_split_source.groovy | 96 +++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy diff --git a/regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy b/regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy new file mode 100644 index 00000000000000..171988d565c6aa --- /dev/null +++ b/regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy @@ -0,0 +1,96 @@ +// 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. + +// Regression for https://github.com/apache/doris/issues/62259 +// +// Querying an Iceberg external table over Arrow Flight SQL in batch split mode used to fail +// (BE crash / "Split source X is released"). Arrow Flight executes a query in two phases: +// GetFlightInfo (plan + submit to BE) then DoGet (the client pulls results from the BE). In +// batch split mode the BE keeps scanning during DoGet and lazily fetches file splits from the +// FE via the fetchSplitBatch RPC, using an async SplitSource that the FE coordinator holds. The +// FE used to release that SplitSource at the end of GetFlightInfo, before the BE's DoGet, so the +// split fetch failed. The MySQL protocol is unaffected because plan + execute share one request. +// +// This test forces batch split mode on the Arrow Flight session and scans the table, which must +// now return all rows. +suite("test_iceberg_arrow_flight_split_source", "p0,external") { + String enabled = context.config.otherConfigs.get("enableIcebergTest") + if (enabled == null || !enabled.equalsIgnoreCase("true")) { + logger.info("disable iceberg test.") + return + } + + // The bug only manifests over the Arrow Flight SQL protocol. Skip when it is not configured. + String arrowFlightHost = context.config.otherConfigs.get("extArrowFlightSqlHost") + if (arrowFlightHost == null || arrowFlightHost.isEmpty()) { + logger.info("extArrowFlightSqlHost is not configured, skip the test.") + return + } + + String rest_port = context.config.otherConfigs.get("iceberg_rest_uri_port") + String minio_port = context.config.otherConfigs.get("iceberg_minio_port") + String externalEnvIp = context.config.otherConfigs.get("externalEnvIp") + String catalog_name = "test_iceberg_arrow_flight_split_source" + // sample_cow_orc has 1000 rows; with num_files_in_batch_mode=1 a plain scan uses batch mode. + String table = "${catalog_name}.format_v2.sample_cow_orc" + + sql """drop catalog if exists ${catalog_name}""" + sql """CREATE CATALOG ${catalog_name} PROPERTIES ( + 'type'='iceberg', + 'iceberg.catalog.type'='rest', + 'uri' = 'http://${externalEnvIp}:${rest_port}', + "s3.access_key" = "admin", + "s3.secret_key" = "password", + "s3.endpoint" = "http://${externalEnvIp}:${minio_port}", + "s3.region" = "us-east-1" + );""" + + try { + // Baseline over the MySQL protocol (works regardless of the bug). + def expected = sql """ select count(*) from ${table}; """ + long expectedRows = (expected[0][0] as long) + assert expectedRows > 0 : "precondition: ${table} should not be empty" + + // Force batch split mode on the Arrow Flight session (a separate session from the MySQL + // connection above, so the variables must be set here). With num_files_in_batch_mode=1 + // even a single-file scan builds the async SplitSource that triggers #62259. + arrow_flight_sql """ set enable_external_table_batch_mode = true; """ + arrow_flight_sql """ set num_files_in_batch_mode = 1; """ + + // Make sure the Arrow Flight session really uses the batch SplitSource path, so the test + // cannot silently pass on the non-batch path. "approximate" only appears in batch mode. + def explainRows = arrow_flight_sql """ explain select * from ${table}; """ + boolean isBatch = explainRows.any { row -> + row.any { cell -> cell != null && cell.toString().contains("approximate") } + } + assert isBatch : "expected batch split mode (approximate) in the Arrow Flight plan, got: ${explainRows}" + + // The regression: a real data scan over Arrow Flight SQL (not count(*), which is pushed + // down and bypasses batch mode). Before the fix this failed with "Split source X is + // released" or crashed the BE; now it must return all rows. + def flightResult = arrow_flight_sql """ select * from ${table}; """ + assertEquals(expectedRows, (flightResult.size() as long)) + + // A second scan on the same connection also exercises cleanup of the previous query's + // deferred coordinator when the next query starts. + def flightLimited = arrow_flight_sql """ select * from ${table} limit 10; """ + assert flightLimited.size() > 0 && flightLimited.size() <= 10 : "unexpected row count: ${flightLimited.size()}" + } finally { + arrow_flight_sql """ set num_files_in_batch_mode = 1024; """ + sql """drop catalog if exists ${catalog_name}""" + } +} From e57f3d65cfd13b3985cc7f7deec72cad4b46fb5d Mon Sep 17 00:00:00 2001 From: morningman Date: Thu, 25 Jun 2026 09:50:29 +0800 Subject: [PATCH 3/6] [test](arrow-flight) Connect the Iceberg split-source test to the cluster's real Arrow Flight port The external regression cluster serves Arrow Flight SQL on the default port (FE 8070) rather than the configured extArrowFlightSqlPort (8081), so the framework's arrow_flight_sql() helper failed at connect with "Connection refused: :8081" before the test ever reached the batch split-source path it is meant to exercise. The error surfaced on the finally line because the reconnect there also failed and masked the original exception. Read the live Arrow Flight port from SHOW FRONTENDS and open a dedicated connection against it (as the remote_doris tests already do), skip when Arrow Flight is disabled, and make the finally best-effort so a dead endpoint can no longer mask the real failure or leak the catalog. --- ...t_iceberg_arrow_flight_split_source.groovy | 47 ++++++++++++++++--- 1 file changed, 41 insertions(+), 6 deletions(-) diff --git a/regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy b/regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy index 171988d565c6aa..e721d0e9d8eab6 100644 --- a/regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy +++ b/regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +import java.sql.Connection +import java.sql.DriverManager + +import org.apache.doris.regression.util.JdbcUtils + // Regression for https://github.com/apache/doris/issues/62259 // // Querying an Iceberg external table over Arrow Flight SQL in batch split mode used to fail @@ -41,6 +46,23 @@ suite("test_iceberg_arrow_flight_split_source", "p0,external") { return } + // The framework's arrow_flight_sql() helper always dials extArrowFlightSqlPort, but that + // configured value does not always match this cluster's real Arrow Flight SQL port (e.g. the + // external pipeline serves Arrow Flight on the default 8070 while extArrowFlightSqlPort is + // 8081). Read the live port from SHOW FRONTENDS and open our own connection against it, like + // the remote_doris tests do, so the test connects to this cluster's actual endpoint. + def frontends = sql """ show frontends """ + String arrowFlightPort = frontends[0][6].toString() + if (!arrowFlightPort.isInteger() || (arrowFlightPort as int) <= 0) { + logger.info("Arrow Flight SQL is disabled on this cluster (port=${arrowFlightPort}), skip the test.") + return + } + String arrowFlightUser = context.config.otherConfigs.get("extArrowFlightSqlUser") + String arrowFlightPassword = context.config.otherConfigs.get("extArrowFlightSqlPassword") + Class.forName("org.apache.arrow.driver.jdbc.ArrowFlightJdbcDriver") + String arrowFlightUrl = "jdbc:arrow-flight-sql://${arrowFlightHost}:${arrowFlightPort}" + + "/?useServerPrepStmts=false&useSSL=false&useEncryption=false" + String rest_port = context.config.otherConfigs.get("iceberg_rest_uri_port") String minio_port = context.config.otherConfigs.get("iceberg_minio_port") String externalEnvIp = context.config.otherConfigs.get("externalEnvIp") @@ -59,21 +81,30 @@ suite("test_iceberg_arrow_flight_split_source", "p0,external") { "s3.region" = "us-east-1" );""" + Connection flightConn = null try { // Baseline over the MySQL protocol (works regardless of the bug). def expected = sql """ select count(*) from ${table}; """ long expectedRows = (expected[0][0] as long) assert expectedRows > 0 : "precondition: ${table} should not be empty" + // A dedicated Arrow Flight SQL connection to this cluster's real port. Run a statement over + // it the same way arrow_flight_sql() does, via JdbcUtils.executeToList. + flightConn = DriverManager.getConnection(arrowFlightUrl, arrowFlightUser, arrowFlightPassword) + def flightSql = { String stmt -> + def (rows, meta) = JdbcUtils.executeToList(flightConn, stmt) + return rows + } + // Force batch split mode on the Arrow Flight session (a separate session from the MySQL // connection above, so the variables must be set here). With num_files_in_batch_mode=1 // even a single-file scan builds the async SplitSource that triggers #62259. - arrow_flight_sql """ set enable_external_table_batch_mode = true; """ - arrow_flight_sql """ set num_files_in_batch_mode = 1; """ + flightSql """ set enable_external_table_batch_mode = true """ + flightSql """ set num_files_in_batch_mode = 1 """ // Make sure the Arrow Flight session really uses the batch SplitSource path, so the test // cannot silently pass on the non-batch path. "approximate" only appears in batch mode. - def explainRows = arrow_flight_sql """ explain select * from ${table}; """ + def explainRows = flightSql """ explain select * from ${table} """ boolean isBatch = explainRows.any { row -> row.any { cell -> cell != null && cell.toString().contains("approximate") } } @@ -82,15 +113,19 @@ suite("test_iceberg_arrow_flight_split_source", "p0,external") { // The regression: a real data scan over Arrow Flight SQL (not count(*), which is pushed // down and bypasses batch mode). Before the fix this failed with "Split source X is // released" or crashed the BE; now it must return all rows. - def flightResult = arrow_flight_sql """ select * from ${table}; """ + def flightResult = flightSql """ select * from ${table} """ assertEquals(expectedRows, (flightResult.size() as long)) // A second scan on the same connection also exercises cleanup of the previous query's // deferred coordinator when the next query starts. - def flightLimited = arrow_flight_sql """ select * from ${table} limit 10; """ + def flightLimited = flightSql """ select * from ${table} limit 10 """ assert flightLimited.size() > 0 && flightLimited.size() <= 10 : "unexpected row count: ${flightLimited.size()}" } finally { - arrow_flight_sql """ set num_files_in_batch_mode = 1024; """ + // Close our own connection (best effort) so a dead endpoint cannot mask the real failure, + // then drop the catalog over the reliable MySQL connection. + if (flightConn != null) { + try { flightConn.close() } catch (Throwable ignore) {} + } sql """drop catalog if exists ${catalog_name}""" } } From d4a974dd0fcbbdd1fc3e8997eebe3c2254531090 Mon Sep 17 00:00:00 2001 From: morningman Date: Thu, 25 Jun 2026 15:19:33 +0800 Subject: [PATCH 4/6] [test](arrow-flight) Add FE unit tests for deferred-coordinator resource cleanup The fix for #62259 keeps an Arrow Flight query's coordinator alive across GetFlightInfo -> DoGet and releases it later via deferred executors held on the ConnectContext. These tests pin the leak-prevention contract of that deferred cleanup so the coordinator (and the external-table batch SplitSource and query queue slot it holds, plus the query registration) cannot leak: - ConnectContextTest: closeFlightSqlDeferredExecutors() finalizes every deferred executor, clears the list (idempotent, no double-close, no retain), and still finalizes the rest when one executor's finalize throws. - FlightSqlConnectPoolMgrTest (new): unregisterConnection() finalizes the deferred executors on connection teardown (idle/query timeout, bearer token expiry, explicit CloseSession), even for a connection that was never registered, so an abandoned connection cannot leak the coordinator. - StmtExecutorTest: finalizeArrowFlightQuery() still unregisters the query (releases the registration) when coord.close() throws (the try/finally). Each test was verified to fail when its corresponding cleanup is removed. --- .../apache/doris/qe/ConnectContextTest.java | 58 +++++++++++++++++ .../org/apache/doris/qe/StmtExecutorTest.java | 36 +++++++++++ .../sessions/FlightSqlConnectPoolMgrTest.java | 62 +++++++++++++++++++ 3 files changed, 156 insertions(+) create mode 100644 fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgrTest.java diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java index 94d922e69876a8..509d1823afc415 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java +++ b/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(); + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java index b7b12bc91c3deb..607cea3b40a302 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/StmtExecutorTest.java +++ b/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, ""); diff --git a/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgrTest.java b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgrTest.java new file mode 100644 index 00000000000000..80dcb5b4346e12 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgrTest.java @@ -0,0 +1,62 @@ +// 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.doris.service.arrowflight.sessions; + +import org.apache.doris.qe.ConnectContext; + +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +public class FlightSqlConnectPoolMgrTest { + + // Arrow Flight SQL keeps a query's coordinator alive across GetFlightInfo -> DoGet (see #62259). + // unregisterConnection() is the catch-all teardown path: idle/query timeout, bearer token expiry + // and explicit CloseSession all reach here. It must finalize the deferred coordinators so an + // abandoned connection cannot leak them (the external-table batch SplitSource and the query + // queue slot the coordinator holds). + @Test + public void testUnregisterConnectionFinalizesDeferredExecutors() { + FlightSqlConnectPoolMgr poolMgr = new FlightSqlConnectPoolMgr(100); + ConnectContext ctx = Mockito.mock(ConnectContext.class); + + poolMgr.unregisterConnection(ctx); + + // The deferred coordinators must be released on teardown even though this connection was + // never registered in the pool (an abandoned connection is still cleaned up, not leaked). + Mockito.verify(ctx).closeFlightSqlDeferredExecutors(); + } + + // Cleanup must run before the connection bookkeeping (closeTxn / map removal), so that a failure + // there cannot strand the deferred coordinators. Verify the deferred executors are finalized + // even when the context is the one stored in the pool. + @Test + public void testUnregisterRegisteredConnectionFinalizesDeferredExecutors() { + FlightSqlConnectPoolMgr poolMgr = new FlightSqlConnectPoolMgr(100); + ConnectContext ctx = Mockito.mock(ConnectContext.class); + Mockito.when(ctx.getConnectionId()).thenReturn(7); + Mockito.when(ctx.getConnectType()).thenReturn(ConnectContext.ConnectType.ARROW_FLIGHT_SQL); + Mockito.when(ctx.getPeerIdentity()).thenReturn("token-7"); + poolMgr.getConnectionMap().put(7, ctx); + + poolMgr.unregisterConnection(ctx); + + Mockito.verify(ctx).closeFlightSqlDeferredExecutors(); + Assert.assertNull(poolMgr.getConnectionMap().get(7)); + } +} From 826d25d0dad61bae05905daa98f19f2a63b9009a Mon Sep 17 00:00:00 2001 From: morningman Date: Sun, 28 Jun 2026 11:06:02 +0800 Subject: [PATCH 5/6] [fix](arrow-flight) Finalize deferred coordinator when GetFlightInfo fails after deferral For an Arrow Flight external-table scan, executeAndSendResult() registers the query's coordinator as a deferred executor on the ConnectContext (kept alive across GetFlightInfo -> DoGet, see #62259) right after submitting it to the BE. GetFlightInfo then still has to fetch the Arrow schema from the BE. If that fetch fails (timeout / non-OK / empty / mismatched schema / RPC error), FlightSqlConnectProcessor.close() skips the deferred executor and the outer catch only rethrows. No FlightInfo is returned, so no DoGet will ever pull the results, yet the coordinator (its external-table batch SplitSource, the query queue slot and the query registration) stays alive until the next query starts or the connection is torn down. Finalize the deferred coordinator on the GetFlightInfo error path: executeQueryStatement's catch now calls closeFlightSqlDeferredExecutors(). At that point the list holds only this failed query's coordinator (the previous one was already finalized at the top of the method), and this covers every post-deferral failure, not just the schema fetch. Add DorisFlightSqlProducerTest.testGetFlightInfoFinalizesDeferredExecutorWhen SchemaFetchFails, which defers a coordinator then fails the schema fetch and asserts the deferred executor is finalized (verified to fail without the fix). --- .../arrowflight/DorisFlightSqlProducer.java | 8 +++ .../DorisFlightSqlProducerTest.java | 65 +++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java index 75415327386461..e64690a54cb10f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducer.java @@ -284,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(); diff --git a/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducerTest.java b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducerTest.java index 3a4b0facace7e1..eede12688517be 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducerTest.java @@ -19,19 +19,23 @@ import org.apache.doris.common.FeConstants; import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; import org.apache.doris.service.arrowflight.results.FlightSqlChannel; import org.apache.doris.service.arrowflight.sessions.FlightSessionsManager; import org.apache.doris.service.arrowflight.sessions.FlightSqlConnectContext; +import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightProducer.CallContext; import org.apache.arrow.flight.FlightProducer.StreamListener; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.Result; import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementRequest; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementQuery; import org.junit.After; import org.junit.Assert; import org.junit.Before; import org.junit.Test; +import org.mockito.MockedConstruction; import org.mockito.Mockito; import java.util.concurrent.CountDownLatch; @@ -142,4 +146,65 @@ public void onCompleted() { producer.close(); } } + + // Arrow Flight SQL keeps a query's coordinator alive across GetFlightInfo -> DoGet (see #62259): + // executeAndSendResult() registers it as a deferred executor on the ConnectContext right after + // submitting it to the BE. GetFlightInfo then still has to fetch the Arrow schema from the BE. + // If that fetch fails (timeout / non-OK / empty / mismatched schema / RPC error), no FlightInfo + // is returned, so no DoGet will ever pull this query's results. The deferred coordinator must be + // finalized on this error path; otherwise its external-table batch SplitSource, query queue slot + // and query registration leak until the next query starts or the connection is torn down. + @Test + public void testGetFlightInfoFinalizesDeferredExecutorWhenSchemaFetchFails() throws Exception { + // A flight ConnectContext whose getFlightSqlChannel() works (the base context throws). + ConnectContext ctx = Mockito.spy(new ConnectContext()); + Mockito.doReturn(Mockito.mock(FlightSqlChannel.class)).when(ctx).getFlightSqlChannel(); + + // Stands in for the just-planned external-table query whose results DoGet would pull from BE. + StmtExecutor deferred = Mockito.mock(StmtExecutor.class); + + FlightSessionsManager sessionsManager = Mockito.mock(FlightSessionsManager.class); + Mockito.when(sessionsManager.getConnectContext(Mockito.anyString())).thenReturn(ctx); + + CallContext callContext = Mockito.mock(CallContext.class); + Mockito.when(callContext.peerIdentity()).thenReturn("token"); + + DorisFlightSqlProducer producer = new DorisFlightSqlProducer( + Location.forGrpcInsecure("127.0.0.1", 9090), sessionsManager); + try (MockedConstruction mocked = Mockito.mockConstruction( + FlightSqlConnectProcessor.class, (mock, context) -> { + // handleQuery plans + submits to BE and defers the coordinator (coordBase == coord), + // exactly as executeAndSendResult() does for an Arrow Flight external-table scan. + Mockito.doAnswer(invocation -> { + ctx.setReturnResultFromLocal(false); + ctx.addFlightSqlDeferredExecutor(deferred); + return null; + }).when(mock).handleQuery(Mockito.anyString()); + // The Arrow schema fetch fails after the coordinator was already deferred. + Mockito.doThrow(new RuntimeException("fetch arrow flight schema timeout")) + .when(mock).fetchArrowFlightSchema(Mockito.anyInt()); + })) { + CommandStatementQuery request = CommandStatementQuery.newBuilder().setQuery("select 1").build(); + FlightDescriptor descriptor = FlightDescriptor.command(new byte[0]); + + try { + producer.getFlightInfoStatement(request, callContext, descriptor); + Assert.fail("expected the schema fetch failure to propagate as a CallStatus"); + } catch (Throwable expected) { + // GetFlightInfo is expected to fail; the point of the test is what happens to the + // deferred coordinator, not the thrown status itself. + } + + // The deferred coordinator of the failed query is finalized on the error path instead of + // leaking until the next query / connection teardown. + Mockito.verify(deferred).finalizeArrowFlightQuery(); + + // It is also removed from the deferred list, so a later teardown does not finalize it + // again (no double-close, no retained reference). + ctx.closeFlightSqlDeferredExecutors(); + Mockito.verify(deferred, Mockito.times(1)).finalizeArrowFlightQuery(); + } finally { + producer.close(); + } + } } From 663491e395f66bc7c56279c9113539cbeb40d014 Mon Sep 17 00:00:00 2001 From: morningman Date: Fri, 17 Jul 2026 15:52:05 +0800 Subject: [PATCH 6/6] [fix](arrow-flight) Correct the deferral comment: it covers all BE-result Flight queries, not only external batch scans The gate that defers the coordinator close is coordBase == coord for any ARROW_FLIGHT_SQL query whose results are produced on the BE, so it captures internal-table and non-batch external SELECTs too, not only external-table batch-split scans. Deferral is only strictly required for the batch-split case (where the BE fetches splits from the FE during DoGet); the others are deferred by the same gate and merely hold their coordinator, query queue slot and query registration until the next query or connection teardown. The old comment implied only external batch scans were affected. Comment-only change. See #62259. --- .../java/org/apache/doris/qe/StmtExecutor.java | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index 159526eb6d8147..795a81c8a76ed7 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java @@ -1491,12 +1491,18 @@ public void executeAndSendResult(boolean isOutfileQuery, boolean isSendFields, if (context.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)) { Preconditions.checkState(!context.isReturnResultFromLocal()); profile.getSummaryProfile().setTempStartTime(); - // The BE keeps scanning and pulling results after GetFlightInfo returns; for an - // external-table scan in batch mode it lazily fetches splits from the FE during the - // later DoGet phase. Keep the coordinator (and the batch SplitSource its scan nodes - // hold) alive by deferring its close to ConnectContext, instead of closing it in the - // finally block below which would release the SplitSource too early. See #62259. - // (Point queries use a different coordBase and are not deferred.) + // 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);