GH-1239: Fix memory leak when C Data import hits allocator limit mid-array (#1240)

`ReferenceCountedArrowArray.unsafeAssociateAllocation` calls `retain()`
before `wrapForeignAllocation`. If `wrapForeignAllocation` throws
(allocator over its limit), the retain is never balanced, the reference
count stays elevated, and the C Data release callback never fires —
leaking the producer's native memory.

**Fix:** call `retain()` after `wrapForeignAllocation` returns. Same
behavior on the success path; on failure the existing `finally` in
`ArrayImporter.importArray` drives the count to zero and fires the
release callback.

**Test:** `ImportOutOfMemoryTest` — exports a batch from a "producer"
allocator, tries to import it into a too-small consumer, and asserts the
producer drains to zero after the OOM. Fails on the original code,
passes with the fix.

Fixes: #1239

---------

Signed-off-by: Sandesh Kumar <kusandes@amazon.com>
Co-authored-by: Sandesh Kumar <kusandes@amazon.com>

Author: sandeshkr419

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 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) {
+              @Override
+              protected void release0() {
+                ReferenceCountedArrowArray.this.release();
+              }
+            });
     retain();
-    return trackingAllocator.wrapForeignAllocation(
-        new ForeignAllocation(capacity, memoryAddress) {
-          @Override
-          protected void release0() {
-            ReferenceCountedArrowArray.this.release();
-          }
-        });
+    return buf;
   }
 }

c/src/test/java/org/apache/arrow/c/ImportOutOfMemoryTest.java

@@ -0,0 +1,139 @@
+/*
+ * 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.
+ *
+ * <p>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.
+    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.
+        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);
+          }
+        }
+
+        // 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)");
+      }
+    }
+  }
+
+  /** 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<FieldVector> 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);
+    }
+  }
+}