From 8ccb73e7a4604329a78eea4fa6f82a3e599075c9 Mon Sep 17 00:00:00 2001 From: Sandesh Kumar Date: Sun, 12 Jul 2026 07:28:05 +0000 Subject: [PATCH 1/2] Fix C Data import leaking the whole array on a mid-import allocator OOM ReferenceCountedArrowArray.unsafeAssociateAllocation retained the imported array's shared reference count *before* calling wrapForeignAllocation, which can throw OutOfMemoryException when the tracking (importing) allocator is over its limit. The retain was not rolled back on that throw, so the reference count stayed permanently elevated: release() never reached zero, the producer's C Data release callback was never invoked, and the entire imported array's buffers leaked in the producer's memory (invisible to the JVM heap). This strands whole batches when importing a wide/variable-width array into a bounded allocator that fills part-way through the array. Retain only after wrapForeignAllocation succeeds, so each retain() is paired with exactly one ForeignAllocation.release0(). If it throws, no ForeignAllocation was created, so skipping the retain is correct: the initial reference held by ArrayImporter.importArray (released in its finally) then drives the count to zero and fires the release callback, freeing the array. Adds ImportOutOfMemoryTest, which exports a multi-buffer batch from a dedicated producer allocator and imports it into an allocator too small to hold it; it asserts the import throws an allocator OOM and that the producer drains to zero (release callback fired). The test fails on the previous retain-before-wrap code (producer strands ~1 MB) and passes with this change. --- .../arrow/c/ReferenceCountedArrowArray.java | 19 ++- .../apache/arrow/c/ImportOutOfMemoryTest.java | 142 ++++++++++++++++++ 2 files changed, 154 insertions(+), 7 deletions(-) create mode 100644 c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java diff --git a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java index cf50f9417b..324e8671ec 100644 --- a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java +++ b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java @@ -64,13 +64,18 @@ void release() { */ ArrowBuf unsafeAssociateAllocation( BufferAllocator trackingAllocator, long capacity, long memoryAddress) { + // Retain AFTER wrap: wrapForeignAllocation throws OutOfMemoryException when the allocator is + // over its limit, and a retain() before that throw would leave the count elevated with no + // matching release0(), preventing the array's release callback from ever firing. + ArrowBuf buf = + trackingAllocator.wrapForeignAllocation( + new ForeignAllocation(capacity, memoryAddress) { + @Override + protected void release0() { + ReferenceCountedArrowArray.this.release(); + } + }); retain(); - return trackingAllocator.wrapForeignAllocation( - new ForeignAllocation(capacity, memoryAddress) { - @Override - protected void release0() { - ReferenceCountedArrowArray.this.release(); - } - }); + return buf; } } diff --git a/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java new file mode 100644 index 0000000000..51e4d4ae71 --- /dev/null +++ b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.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.arrow.c; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.OutOfMemoryException; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Regression test: a mid-import {@link OutOfMemoryException} must not leak the imported array. + * + *

A "producer" allocator owns the exported batch; if the C Data release callback fires, the + * producer drains to zero. A too-small consumer allocator forces an OOM part-way through the + * import. The test asserts the producer drains, confirming the release callback fired despite + * the failure. + */ +final class ImportOutOfMemoryTest { + private static final int ROWS = 1024; + private static final int VALUE_BYTES = 256; + private static final int COLUMNS = 4; + // Far smaller than the exported batch, so the import OOMs part-way through the buffers. + private static final long TINY_LIMIT = 16 * 1024; + + private RootAllocator root; + + @BeforeEach + public void setUp() { + root = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + public void tearDown() { + root.close(); + } + + @Test + public void importOomDoesNotLeakExportedArray() { + // "producer" owns only the exported batch buffers; the C Data struct containers live on a + // separate allocator (they are consumed/closed by import, which would otherwise muddy the + // producer's balance). So producer draining to zero is an exact signal that the array's release + // callback fired. + BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); + BufferAllocator structs = root.newChildAllocator("structs", 0, Long.MAX_VALUE); + try (ArrowArray array = ArrowArray.allocateNew(structs); + ArrowSchema schema = ArrowSchema.allocateNew(structs)) { + exportBatch(producer, array, schema); + assertTrue( + producer.getAllocatedMemory() > 0, "producer holds the exported batch before import"); + + // A consumer allocator far too small to hold the batch: the import throws part-way through. + BufferAllocator consumer = root.newChildAllocator("consumer", 0, TINY_LIMIT); + try (CDataDictionaryProvider provider = new CDataDictionaryProvider()) { + Schema importSchema = Data.importSchema(consumer, schema, provider); + try (VectorSchemaRoot importRoot = VectorSchemaRoot.create(importSchema, consumer)) { + Exception thrown = + assertThrows( + Exception.class, + () -> Data.importIntoVectorSchemaRoot(consumer, array, importRoot, provider)); + assertTrue( + hasOutOfMemoryCause(thrown), + "mid-import failure must be an allocator OOM: " + thrown); + } + } + consumer.close(); + + // The array's release callback must have fired despite the mid-import OOM, freeing the whole + // exported batch. On the unfixed retain-before-wrap code the batch is stranded instead. + assertEquals( + 0L, + producer.getAllocatedMemory(), + "import OOM leaked the exported batch (producer not drained)"); + } + // Closing these (and the RootAllocator in tearDown) additionally asserts nothing leaked at all. + producer.close(); + structs.close(); + } + + /** True if {@code t} is, or is caused by, an Arrow {@link OutOfMemoryException}. */ + private static boolean hasOutOfMemoryCause(Throwable t) { + for (Throwable cause = t; cause != null; cause = cause.getCause()) { + if (cause instanceof OutOfMemoryException) { + return true; + } + } + return false; + } + + /** + * Builds a wide multi-column VarChar batch on {@code alloc} and exports it into the C structs. + */ + private void exportBatch(BufferAllocator alloc, ArrowArray array, ArrowSchema schema) { + byte[] value = new byte[VALUE_BYTES]; + for (int i = 0; i < value.length; i++) { + value[i] = (byte) 'x'; + } + List vectors = new ArrayList<>(COLUMNS); + for (int c = 0; c < COLUMNS; c++) { + VarCharVector vector = new VarCharVector("col" + c, alloc); + vector.allocateNew((long) ROWS * VALUE_BYTES, ROWS); + for (int r = 0; r < ROWS; r++) { + vector.setSafe(r, value); + } + vector.setValueCount(ROWS); + vectors.add(vector); + } + try (VectorSchemaRoot source = new VectorSchemaRoot(vectors)) { + long total = 0; + for (FieldVector vector : source.getFieldVectors()) { + total += vector.getBufferSize(); + } + assertTrue(total > TINY_LIMIT, "test setup: batch must exceed the consumer limit"); + Data.exportVectorSchemaRoot(alloc, source, null, array, schema); + } + } +} From 9a7c5c38323e8cc9fcb1e0fa6be8c6f7673386b6 Mon Sep 17 00:00:00 2001 From: Sandesh Kumar Date: Tue, 21 Jul 2026 18:10:23 +0000 Subject: [PATCH 2/2] Address review: clarify retain ordering, wrap test allocators in try-with-resources Expand the comment in unsafeAssociateAllocation to explain why retain-after is correct for both throw paths in BaseAllocator.wrapForeignAllocation: the OOM path throws before the ForeignAllocation is associated (release0 is not called), and the unexpected-error path has BaseAllocator's own catch call release0 before rethrowing (a retain-before + catch pattern would double-release on that path). Wrap all child allocators in ImportOutOfMemoryTest in try-with-resources so an early assertion failure does not leave them open for tearDown to find. Signed-off-by: Sandesh Kumar --- .../arrow/c/ReferenceCountedArrowArray.java | 6 +- .../apache/arrow/c/ImportOutOfMemoryTest.java | 61 +++++++++---------- 2 files changed, 32 insertions(+), 35 deletions(-) diff --git a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java index 324e8671ec..f51fb25105 100644 --- a/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java +++ b/c/src/main/java/org/apache/arrow/c/ReferenceCountedArrowArray.java @@ -64,9 +64,9 @@ void release() { */ ArrowBuf unsafeAssociateAllocation( BufferAllocator trackingAllocator, long capacity, long memoryAddress) { - // Retain AFTER wrap: wrapForeignAllocation throws OutOfMemoryException when the allocator is - // over its limit, and a retain() before that throw would leave the count elevated with no - // matching release0(), preventing the array's release callback from ever firing. + // Retain only after wrapForeignAllocation succeeds. On the allocator-limit OOM path, + // wrapForeignAllocation throws before the ForeignAllocation is associated, so release0() + // is not called; retaining first would leave the count elevated with no matching release0(). ArrowBuf buf = trackingAllocator.wrapForeignAllocation( new ForeignAllocation(capacity, memoryAddress) { diff --git a/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java index 51e4d4ae71..7c099f2ef0 100644 --- a/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java +++ b/c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java @@ -38,8 +38,8 @@ * *

A "producer" allocator owns the exported batch; if the C Data release callback fires, the * producer drains to zero. A too-small consumer allocator forces an OOM part-way through the - * import. The test asserts the producer drains, confirming the release callback fired despite - * the failure. + * import. The test asserts the producer drains, confirming the release callback fired despite the + * failure. */ final class ImportOutOfMemoryTest { private static final int ROWS = 1024; @@ -66,40 +66,37 @@ public void importOomDoesNotLeakExportedArray() { // separate allocator (they are consumed/closed by import, which would otherwise muddy the // producer's balance). So producer draining to zero is an exact signal that the array's release // callback fired. - BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); - BufferAllocator structs = root.newChildAllocator("structs", 0, Long.MAX_VALUE); - try (ArrowArray array = ArrowArray.allocateNew(structs); - ArrowSchema schema = ArrowSchema.allocateNew(structs)) { - exportBatch(producer, array, schema); - assertTrue( - producer.getAllocatedMemory() > 0, "producer holds the exported batch before import"); + try (BufferAllocator producer = root.newChildAllocator("producer", 0, Long.MAX_VALUE); + BufferAllocator structs = root.newChildAllocator("structs", 0, Long.MAX_VALUE)) { + try (ArrowArray array = ArrowArray.allocateNew(structs); + ArrowSchema schema = ArrowSchema.allocateNew(structs)) { + exportBatch(producer, array, schema); + assertTrue( + producer.getAllocatedMemory() > 0, "producer holds the exported batch before import"); - // A consumer allocator far too small to hold the batch: the import throws part-way through. - BufferAllocator consumer = root.newChildAllocator("consumer", 0, TINY_LIMIT); - try (CDataDictionaryProvider provider = new CDataDictionaryProvider()) { - Schema importSchema = Data.importSchema(consumer, schema, provider); - try (VectorSchemaRoot importRoot = VectorSchemaRoot.create(importSchema, consumer)) { - Exception thrown = - assertThrows( - Exception.class, - () -> Data.importIntoVectorSchemaRoot(consumer, array, importRoot, provider)); - assertTrue( - hasOutOfMemoryCause(thrown), - "mid-import failure must be an allocator OOM: " + thrown); + // A consumer allocator far too small to hold the batch: the import throws part-way through. + try (BufferAllocator consumer = root.newChildAllocator("consumer", 0, TINY_LIMIT); + CDataDictionaryProvider provider = new CDataDictionaryProvider()) { + Schema importSchema = Data.importSchema(consumer, schema, provider); + try (VectorSchemaRoot importRoot = VectorSchemaRoot.create(importSchema, consumer)) { + Exception thrown = + assertThrows( + Exception.class, + () -> Data.importIntoVectorSchemaRoot(consumer, array, importRoot, provider)); + assertTrue( + hasOutOfMemoryCause(thrown), + "mid-import failure must be an allocator OOM: " + thrown); + } } - } - consumer.close(); - // The array's release callback must have fired despite the mid-import OOM, freeing the whole - // exported batch. On the unfixed retain-before-wrap code the batch is stranded instead. - assertEquals( - 0L, - producer.getAllocatedMemory(), - "import OOM leaked the exported batch (producer not drained)"); + // The array's release callback must have fired despite the mid-import OOM, freeing the + // whole exported batch. On the unfixed retain-before-wrap code the batch is stranded. + assertEquals( + 0L, + producer.getAllocatedMemory(), + "import OOM leaked the exported batch (producer not drained)"); + } } - // Closing these (and the RootAllocator in tearDown) additionally asserts nothing leaked at all. - producer.close(); - structs.close(); } /** True if {@code t} is, or is caused by, an Arrow {@link OutOfMemoryException}. */