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 ebc3347f07f27d..e2d197cb537ea3 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 @@ -356,14 +356,22 @@ public void createPreparedStatement(final ActionCreatePreparedStatementRequest r final ByteString handle = ByteString.copyFromUtf8(context.peerIdentity() + ":" + preparedStatementId); connectContext.addPreparedQuery(preparedStatementId, query); - VectorSchemaRoot emptyVectorSchemaRoot = new VectorSchemaRoot(new ArrayList<>(), new ArrayList<>()); - final Schema parameterSchema = emptyVectorSchemaRoot.getSchema(); + // Close the temporary VectorSchemaRoot after extracting its Schema, otherwise the + // off-heap buffers backing its vectors are leaked on every prepare (FE direct memory leak). + final Schema parameterSchema; + try (VectorSchemaRoot emptyVectorSchemaRoot = + new VectorSchemaRoot(new ArrayList<>(), new ArrayList<>())) { + parameterSchema = emptyVectorSchemaRoot.getSchema(); + } // TODO FE does not have the ability to convert root fragment output expr into arrow schema. // However, the metaData schema returned by createPreparedStatement is usually not used by the client, // but it cannot be empty, otherwise it will be mistaken by the client as an updata statement. // see: https://github.com/apache/arrow/issues/38911 - Schema metaData = connectContext.getFlightSqlChannel() - .createOneOneSchemaRoot("ResultMeta", "UNIMPLEMENTED").getSchema(); + final Schema metaData; + try (VectorSchemaRoot metaSchemaRoot = connectContext.getFlightSqlChannel() + .createOneOneSchemaRoot("ResultMeta", "UNIMPLEMENTED")) { + metaData = metaSchemaRoot.getSchema(); + } listener.onNext(new Result( Any.pack(buildCreatePreparedStatementResult(handle, parameterSchema, metaData)).toByteArray())); } catch (Exception e) { 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 8ce5f72b4000eb..4f54132876b0c2 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 @@ -156,6 +156,12 @@ public long resultNum() { return resultCache.size(); } + // Returns the off-heap memory (bytes) currently held by this channel's Arrow allocator. + // Exposed so tests can assert prepared-statement handling does not leak VectorSchemaRoot buffers. + public long getAllocatedMemory() { + return allocator.getAllocatedMemory(); + } + public void reset() { resultCache.invalidateAll(); } 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 new file mode 100644 index 00000000000000..3a4b0facace7e1 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/service/arrowflight/DorisFlightSqlProducerTest.java @@ -0,0 +1,145 @@ +// 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; + +import org.apache.doris.common.FeConstants; +import org.apache.doris.qe.ConnectContext; +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.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.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mockito; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +public class DorisFlightSqlProducerTest { + + private boolean prevRunningUnitTest; + + @Before + public void setUp() { + // FlightSqlConnectContext.init() only reaches Env when this is false; keep it true so the + // context can be built without a running FE. + prevRunningUnitTest = FeConstants.runningUnitTest; + FeConstants.runningUnitTest = true; + } + + @After + public void tearDown() { + FeConstants.runningUnitTest = prevRunningUnitTest; + } + + /** + * Regression test for the FE direct-memory leak in {@code createPreparedStatement} + * (issue apache/doris#65305, fixed by PR #65311). + * + *
Before the fix, each prepare allocated a {@link org.apache.arrow.vector.VectorSchemaRoot} + * via {@code FlightSqlChannel.createOneOneSchemaRoot("ResultMeta", ...)} and read only its + * {@code Schema}, never closing the root, leaking one off-heap {@code VarCharVector} buffer per + * prepare. Because Arrow Flight has no parameter binding, every client query triggers a fresh + * prepare, so the leak grows monotonically until {@code MaxDirectMemorySize} is exhausted. + * + *
This test drives {@code createPreparedStatement} against a real {@link FlightSqlChannel} many
+ * times and asserts the channel's Arrow allocator holds zero bytes afterwards. It goes red if the
+ * try-with-resources that closes the temporary roots is removed.
+ */
+ @Test
+ public void createPreparedStatementDoesNotLeakChannelAllocator() throws Exception {
+ // A real flight session context owns a real FlightSqlChannel (and thus a real Arrow allocator),
+ // so allocator bookkeeping is exercised for real instead of mocked away.
+ FlightSqlConnectContext connectContext = new FlightSqlConnectContext("test-peer-identity");
+ FlightSqlChannel channel = connectContext.getFlightSqlChannel();
+ Assert.assertEquals("channel allocator should start empty", 0L, channel.getAllocatedMemory());
+
+ FlightSessionsManager sessionsManager = new FlightSessionsManager() {
+ @Override
+ public ConnectContext getConnectContext(String peerIdentity) {
+ return connectContext;
+ }
+
+ @Override
+ public ConnectContext createConnectContext(String peerIdentity) {
+ return connectContext;
+ }
+
+ @Override
+ public void closeConnectContext(String peerIdentity) {
+ // not exercised by this test
+ }
+ };
+ DorisFlightSqlProducer producer =
+ new DorisFlightSqlProducer(Location.forGrpcInsecure("127.0.0.1", 9090), sessionsManager);
+
+ CallContext callContext = Mockito.mock(CallContext.class);
+ Mockito.when(callContext.peerIdentity()).thenReturn("test-peer-identity");
+
+ final int rounds = 100;
+ AtomicInteger errors = new AtomicInteger(0);
+ try {
+ // createPreparedStatement runs asynchronously and mutates a shared, non-thread-safe
+ // ConnectContext, so drive it serially: each prepare completes before the next is issued.
+ // The leak, if any, still accumulates on the single channel allocator across rounds.
+ for (int i = 0; i < rounds; i++) {
+ CountDownLatch finished = new CountDownLatch(1);
+ StreamListener