This is an automated email from the ASF dual-hosted git repository.
morningman pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/branch-4.0 by this push:
new 058dc08e620 [branch-4.0] pick some arrow memory leak fix to branch-4.0
(#66462)
058dc08e620 is described below
commit 058dc08e620942eefe729f359193a73148f20419
Author: yiguolei <[email protected]>
AuthorDate: Wed Aug 5 15:15:25 2026 +0800
[branch-4.0] pick some arrow memory leak fix to branch-4.0 (#66462)
### What problem does this PR solve?
pick #64799 #66437
---
.../java/org/apache/doris/qe/ConnectContext.java | 34 ++++++
.../java/org/apache/doris/qe/StmtExecutor.java | 44 ++++++-
.../arrowflight/DorisFlightSqlProducer.java | 12 ++
.../arrowflight/FlightSqlConnectProcessor.java | 21 +++-
.../arrowflight/results/FlightSqlChannel.java | 1 +
.../results/FlightSqlResultCacheEntry.java | 2 +-
.../sessions/FlightSqlConnectContext.java | 3 -
.../sessions/FlightSqlConnectPoolMgr.java | 20 ++++
.../org/apache/doris/qe/ConnectContextTest.java | 58 +++++++++
.../java/org/apache/doris/qe/StmtExecutorTest.java | 36 ++++++
.../arrowflight/DorisFlightSqlProducerTest.java | 65 ++++++++++
.../sessions/FlightSqlConnectPoolMgrTest.java | 69 +++++++++++
.../test_iceberg_arrow_flight_split_source.groovy | 131 +++++++++++++++++++++
13 files changed, 485 insertions(+), 11 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 de2e7c220b1..f9114313cd8 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
@@ -941,6 +941,40 @@ public class ConnectContext {
return plSqlOperation;
}
+ // 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.
*/
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 f97d02423a4..0037fc28aa5 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
@@ -187,6 +187,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;
@@ -944,6 +948,23 @@ public class StmtExecutor {
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;
@@ -1372,6 +1393,22 @@ public class StmtExecutor {
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;
}
@@ -1487,7 +1524,12 @@ public class StmtExecutor {
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 e2d197cb537..e64690a54cb 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 @@ public class DorisFlightSqlProducer implements
FlightSqlProducer, AutoCloseable
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 @@ public class DorisFlightSqlProducer implements
FlightSqlProducer, AutoCloseable
}
}
} 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/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java
b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/FlightSqlConnectProcessor.java
index 2e5dfb3259e..9c59e8db606 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
@@ -145,11 +145,11 @@ public class FlightSqlConnectProcessor extends
ConnectProcessor implements AutoC
}
endpointLoc.setResultPublicAccessAddr(resultPublicAccessAddr);
if (pResult.hasSchema() && pResult.getSchema().size() > 0) {
- RootAllocator rootAllocator = new
RootAllocator(Integer.MAX_VALUE);
- ArrowStreamReader arrowStreamReader = new
ArrowStreamReader(
- new
ByteArrayInputStream(pResult.getSchema().toByteArray()), rootAllocator);
- try {
+ try (RootAllocator rootAllocator = new
RootAllocator(Integer.MAX_VALUE);
+ ArrowStreamReader arrowStreamReader = new
ArrowStreamReader(
+ new
ByteArrayInputStream(pResult.getSchema().toByteArray()), rootAllocator)) {
Schema schema;
+ // SchemaRoot belongs to ArrowStreamReader, it will be
released when ArrowStreamReader is closed
VectorSchemaRoot root =
arrowStreamReader.getVectorSchemaRoot();
List<FieldVector> fieldVectors =
root.getFieldVectors();
if (fieldVectors.size() != resultOutputExprs.size()) {
@@ -195,11 +195,20 @@ public class FlightSqlConnectProcessor extends
ConnectProcessor implements AutoC
@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/results/FlightSqlChannel.java
b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/results/FlightSqlChannel.java
index 4f54132876b..2781994dfa1 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/results/FlightSqlChannel.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/results/FlightSqlChannel.java
@@ -168,6 +168,7 @@ public class FlightSqlChannel {
public void close() {
reset();
+ allocator.close();
}
private static class ResultRemovalListener implements
RemovalListener<String, FlightSqlResultCacheEntry> {
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/results/FlightSqlResultCacheEntry.java
b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/results/FlightSqlResultCacheEntry.java
index 12ce04ca8ed..6edc868b0dd 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/results/FlightSqlResultCacheEntry.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/results/FlightSqlResultCacheEntry.java
@@ -42,7 +42,7 @@ public final class FlightSqlResultCacheEntry implements
AutoCloseable {
@Override
public void close() throws Exception {
- vectorSchemaRoot.clear();
+ vectorSchemaRoot.close();
}
@Override
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectContext.java
b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectContext.java
index 75f1c0ee4bf..ceddfaa5639 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectContext.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectContext.java
@@ -62,9 +62,6 @@ public class FlightSqlConnectContext extends ConnectContext {
@Override
protected void closeChannel() {
- if (flightSqlChannel != null) {
- flightSqlChannel.close();
- }
connectScheduler.getFlightSqlConnectPoolMgr().unregisterConnection(this);
}
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 3002116b386..04982a7fedc 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
@@ -20,6 +20,7 @@ package org.apache.doris.service.arrowflight.sessions;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.ConnectContext.ConnectType;
import org.apache.doris.qe.ConnectPoolMgr;
+import org.apache.doris.service.arrowflight.results.FlightSqlChannel;
import com.google.common.collect.Maps;
import org.apache.logging.log4j.LogManager;
@@ -55,6 +56,25 @@ public class FlightSqlConnectPoolMgr extends ConnectPoolMgr {
@Override
public void unregisterConnection(ConnectContext ctx) {
+ // All Flight SQL session teardown paths (idle/query timeout, bearer
token expiry, and
+ // explicit CloseSession) reach here. Release channel-cached Arrow
results before removing
+ // the context from the pool.
+ FlightSqlChannel flightSqlChannel = ctx.getFlightSqlChannel();
+ if (flightSqlChannel != null) {
+ try {
+ flightSqlChannel.close();
+ } catch (Throwable t) {
+ // RootAllocator.close() marks the allocator closed before it
reports outstanding
+ // bytes. The error is actionable, but session teardown must
still release the
+ // coordinator, transaction and pool/token bookkeeping below.
+ LOG.warn("failed to close Flight SQL channel while
unregistering connection {}, peer identity {}",
+ ctx.getConnectionId(), ctx.getPeerIdentity(), t);
+ }
+ }
+ // 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).
+ ctx.closeFlightSqlDeferredExecutors();
ctx.closeTxn();
if (connectionMap.remove(ctx.getConnectionId()) != null) {
numberConnection.decrementAndGet();
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 f02e22ebfa1..b626f6acf8a 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
@@ -871,4 +871,62 @@ public class ConnectContextTest {
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 8635e1c0082..296a8df7cb8 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
@@ -26,6 +26,8 @@ import org.apache.doris.mysql.MysqlChannel;
import org.apache.doris.mysql.MysqlSerializer;
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;
@@ -66,6 +68,40 @@ public class StmtExecutorTest extends TestWithFeService {
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/DorisFlightSqlProducerTest.java
b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducerTest.java
index 3a4b0facace..eede1268851 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 @@ package org.apache.doris.service.arrowflight;
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 class DorisFlightSqlProducerTest {
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<FlightSqlConnectProcessor> 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();
+ }
+ }
}
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 00000000000..0513221569d
--- /dev/null
+++
b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/sessions/FlightSqlConnectPoolMgrTest.java
@@ -0,0 +1,69 @@
+// 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.apache.doris.service.arrowflight.results.FlightSqlChannel;
+
+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);
+ FlightSqlChannel channel = Mockito.mock(FlightSqlChannel.class);
+ Mockito.when(ctx.getFlightSqlChannel()).thenReturn(channel);
+
+ 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(channel).close();
+ 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);
+ FlightSqlChannel channel = Mockito.mock(FlightSqlChannel.class);
+ Mockito.when(ctx.getFlightSqlChannel()).thenReturn(channel);
+ 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(channel).close();
+ Mockito.verify(ctx).closeFlightSqlDeferredExecutors();
+ Assert.assertNull(poolMgr.getConnectionMap().get(7));
+ }
+}
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 00000000000..e721d0e9d8e
--- /dev/null
+++
b/regression-test/suites/external_table_p0/iceberg/test_iceberg_arrow_flight_split_source.groovy
@@ -0,0 +1,131 @@
+// 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.
+
+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
+// (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
+ }
+
+ // 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")
+ 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"
+ );"""
+
+ 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.
+ 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 = flightSql """ 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 = 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 = flightSql """ select * from ${table} limit 10 """
+ assert flightLimited.size() > 0 && flightLimited.size() <= 10 :
"unexpected row count: ${flightLimited.size()}"
+ } finally {
+ // 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}"""
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
