From 1fcf06aec97eeda10419b1e274c9b2e0008b3682 Mon Sep 17 00:00:00 2001 From: lisen Date: Wed, 16 Sep 2026 09:45:09 +0800 Subject: [PATCH] [client] Close Arrow read context after LIMIT log scan LimitBatchScanner dropped LogRecordReadContext after parsing Arrow batches, leaking the allocator on every log-table LIMIT scan. Co-authored-by: Cursor --- .../scanner/batch/LimitBatchScanner.java | 22 ++- .../scanner/batch/LimitBatchScannerTest.java | 142 ++++++++++++++++++ 2 files changed, 156 insertions(+), 8 deletions(-) create mode 100644 fluss-client/src/test/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScannerTest.java diff --git a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java index 58e39241dc..f60b8e9b0e 100644 --- a/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java +++ b/fluss-client/src/main/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScanner.java @@ -193,16 +193,22 @@ private List parseLimitScanResponse(LimitScanResponse limitScanResp null, schemaGetter, chunkedFactory); - LogRecords records = MemoryLogRecords.pointToByteBuffer(recordsBuffer); - for (LogRecordBatch logRecordBatch : records.batches()) { - // A batch of log record maybe little more than limit, thus we need slice the - // last limit number. - try (CloseableIterator logRecordIterator = - logRecordBatch.records(readContext)) { - while (logRecordIterator.hasNext()) { - scanRows.add(maybeProject(logRecordIterator.next().getRow())); + try { + LogRecords records = MemoryLogRecords.pointToByteBuffer(recordsBuffer); + for (LogRecordBatch logRecordBatch : records.batches()) { + // A batch of log record maybe little more than limit, thus we need slice the + // last limit number. + try (CloseableIterator logRecordIterator = + logRecordBatch.records(readContext)) { + while (logRecordIterator.hasNext()) { + scanRows.add(maybeProject(logRecordIterator.next().getRow())); + } } } + } finally { + // Release the Arrow allocator created for this parse; closing the + // record iterator alone does not close LogRecordReadContext. + readContext.close(); } } if (scanRows.size() > limit) { diff --git a/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScannerTest.java b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScannerTest.java new file mode 100644 index 0000000000..3a910decf7 --- /dev/null +++ b/fluss-client/src/test/java/org/apache/fluss/client/table/scanner/batch/LimitBatchScannerTest.java @@ -0,0 +1,142 @@ +/* + * 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.fluss.client.table.scanner.batch; + +import org.apache.fluss.client.metadata.MetadataUpdater; +import org.apache.fluss.cluster.Cluster; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.metadata.LogFormat; +import org.apache.fluss.metadata.SchemaGetter; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.record.MemoryLogRecords; +import org.apache.fluss.record.TestingSchemaGetter; +import org.apache.fluss.row.InternalRow; +import org.apache.fluss.rpc.TestingTabletGatewayService; +import org.apache.fluss.rpc.gateway.TabletServerGateway; +import org.apache.fluss.rpc.messages.LimitScanRequest; +import org.apache.fluss.rpc.messages.LimitScanResponse; +import org.apache.fluss.utils.CloseableIterator; + +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.CompletableFuture; + +import static org.apache.fluss.record.TestData.DATA1_ROW_TYPE; +import static org.apache.fluss.record.TestData.DATA1_SCHEMA; +import static org.apache.fluss.record.TestData.DATA1_TABLE_ID; +import static org.apache.fluss.record.TestData.DATA1_TABLE_INFO; +import static org.apache.fluss.record.TestData.DEFAULT_MAGIC; +import static org.apache.fluss.record.TestData.DEFAULT_SCHEMA_ID; +import static org.apache.fluss.testutils.DataTestUtils.createRecordsWithoutBaseLogOffset; +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests that {@link LimitBatchScanner} closes {@link org.apache.fluss.record.LogRecordReadContext} + * after parsing Arrow log records. + */ +class LimitBatchScannerTest { + + private static final TableBucket BUCKET_0 = new TableBucket(DATA1_TABLE_ID, 0); + private static final Duration POLL_TIMEOUT = Duration.ofSeconds(5); + private static final SchemaGetter SCHEMA_GETTER = + new TestingSchemaGetter(DEFAULT_SCHEMA_ID, DATA1_SCHEMA); + + @Test + void pollBatchClosesArrowReadContextSoScannerCloseSucceeds() throws Exception { + MemoryLogRecords records = + createRecordsWithoutBaseLogOffset( + DATA1_ROW_TYPE, + DEFAULT_SCHEMA_ID, + 0L, + 1000L, + DEFAULT_MAGIC, + Arrays.asList(new Object[] {1, "a"}, new Object[] {2, "b"}), + LogFormat.ARROW); + byte[] recordBytes = new byte[records.sizeInBytes()]; + records.getMemorySegment().get(records.getPosition(), recordBytes); + + LimitScanResponse response = + new LimitScanResponse().setIsLogTable(true).setRecords(recordBytes); + LimitGateway gateway = new LimitGateway(response); + + LimitBatchScanner scanner = + new LimitBatchScanner( + DATA1_TABLE_INFO, + BUCKET_0, + SCHEMA_GETTER, + new TestMetadataUpdater(gateway), + null, + 10); + try { + CloseableIterator batch = scanner.pollBatch(POLL_TIMEOUT); + assertThat(batch).isNotNull(); + List rows = new ArrayList<>(); + while (batch.hasNext()) { + rows.add(batch.next()); + } + batch.close(); + assertThat(rows).hasSize(2); + assertThat(rows.get(0).getInt(0)).isEqualTo(1); + assertThat(rows.get(1).getInt(0)).isEqualTo(2); + } finally { + scanner.close(); + } + } + + private static final class LimitGateway extends TestingTabletGatewayService { + private final LimitScanResponse response; + + private LimitGateway(LimitScanResponse response) { + this.response = response; + } + + @Override + public CompletableFuture limitScan(LimitScanRequest request) { + return CompletableFuture.completedFuture(response); + } + } + + private static final class TestMetadataUpdater extends MetadataUpdater { + private final TabletServerGateway gateway; + + TestMetadataUpdater(TabletServerGateway gateway) { + super(null, new Configuration(), Cluster.empty()); + this.gateway = gateway; + } + + @Override + public void checkAndUpdateMetadata(TablePath tablePath, TableBucket tableBucket) {} + + @Override + public int leaderFor(TablePath tablePath, TableBucket tableBucket) { + return 0; + } + + @Override + public @Nullable TabletServerGateway newTabletServerClientForNode(int serverId) { + return gateway; + } + } +}