Skip to content

Commit 714ddc3

Browse files
hutiefang76hutiefang76
committed
GH-1385: Match the JNI return type when closing exported C Data
Use CallVoidMethod for PrivateData.close in array/schema and stream release callbacks. Keep cleanup intact and exercise all three exported structures, including close failures and nested retained buffers. Generated-by: DeepSeek V4 Flash and Codex
1 parent 617edd5 commit 714ddc3

2 files changed

Lines changed: 237 additions & 2 deletions

File tree

‎c/src/main/cpp/jni_wrapper.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ void release_exported(T* base) {
170170
JNIEnvGuard guard(private_data->vm_);
171171
JNIEnv* env = guard.env();
172172

173-
env->CallObjectMethod(private_data->j_private_data_, kPrivateDataCloseMethod);
173+
env->CallVoidMethod(private_data->j_private_data_, kPrivateDataCloseMethod);
174174
if (env->ExceptionCheck()) {
175175
// Can't signal this to caller, so log and then try to free things
176176
// as best we can
@@ -295,7 +295,7 @@ void ArrowArrayStreamRelease(ArrowArrayStream* stream) {
295295
JNIEnvGuard guard(private_data->vm_);
296296
JNIEnv* env = guard.env();
297297

298-
env->CallObjectMethod(private_data->j_private_data_, kPrivateDataCloseMethod);
298+
env->CallVoidMethod(private_data->j_private_data_, kPrivateDataCloseMethod);
299299
if (env->ExceptionCheck()) {
300300
env->ExceptionDescribe();
301301
env->ExceptionClear();

‎c/src/test/java/org/apache/arrow/c/ReleaseTest.java‎

Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.c;
18+
19+
import static org.assertj.core.api.Assertions.assertThat;
20+
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
21+
22+
import java.util.Arrays;
23+
import java.util.Collections;
24+
import org.apache.arrow.c.jni.JniWrapper;
25+
import org.apache.arrow.c.jni.PrivateData;
26+
import org.apache.arrow.memory.BufferAllocator;
27+
import org.apache.arrow.memory.RootAllocator;
28+
import org.apache.arrow.vector.FieldVector;
29+
import org.apache.arrow.vector.IntVector;
30+
import org.apache.arrow.vector.VectorSchemaRoot;
31+
import org.apache.arrow.vector.types.pojo.ArrowType;
32+
import org.apache.arrow.vector.types.pojo.Field;
33+
import org.apache.arrow.vector.types.pojo.FieldType;
34+
import org.apache.arrow.vector.types.pojo.Schema;
35+
import org.junit.jupiter.api.Test;
36+
import org.junit.jupiter.params.ParameterizedTest;
37+
import org.junit.jupiter.params.provider.EnumSource;
38+
39+
/**
40+
* Release lifecycle of exported C Data Interface structures: running the native release callback
41+
* closes the Java {@link PrivateData} exactly once, clears the structure's release markers so a
42+
* second release is a no-op, and frees everything the export retained. Covers both callbacks, the
43+
* one shared by {@link ArrowSchema}/{@link ArrowArray} and the {@link ArrowArrayStream} one. See
44+
* https://github.com/apache/arrow-java/issues/1385.
45+
*/
46+
final class ReleaseTest {
47+
48+
/** Exported structure kinds, one per native release callback. */
49+
private enum Kind {
50+
SCHEMA,
51+
ARRAY,
52+
STREAM
53+
}
54+
55+
/** Release callback and private data pointers of an exported structure. */
56+
private static final class State {
57+
final long releaseCallback;
58+
final long privateData;
59+
60+
State(long releaseCallback, long privateData) {
61+
this.releaseCallback = releaseCallback;
62+
this.privateData = privateData;
63+
}
64+
}
65+
66+
/** Private data that counts how often the native release callback closed it. */
67+
private static final class CountingPrivateData implements PrivateData {
68+
private final boolean failOnClose;
69+
int closeCount;
70+
71+
CountingPrivateData(boolean failOnClose) {
72+
this.failOnClose = failOnClose;
73+
}
74+
75+
@Override
76+
public void close() {
77+
closeCount++;
78+
if (failOnClose) {
79+
throw new IllegalStateException("Deliberate failure from PrivateData.close");
80+
}
81+
}
82+
}
83+
84+
@ParameterizedTest
85+
@EnumSource(Kind.class)
86+
void releaseClosesPrivateDataOnce(Kind kind) {
87+
CountingPrivateData privateData = new CountingPrivateData(false);
88+
89+
try (BufferAllocator allocator = new RootAllocator()) {
90+
try (BaseStruct struct = export(allocator, kind, privateData)) {
91+
final State exported = state(struct);
92+
assertThat(exported.releaseCallback).isNotZero();
93+
assertThat(exported.privateData).isNotZero();
94+
95+
struct.release();
96+
97+
assertThat(privateData.closeCount).isEqualTo(1);
98+
final State released = state(struct);
99+
assertThat(released.releaseCallback).isZero();
100+
assertThat(released.privateData).isZero();
101+
102+
// Releasing again must not run the release callback (and thus not close twice)
103+
struct.release();
104+
105+
assertThat(privateData.closeCount).isEqualTo(1);
106+
assertThat(state(struct).releaseCallback).isZero();
107+
}
108+
109+
assertThat(allocator.getAllocatedMemory()).isZero();
110+
}
111+
}
112+
113+
@ParameterizedTest
114+
@EnumSource(Kind.class)
115+
void releaseCleansUpWhenCloseThrows(Kind kind) {
116+
CountingPrivateData privateData = new CountingPrivateData(true);
117+
118+
try (BufferAllocator allocator = new RootAllocator()) {
119+
try (BaseStruct struct = export(allocator, kind, privateData)) {
120+
assertThat(state(struct).releaseCallback).isNotZero();
121+
122+
// The native side describes and clears the failure: release must not propagate it
123+
assertDoesNotThrow(struct::release);
124+
125+
assertThat(privateData.closeCount).isEqualTo(1);
126+
final State released = state(struct);
127+
assertThat(released.releaseCallback).isZero();
128+
assertThat(released.privateData).isZero();
129+
}
130+
131+
assertThat(allocator.getAllocatedMemory()).isZero();
132+
}
133+
}
134+
135+
@Test
136+
void releaseOfExportedVectorClearsNestedState() {
137+
final Schema schema =
138+
new Schema(
139+
Collections.singletonList(
140+
new Field(
141+
"outer",
142+
FieldType.nullable(new ArrowType.Struct()),
143+
Arrays.asList(
144+
Field.nullable("ints", new ArrowType.Int(32, true)),
145+
Field.nullable("ints2", new ArrowType.Int(32, true))))));
146+
147+
try (BufferAllocator allocator = new RootAllocator()) {
148+
try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator);
149+
ArrowArray array = ArrowArray.allocateNew(allocator);
150+
ArrowSchema exportedSchema = ArrowSchema.allocateNew(allocator)) {
151+
final FieldVector vector = root.getVector("outer");
152+
fill(vector, 4);
153+
root.setRowCount(4);
154+
155+
Data.exportVector(allocator, vector, null, array, exportedSchema);
156+
157+
final ArrowArray.Snapshot exportedArray = array.snapshot();
158+
assertThat(exportedArray.release).isNotZero();
159+
assertThat(exportedArray.private_data).isNotZero();
160+
assertThat(exportedArray.n_children).isEqualTo(2);
161+
162+
final ArrowSchema.Snapshot exportedField = exportedSchema.snapshot();
163+
assertThat(exportedField.release).isNotZero();
164+
assertThat(exportedField.private_data).isNotZero();
165+
assertThat(exportedField.n_children).isEqualTo(2);
166+
167+
array.release();
168+
exportedSchema.release();
169+
170+
assertThat(array.snapshot().release).isZero();
171+
assertThat(array.snapshot().private_data).isZero();
172+
assertThat(exportedSchema.snapshot().release).isZero();
173+
assertThat(exportedSchema.snapshot().private_data).isZero();
174+
}
175+
176+
// The nested private data close callbacks free the buffers the export retained.
177+
assertThat(allocator.getAllocatedMemory()).isZero();
178+
}
179+
}
180+
181+
private static void fill(FieldVector vector, int valueCount) {
182+
for (FieldVector child : vector.getChildrenFromFields()) {
183+
final IntVector intVector = (IntVector) child;
184+
intVector.allocateNew(valueCount);
185+
for (int i = 0; i < valueCount; i++) {
186+
intVector.setSafe(i, i);
187+
}
188+
intVector.setValueCount(valueCount);
189+
}
190+
vector.setValueCount(valueCount);
191+
}
192+
193+
/** Exports a fresh structure of the given kind, so that native code owns its release. */
194+
private static BaseStruct export(BufferAllocator allocator, Kind kind, PrivateData privateData) {
195+
switch (kind) {
196+
case SCHEMA:
197+
{
198+
final ArrowSchema schema = ArrowSchema.allocateNew(allocator);
199+
schema.save(new ArrowSchema.Snapshot());
200+
JniWrapper.get().exportSchema(schema.memoryAddress(), privateData);
201+
return schema;
202+
}
203+
case ARRAY:
204+
{
205+
final ArrowArray array = ArrowArray.allocateNew(allocator);
206+
array.save(new ArrowArray.Snapshot());
207+
JniWrapper.get().exportArray(array.memoryAddress(), privateData);
208+
return array;
209+
}
210+
case STREAM:
211+
{
212+
final ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator);
213+
stream.save(new ArrowArrayStream.Snapshot());
214+
JniWrapper.get().exportArrayStream(stream.memoryAddress(), privateData);
215+
return stream;
216+
}
217+
default:
218+
throw new IllegalArgumentException("Unsupported structure kind: " + kind);
219+
}
220+
}
221+
222+
/** Reads the native release callback and private data pointer of an exported structure. */
223+
private static State state(BaseStruct struct) {
224+
if (struct instanceof ArrowSchema) {
225+
final ArrowSchema.Snapshot snapshot = ((ArrowSchema) struct).snapshot();
226+
return new State(snapshot.release, snapshot.private_data);
227+
}
228+
if (struct instanceof ArrowArray) {
229+
final ArrowArray.Snapshot snapshot = ((ArrowArray) struct).snapshot();
230+
return new State(snapshot.release, snapshot.private_data);
231+
}
232+
final ArrowArrayStream.Snapshot snapshot = ((ArrowArrayStream) struct).snapshot();
233+
return new State(snapshot.release, snapshot.private_data);
234+
}
235+
}

0 commit comments

Comments
 (0)