This is an automated email from the ASF dual-hosted git repository.
nzw921rx pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/seatunnel.git
The following commit(s) were added to refs/heads/dev by this push:
new 4a6f0fe829 [Fix][Connector-V2] Fix Arrow memory leak by correcting
close order in ArrowToSeatunnelRowReader (#11929)
4a6f0fe829 is described below
commit 4a6f0fe82964f430877bc8b689e2ed75883d9136
Author: David Zollo <[email protected]>
AuthorDate: Sat Oct 17 03:32:26 2026 -0400
[Fix][Connector-V2] Fix Arrow memory leak by correcting close order in
ArrowToSeatunnelRowReader (#11929)
Co-authored-by: Daniel <[email protected]>
Co-authored-by: DanielLeens <[email protected]>
---
.github/workflows/backend.yml | 14 +-
.../arrow/reader/ArrowToSeatunnelRowReader.java | 58 ++++-
.../reader/ArrowToSeatunnelRowReaderCloseTest.java | 272 +++++++++++++++++++++
3 files changed, 330 insertions(+), 14 deletions(-)
diff --git a/.github/workflows/backend.yml b/.github/workflows/backend.yml
index 6a26ce2f29..0a64afd075 100644
--- a/.github/workflows/backend.yml
+++ b/.github/workflows/backend.yml
@@ -771,7 +771,7 @@ jobs:
matrix:
java: [ '11', '17' ]
os: [ 'ubuntu-latest' ]
- timeout-minutes: 180
+ timeout-minutes: 300
steps:
- uses: actions/checkout@v4
- name: Set up JDK ${{ matrix.java }}
@@ -811,7 +811,7 @@ jobs:
matrix:
java: [ '11', '17' ]
os: [ 'ubuntu-latest' ]
- timeout-minutes: 180
+ timeout-minutes: 300
steps:
- uses: actions/checkout@v4
- name: Set up JDK ${{ matrix.java }}
@@ -891,7 +891,7 @@ jobs:
matrix:
java: [ '11', '17' ]
os: [ 'ubuntu-latest' ]
- timeout-minutes: 200
+ timeout-minutes: 300
steps:
- uses: actions/checkout@v4
- name: Set up JDK ${{ matrix.java }}
@@ -930,7 +930,7 @@ jobs:
matrix:
java: [ '11', '17' ]
os: [ 'ubuntu-latest' ]
- timeout-minutes: 180
+ timeout-minutes: 300
steps:
- uses: actions/checkout@v4
- name: Set up JDK ${{ matrix.java }}
@@ -969,7 +969,7 @@ jobs:
matrix:
java: [ '11', '17' ]
os: [ 'ubuntu-latest' ]
- timeout-minutes: 210
+ timeout-minutes: 300
steps:
- uses: actions/checkout@v4
- name: Set up JDK ${{ matrix.java }}
@@ -1008,7 +1008,7 @@ jobs:
matrix:
java: [ '11', '17' ]
os: [ 'ubuntu-latest' ]
- timeout-minutes: 210
+ timeout-minutes: 300
steps:
- uses: actions/checkout@v4
- name: Set up JDK ${{ matrix.java }}
@@ -1048,7 +1048,7 @@ jobs:
matrix:
java: [ '11', '17' ]
os: [ 'ubuntu-latest' ]
- timeout-minutes: 210
+ timeout-minutes: 300
steps:
- uses: actions/checkout@v4
- name: Set up JDK ${{ matrix.java }}
diff --git
a/seatunnel-connectors-v2/connector-common/src/main/java/org/apache/seatunnel/connectors/seatunnel/common/source/arrow/reader/ArrowToSeatunnelRowReader.java
b/seatunnel-connectors-v2/connector-common/src/main/java/org/apache/seatunnel/connectors/seatunnel/common/source/arrow/reader/ArrowToSeatunnelRowReader.java
index f7821ed5fd..327fe3da33 100644
---
a/seatunnel-connectors-v2/connector-common/src/main/java/org/apache/seatunnel/connectors/seatunnel/common/source/arrow/reader/ArrowToSeatunnelRowReader.java
+++
b/seatunnel-connectors-v2/connector-common/src/main/java/org/apache/seatunnel/connectors/seatunnel/common/source/arrow/reader/ArrowToSeatunnelRowReader.java
@@ -17,6 +17,7 @@
package org.apache.seatunnel.connectors.seatunnel.common.source.arrow.reader;
+import
org.apache.seatunnel.shade.com.google.common.annotations.VisibleForTesting;
import org.apache.seatunnel.shade.org.apache.arrow.memory.RootAllocator;
import org.apache.seatunnel.shade.org.apache.arrow.vector.FieldVector;
import org.apache.seatunnel.shade.org.apache.arrow.vector.VectorSchemaRoot;
@@ -82,6 +83,27 @@ public class ArrowToSeatunnelRowReader implements
AutoCloseable {
initArrowReader(byteArray);
}
+ /**
+ * Creates a reader with prebuilt Arrow resources.
+ *
+ * <p>The constructor keeps close-order regression tests focused on
resource ownership while
+ * preserving the production byte-array constructor as the normal entry
point. It skips
+ * initArrowReader(byte[]) on purpose, so it must never be promoted to a
production entry point:
+ * hasNext()/next() are only safe after readArrow() ran against real
stream state.
+ */
+ @VisibleForTesting
+ ArrowToSeatunnelRowReader(
+ ArrowStreamReader arrowStreamReader,
+ RootAllocator rootAllocator,
+ SeaTunnelRowType seaTunnelRowType) {
+ // Take ownership of the caller-supplied resources before any schema
processing that
+ // can throw, so close() stays able to release them even if
initialization fails.
+ this.arrowStreamReader = arrowStreamReader;
+ this.rootAllocator = rootAllocator;
+ this.seaTunnelDataTypes = seaTunnelRowType.getFieldTypes();
+ initFieldIndexMap(seaTunnelRowType);
+ }
+
private void initFieldIndexMap(SeaTunnelRowType seaTunnelRowType) {
for (int i = 0; i < seaTunnelRowType.getFieldNames().length; i++) {
fieldIndexMap.put(seaTunnelRowType.getFieldNames()[i], i);
@@ -319,18 +341,40 @@ public class ArrowToSeatunnelRowReader implements
AutoCloseable {
@Override
public void close() {
+ RuntimeException closeException = null;
try {
- if (root != null) {
- root.close();
- }
- if (rootAllocator != null) {
- rootAllocator.close();
- }
+ // ArrowStreamReader must be closed before RootAllocator.
+ // ArrowStreamReader internally closes VectorSchemaRoot and
releases all Arrow
+ // buffer allocations back to the allocator. If the allocator is
closed first,
+ // it detects unreleased memory and throws IllegalStateException
("Memory was
+ // leaked by query"), which is the root cause of issue #9863.
if (arrowStreamReader != null) {
arrowStreamReader.close();
}
} catch (IOException e) {
- throw new RuntimeException("failed to close arrow stream reader.",
e);
+ closeException = new RuntimeException("failed to close arrow
stream reader.", e);
+ } catch (RuntimeException e) {
+ // Arrow's ArrowReader.close() can also fail with unchecked
exceptions (for example
+ // from VectorSchemaRoot or dictionary vector cleanup). Route them
through the same
+ // aggregation as IOException; otherwise they would propagate
directly and any
+ // allocator close failure below would be dropped without a
suppressed link, losing
+ // the leak diagnostic this close path exists to preserve.
+ closeException = e;
+ } finally {
+ try {
+ if (rootAllocator != null) {
+ rootAllocator.close();
+ }
+ } catch (RuntimeException e) {
+ if (closeException == null) {
+ closeException = e;
+ } else {
+ closeException.addSuppressed(e);
+ }
+ }
+ }
+ if (closeException != null) {
+ throw closeException;
}
}
}
diff --git
a/seatunnel-connectors-v2/connector-common/src/test/java/org/apache/seatunnel/connectors/seatunnel/common/source/arrow/reader/ArrowToSeatunnelRowReaderCloseTest.java
b/seatunnel-connectors-v2/connector-common/src/test/java/org/apache/seatunnel/connectors/seatunnel/common/source/arrow/reader/ArrowToSeatunnelRowReaderCloseTest.java
new file mode 100644
index 0000000000..6202ac5691
--- /dev/null
+++
b/seatunnel-connectors-v2/connector-common/src/test/java/org/apache/seatunnel/connectors/seatunnel/common/source/arrow/reader/ArrowToSeatunnelRowReaderCloseTest.java
@@ -0,0 +1,272 @@
+/*
+ * 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.seatunnel.connectors.seatunnel.common.source.arrow.reader;
+
+import org.apache.seatunnel.shade.org.apache.arrow.memory.BufferAllocator;
+import org.apache.seatunnel.shade.org.apache.arrow.memory.RootAllocator;
+import org.apache.seatunnel.shade.org.apache.arrow.vector.FieldVector;
+import org.apache.seatunnel.shade.org.apache.arrow.vector.IntVector;
+import org.apache.seatunnel.shade.org.apache.arrow.vector.VarCharVector;
+import org.apache.seatunnel.shade.org.apache.arrow.vector.VectorSchemaRoot;
+import
org.apache.seatunnel.shade.org.apache.arrow.vector.dictionary.Dictionary;
+import
org.apache.seatunnel.shade.org.apache.arrow.vector.dictionary.DictionaryEncoder;
+import
org.apache.seatunnel.shade.org.apache.arrow.vector.dictionary.DictionaryProvider;
+import
org.apache.seatunnel.shade.org.apache.arrow.vector.ipc.ArrowStreamReader;
+import
org.apache.seatunnel.shade.org.apache.arrow.vector.ipc.ArrowStreamWriter;
+import org.apache.seatunnel.shade.org.apache.arrow.vector.types.pojo.ArrowType;
+import
org.apache.seatunnel.shade.org.apache.arrow.vector.types.pojo.DictionaryEncoding;
+
+import org.apache.seatunnel.api.table.type.BasicType;
+import org.apache.seatunnel.api.table.type.SeaTunnelDataType;
+import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.channels.Channels;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * Verifies that Arrow reader cleanup releases the allocator even when stream
reader close fails,
+ * and that the production read path can drain real Arrow streams without
tripping the allocator
+ * leak check from issue #9863.
+ */
+public class ArrowToSeatunnelRowReaderCloseTest {
+
+ /**
+ * Covers the normal close path and preserves the required
reader-before-allocator order.
+ *
+ * <p>This protects the Arrow allocator from being closed before the
stream reader has released
+ * vector buffers.
+ */
+ @Test
+ public void testCloseReleasesAllocatorWhenStreamReaderCloseSucceeds()
throws Exception {
+ ArrowStreamReader arrowStreamReader =
Mockito.mock(ArrowStreamReader.class);
+ RootAllocator rootAllocator = Mockito.mock(RootAllocator.class);
+ ArrowToSeatunnelRowReader reader =
+ new ArrowToSeatunnelRowReader(
+ arrowStreamReader, rootAllocator,
createSeatunnelRowType());
+
+ reader.close();
+
+ InOrder inOrder = Mockito.inOrder(arrowStreamReader, rootAllocator);
+ inOrder.verify(arrowStreamReader).close();
+ inOrder.verify(rootAllocator).close();
+ }
+
+ /**
+ * Covers the regression where an ArrowStreamReader close failure skipped
allocator cleanup.
+ *
+ * <p>The reader close failure is still propagated, but the allocator must
always receive its
+ * close call.
+ */
+ @Test
+ public void testCloseReleasesAllocatorWhenStreamReaderCloseFails() throws
Exception {
+ IOException closeFailure = new IOException("close failure");
+ ArrowStreamReader arrowStreamReader =
Mockito.mock(ArrowStreamReader.class);
+ Mockito.doThrow(closeFailure).when(arrowStreamReader).close();
+ RootAllocator rootAllocator = Mockito.mock(RootAllocator.class);
+ ArrowToSeatunnelRowReader reader =
+ new ArrowToSeatunnelRowReader(
+ arrowStreamReader, rootAllocator,
createSeatunnelRowType());
+
+ RuntimeException exception =
Assertions.assertThrows(RuntimeException.class, reader::close);
+
+ Assertions.assertSame(closeFailure, exception.getCause());
+ InOrder inOrder = Mockito.inOrder(arrowStreamReader, rootAllocator);
+ inOrder.verify(arrowStreamReader).close();
+ inOrder.verify(rootAllocator).close();
+ }
+
+ /**
+ * Covers unchecked reader close failures so the allocator failure is
never silently dropped.
+ *
+ * <p>An unchecked exception from ArrowStreamReader.close() must become
the primary close
+ * failure and any allocator close failure must be attached as a
suppressed exception, otherwise
+ * the allocator diagnostic for a real leak is lost during error handling.
+ */
+ @Test
+ public void
testCloseAggregatesUncheckedReaderFailureWithAllocatorFailure() throws
Exception {
+ RuntimeException readerFailure = new IllegalStateException("reader
close failure");
+ IllegalStateException allocatorFailure = new
IllegalStateException("Memory was leaked");
+ ArrowStreamReader arrowStreamReader =
Mockito.mock(ArrowStreamReader.class);
+ Mockito.doThrow(readerFailure).when(arrowStreamReader).close();
+ RootAllocator rootAllocator = Mockito.mock(RootAllocator.class);
+ Mockito.doThrow(allocatorFailure).when(rootAllocator).close();
+ ArrowToSeatunnelRowReader reader =
+ new ArrowToSeatunnelRowReader(
+ arrowStreamReader, rootAllocator,
createSeatunnelRowType());
+
+ RuntimeException exception =
Assertions.assertThrows(RuntimeException.class, reader::close);
+
+ Assertions.assertSame(readerFailure, exception);
+ Assertions.assertArrayEquals(new Throwable[] {allocatorFailure},
exception.getSuppressed());
+ InOrder inOrder = Mockito.inOrder(arrowStreamReader, rootAllocator);
+ inOrder.verify(arrowStreamReader).close();
+ inOrder.verify(rootAllocator).close();
+ }
+
+ /**
+ * Round-trip regression for the plain read path with a real allocator.
+ *
+ * <p>The reader is built through the production byte-array constructor,
so readArrow() runs the
+ * same load, convert, and close sequence Doris and StarRocks sources use.
Any close-order
+ * regression that leaves allocations behind surfaces here as an
IllegalStateException from
+ * RootAllocator.close().
+ */
+ @Test
+ public void testReadArrowDrainsRealStreamAndClosesWithoutLeak() throws
Exception {
+ byte[] payload = writePlainIntStream();
+
+ ArrowToSeatunnelRowReader reader =
+ new ArrowToSeatunnelRowReader(payload,
createSeatunnelRowType());
+
+ Assertions.assertDoesNotThrow(reader::readArrow);
+ Assertions.assertEquals(3, reader.getReadRowCount());
+ Assertions.assertEquals(Arrays.asList(1, 2, 3), drainIds(reader));
+ }
+
+ /**
+ * Round-trip regression for the exact issue #9863 failure signature.
+ *
+ * <p>Dictionary batches are allocated by ArrowStreamReader outside the
VectorSchemaRoot, so
+ * closing the allocator before the stream reader (the pre-fix order)
fails with "Memory was
+ * leaked by query" even for a fully drained stream. This test writes a
dictionary-encoded
+ * column and must fail if the reader-before-allocator close order is ever
reverted. The encoded
+ * column is intentionally absent from the SeaTunnel schema; the reader
skips unknown fields,
+ * keeping the assertion focused on resource release.
+ */
+ @Test
+ public void testReadArrowReleasesDictionaryAllocationsOnClose() throws
Exception {
+ byte[] payload = writeDictionaryEncodedStream();
+
+ ArrowToSeatunnelRowReader reader =
+ new ArrowToSeatunnelRowReader(payload,
createSeatunnelRowType());
+
+ Assertions.assertDoesNotThrow(reader::readArrow);
+ Assertions.assertEquals(3, reader.getReadRowCount());
+ Assertions.assertEquals(Arrays.asList(1, 2, 3), drainIds(reader));
+ }
+
+ /**
+ * Writes a single-batch Arrow stream holding one int column named "id"
with values 1, 2, 3.
+ *
+ * <p>The writer side uses its own allocator; leaked writer allocations
would fail the test when
+ * the try-with-resources closes that allocator.
+ */
+ private byte[] writePlainIntStream() throws IOException {
+ ByteArrayOutputStream out = new ByteArrayOutputStream();
+ try (BufferAllocator writerAllocator = new
RootAllocator(Integer.MAX_VALUE);
+ IntVector idVector = new IntVector("id", writerAllocator)) {
+ fillIds(idVector);
+ try (VectorSchemaRoot writerRoot = VectorSchemaRoot.of(idVector)) {
+ writerRoot.setRowCount(3);
+ try (ArrowStreamWriter writer =
+ new ArrowStreamWriter(writerRoot, null,
Channels.newChannel(out))) {
+ writer.start();
+ writer.writeBatch();
+ writer.end();
+ }
+ }
+ }
+ return out.toByteArray();
+ }
+
+ /**
+ * Writes a single-batch Arrow stream with a plain "id" column plus a
dictionary-encoded "tag"
+ * column, forcing the reader to hold dictionary allocations outside the
root.
+ */
+ private byte[] writeDictionaryEncodedStream() throws IOException {
+ ByteArrayOutputStream out = new ByteArrayOutputStream();
+ try (BufferAllocator writerAllocator = new
RootAllocator(Integer.MAX_VALUE);
+ VarCharVector dictionaryValues =
+ new VarCharVector("tag_dictionary", writerAllocator);
+ VarCharVector rawTags = new VarCharVector("tag",
writerAllocator);
+ IntVector idVector = new IntVector("id", writerAllocator)) {
+ dictionaryValues.allocateNew(2);
+ dictionaryValues.set(0, "apple".getBytes(StandardCharsets.UTF_8));
+ dictionaryValues.set(1, "banana".getBytes(StandardCharsets.UTF_8));
+ dictionaryValues.setValueCount(2);
+ Dictionary dictionary =
+ new Dictionary(
+ dictionaryValues,
+ new DictionaryEncoding(1L, false, new
ArrowType.Int(32, true)));
+
+ rawTags.allocateNew(3);
+ rawTags.set(0, "apple".getBytes(StandardCharsets.UTF_8));
+ rawTags.set(1, "banana".getBytes(StandardCharsets.UTF_8));
+ rawTags.set(2, "apple".getBytes(StandardCharsets.UTF_8));
+ rawTags.setValueCount(3);
+
+ fillIds(idVector);
+
+ try (FieldVector encodedTags =
+ (FieldVector) DictionaryEncoder.encode(rawTags,
dictionary)) {
+ DictionaryProvider.MapDictionaryProvider provider =
+ new
DictionaryProvider.MapDictionaryProvider(dictionary);
+ try (VectorSchemaRoot writerRoot =
VectorSchemaRoot.of(idVector, encodedTags)) {
+ writerRoot.setRowCount(3);
+ try (ArrowStreamWriter writer =
+ new ArrowStreamWriter(writerRoot, provider,
Channels.newChannel(out))) {
+ writer.start();
+ writer.writeBatch();
+ writer.end();
+ }
+ }
+ }
+ }
+ return out.toByteArray();
+ }
+
+ /** Populates the shared "id" test column with values 1, 2, 3. */
+ private void fillIds(IntVector idVector) {
+ idVector.allocateNew(3);
+ idVector.set(0, 1);
+ idVector.set(1, 2);
+ idVector.set(2, 3);
+ idVector.setValueCount(3);
+ }
+
+ /** Drains all rows from the reader and returns the "id" column values in
read order. */
+ private List<Integer> drainIds(ArrowToSeatunnelRowReader reader) {
+ List<Integer> ids = new ArrayList<>();
+ while (reader.hasNext()) {
+ ids.add((Integer) reader.next().getField(0));
+ }
+ return ids;
+ }
+
+ /**
+ * Builds the minimal SeaTunnel schema needed to initialize the reader
resource owner.
+ *
+ * <p>Only the "id" column is declared on purpose: the dictionary-encoded
"tag" column must stay
+ * out of the schema so the round-trip test exercises the unknown-field
skip path while still
+ * forcing dictionary allocations on the reader side.
+ */
+ private SeaTunnelRowType createSeatunnelRowType() {
+ return new SeaTunnelRowType(
+ new String[] {"id"}, new SeaTunnelDataType[]
{BasicType.INT_TYPE});
+ }
+}
