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..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 @@ -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(); + } } } 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..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 @@ -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(); 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(); 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/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(); + } + } } 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)); + } +} 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..e721d0e9d8eab6 --- /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}""" + } +}