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 listener = new StreamListener() { + @Override + public void onNext(Result val) { + // discard the placeholder prepared-statement result + } + + @Override + public void onError(Throwable t) { + errors.incrementAndGet(); + finished.countDown(); + } + + @Override + public void onCompleted() { + finished.countDown(); + } + }; + ActionCreatePreparedStatementRequest request = ActionCreatePreparedStatementRequest.newBuilder() + .setQuery("select * from t where id = " + i).build(); + producer.createPreparedStatement(request, callContext, listener); + Assert.assertTrue("createPreparedStatement #" + i + " did not finish in time", + finished.await(30, TimeUnit.SECONDS)); + } + + // Guard against a false pass: if a prepare failed before reaching the allocation, no buffer + // would be leaked and the memory assertion below could not detect a regression. + Assert.assertEquals("no createPreparedStatement call should fail", 0, errors.get()); + // Every temporary VectorSchemaRoot must have been closed, so the channel's Arrow allocator + // is back to zero. Reverting the fix leaves `rounds` ResultMeta buffers allocated here. + Assert.assertEquals("createPreparedStatement leaked off-heap memory in the channel allocator", + 0L, channel.getAllocatedMemory()); + } finally { + producer.close(); + } + } +}