From 8e09280453301ebd4dfcefe3aa6f951c588ff055 Mon Sep 17 00:00:00 2001 From: David Li Date: Wed, 24 Jun 2026 12:24:59 +0900 Subject: [PATCH 1/3] feat(java/driver/jni): prevent closing statement before connection --- .../arrow/adbc/driver/jni/JniConnection.java | 51 +++++++-- .../arrow/adbc/driver/jni/JniDatabase.java | 18 ++- .../arrow/adbc/driver/jni/JniStatement.java | 28 ++++- .../adbc/driver/jni/impl/ChildReferences.java | 55 +++++++++ .../driver/jni/impl/HasChildReferences.java | 21 ++++ .../driver/jni/impl/NativeQueryResult.java | 6 +- .../adbc/driver/jni/impl/TiedArrowReader.java | 105 ++++++++++++++++++ .../arrow/adbc/driver/jni/JniDriverTest.java | 49 +++++++- .../arrow/adbc/driver/jni/impl/ImplTest.java | 34 ++++++ 9 files changed, 343 insertions(+), 24 deletions(-) create mode 100644 java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/ChildReferences.java create mode 100644 java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/HasChildReferences.java create mode 100644 java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/TiedArrowReader.java diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java index a67b210f57..bf35d62f65 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java @@ -25,26 +25,42 @@ import org.apache.arrow.adbc.core.IngestOption; import org.apache.arrow.adbc.core.IsolationLevel; import org.apache.arrow.adbc.core.TypedKey; +import org.apache.arrow.adbc.driver.jni.impl.ChildReferences; +import org.apache.arrow.adbc.driver.jni.impl.HasChildReferences; import org.apache.arrow.adbc.driver.jni.impl.JniLoader; import org.apache.arrow.adbc.driver.jni.impl.NativeConnectionHandle; import org.apache.arrow.adbc.driver.jni.impl.NativeStatementHandle; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Schema; import org.checkerframework.checker.nullness.qual.Nullable; -public class JniConnection implements AdbcConnection { +public class JniConnection implements AdbcConnection, HasChildReferences { private final BufferAllocator allocator; private final NativeConnectionHandle handle; + private final ChildReferences childReferences; + // Hold the owning database alive, and try to ensure this connection gets cleaned up before the + // database does + private @Nullable HasChildReferences parent; - public JniConnection(BufferAllocator allocator, NativeConnectionHandle handle) { + public JniConnection( + BufferAllocator allocator, HasChildReferences parent, NativeConnectionHandle handle) { this.allocator = allocator; this.handle = handle; + this.childReferences = new ChildReferences(); + this.parent = parent; + parent.getChildReferences().addReference(this); + } + + @Override + public ChildReferences getChildReferences() { + return childReferences; } @Override public AdbcStatement createStatement() throws AdbcException { - return new JniStatement(allocator, JniLoader.INSTANCE.openStatement(handle)); + return new JniStatement(allocator, this, JniLoader.INSTANCE.openStatement(handle)); } @Override @@ -108,7 +124,7 @@ AdbcStatement bulkIngestImpl(String targetTableName, BulkIngestMode mode, Ingest } } - return new JniStatement(allocator, stmtHandle); + return new JniStatement(allocator, this, stmtHandle); } catch (Exception e) { stmtHandle.close(); throw e; @@ -117,7 +133,7 @@ AdbcStatement bulkIngestImpl(String targetTableName, BulkIngestMode mode, Ingest @Override public ArrowReader getInfo(int @Nullable [] infoCodes) throws AdbcException { - return JniLoader.INSTANCE.connectionGetInfo(handle, infoCodes).importStream(allocator); + return JniLoader.INSTANCE.connectionGetInfo(handle, infoCodes).importStream(allocator, this); } @Override @@ -138,7 +154,7 @@ public ArrowReader getObjects( tableNamePattern, tableTypes, columnNamePattern) - .importStream(allocator); + .importStream(allocator, this); } @Override @@ -151,7 +167,7 @@ public Schema getTableSchema(String catalog, String dbSchema, String tableName) @Override public ArrowReader getTableTypes() throws AdbcException { - return JniLoader.INSTANCE.connectionGetTableTypes(handle).importStream(allocator); + return JniLoader.INSTANCE.connectionGetTableTypes(handle).importStream(allocator, this); } @Override @@ -257,7 +273,9 @@ public void setCurrentDbSchema(String dbSchema) throws AdbcException { @Override public ArrowReader readPartition(ByteBuffer descriptor) throws AdbcException { - return JniLoader.INSTANCE.connectionReadPartition(handle, descriptor).importStream(allocator); + return JniLoader.INSTANCE + .connectionReadPartition(handle, descriptor) + .importStream(allocator, this); } @Override @@ -267,17 +285,26 @@ public ArrowReader getStatistics( return JniLoader.INSTANCE .connectionGetStatistics( handle, catalogPattern, dbSchemaPattern, tableNamePattern, approximate) - .importStream(allocator); + .importStream(allocator, this); } @Override public ArrowReader getStatisticNames() throws AdbcException { - return JniLoader.INSTANCE.connectionGetStatisticNames(handle).importStream(allocator); + return JniLoader.INSTANCE.connectionGetStatisticNames(handle).importStream(allocator, this); } @Override - public void close() { - handle.close(); + public void close() throws AdbcException { + try { + AutoCloseables.close(childReferences, handle); + } catch (Exception e) { + throw AdbcException.internal("[jni] failed to close connection").withCause(e); + } + final var parent = this.parent; + if (parent != null) { + parent.getChildReferences().releaseReference(this); + this.parent = null; + } } @Override diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniDatabase.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniDatabase.java index 9468107056..482330fb02 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniDatabase.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniDatabase.java @@ -21,27 +21,37 @@ import org.apache.arrow.adbc.core.AdbcDatabase; import org.apache.arrow.adbc.core.AdbcException; import org.apache.arrow.adbc.core.TypedKey; +import org.apache.arrow.adbc.driver.jni.impl.ChildReferences; +import org.apache.arrow.adbc.driver.jni.impl.HasChildReferences; import org.apache.arrow.adbc.driver.jni.impl.JniLoader; import org.apache.arrow.adbc.driver.jni.impl.NativeDatabaseHandle; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.util.AutoCloseables; -public class JniDatabase implements AdbcDatabase { +public class JniDatabase implements AdbcDatabase, HasChildReferences { private final BufferAllocator allocator; private final NativeDatabaseHandle handle; + private final ChildReferences childReferences; public JniDatabase(BufferAllocator allocator, NativeDatabaseHandle handle) { this.allocator = allocator; this.handle = handle; + this.childReferences = new ChildReferences(); + } + + @Override + public ChildReferences getChildReferences() { + return childReferences; } @Override public AdbcConnection connect() throws AdbcException { - return new JniConnection(allocator, JniLoader.INSTANCE.openConnection(handle)); + return new JniConnection(allocator, this, JniLoader.INSTANCE.openConnection(handle)); } @Override - public void close() { - handle.close(); + public void close() throws Exception { + AutoCloseables.close(childReferences, handle); } @Override diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java index 0e5da15b2b..8971336466 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java @@ -24,6 +24,8 @@ import org.apache.arrow.adbc.core.AdbcException; import org.apache.arrow.adbc.core.AdbcStatement; import org.apache.arrow.adbc.core.TypedKey; +import org.apache.arrow.adbc.driver.jni.impl.ChildReferences; +import org.apache.arrow.adbc.driver.jni.impl.HasChildReferences; import org.apache.arrow.adbc.driver.jni.impl.JniLoader; import org.apache.arrow.adbc.driver.jni.impl.NativePartitionResult; import org.apache.arrow.adbc.driver.jni.impl.NativeQueryResult; @@ -39,15 +41,28 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.checkerframework.checker.nullness.qual.Nullable; -public class JniStatement implements AdbcStatement { +public class JniStatement implements AdbcStatement, HasChildReferences { private final BufferAllocator allocator; private final NativeStatementHandle handle; + private final ChildReferences childReferences; + // Hold the owning connection alive, and try to ensure this statement gets cleaned up before the + // connection does + private @Nullable HasChildReferences parent; private @Nullable VectorSchemaRoot bindRoot; private @Nullable ArrowReader bindStream; - public JniStatement(BufferAllocator allocator, NativeStatementHandle handle) { + public JniStatement( + BufferAllocator allocator, HasChildReferences parent, NativeStatementHandle handle) { this.allocator = allocator; this.handle = handle; + this.childReferences = new ChildReferences(); + this.parent = parent; + parent.getChildReferences().addReference(this); + } + + @Override + public ChildReferences getChildReferences() { + return childReferences; } @Override @@ -126,7 +141,7 @@ public PartitionResult executePartitioned() throws AdbcException { public QueryResult executeQuery() throws AdbcException { exportBind(); NativeQueryResult result = JniLoader.INSTANCE.statementExecuteQuery(handle); - return new QueryResult(result.rowsAffected(), result.importStream(allocator)); + return new QueryResult(result.rowsAffected(), result.importStream(allocator, this)); } @Override @@ -171,10 +186,15 @@ public Iterator pollPartitioned() throws AdbcException { @Override public void close() throws AdbcException { try { - AutoCloseables.close(handle, bindStream); + AutoCloseables.close(childReferences, handle, bindStream); } catch (Exception e) { throw AdbcException.internal("[jni] failed to close statement").withCause(e); } + final var parent = this.parent; + if (parent != null) { + parent.getChildReferences().releaseReference(this); + this.parent = null; + } } @Override diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/ChildReferences.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/ChildReferences.java new file mode 100644 index 0000000000..af6ad9a512 --- /dev/null +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/ChildReferences.java @@ -0,0 +1,55 @@ +/* + * 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.adbc.driver.jni.impl; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Set; +import java.util.WeakHashMap; +import org.apache.arrow.util.AutoCloseables; + +/** + * Track child resources for the ADBC FFI. + * + *

You are supposed to close statements before closing the connection (etc.). This class helps + * track those references to prevent misuse at runtime. + */ +public final class ChildReferences implements AutoCloseable { + private final Set openReferences; + + public ChildReferences() { + this.openReferences = Collections.newSetFromMap(new WeakHashMap<>()); + } + + public void close() throws Exception { + try { + var closeables = new ArrayList<>(openReferences); + AutoCloseables.close(closeables); + } finally { + openReferences.clear(); + } + } + + public void addReference(AutoCloseable any) { + openReferences.add(any); + } + + public void releaseReference(AutoCloseable any) { + openReferences.remove(any); + } +} diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/HasChildReferences.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/HasChildReferences.java new file mode 100644 index 0000000000..cbdca1529f --- /dev/null +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/HasChildReferences.java @@ -0,0 +1,21 @@ +/* + * 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.adbc.driver.jni.impl; + +public interface HasChildReferences { + ChildReferences getChildReferences(); +} diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeQueryResult.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeQueryResult.java index 6399dfbb36..c7409e36a0 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeQueryResult.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/NativeQueryResult.java @@ -21,6 +21,7 @@ import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.ipc.ArrowReader; +import org.checkerframework.checker.nullness.qual.Nullable; public class NativeQueryResult { private final long rowsAffected; @@ -39,10 +40,11 @@ public long rowsAffected() { } /** Import the C Data stream into a Java ArrowReader. */ - public ArrowReader importStream(BufferAllocator allocator) { + public ArrowReader importStream(BufferAllocator allocator, @Nullable HasChildReferences parent) { try (final ArrowArrayStream cStream = ArrowArrayStream.allocateNew(allocator)) { cStream.save(streamSnapshot); - return Data.importArrayStream(allocator, cStream); + final var reader = Data.importArrayStream(allocator, cStream); + return new TiedArrowReader(allocator, reader, parent); } } } diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/TiedArrowReader.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/TiedArrowReader.java new file mode 100644 index 0000000000..db9006d37d --- /dev/null +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/TiedArrowReader.java @@ -0,0 +1,105 @@ +/* + * 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.adbc.driver.jni.impl; + +import java.io.IOException; +import java.util.Map; +import java.util.Set; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.dictionary.Dictionary; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.apache.arrow.vector.types.pojo.Schema; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** A proxy {@link ArrowReader} that keeps an associated ADBC resource alive. */ +class TiedArrowReader extends ArrowReader { + private final ArrowReader delegate; + private @Nullable HasChildReferences parent; + + TiedArrowReader( + BufferAllocator allocator, ArrowReader delegate, @Nullable HasChildReferences parent) { + // XXX: ArrowReader being an abstract class and not an interface is a massive wart in arrow-java + // design + super(allocator); + this.delegate = delegate; + this.parent = parent; + if (parent != null) { + parent.getChildReferences().addReference(this); + } + } + + @Override + public void close(boolean closeReadSource) throws IOException { + delegate.close(closeReadSource); + if (parent != null) { + parent.getChildReferences().releaseReference(this); + } + parent = null; + } + + @Override + public void close() throws IOException { + delegate.close(); + if (parent != null) { + parent.getChildReferences().releaseReference(this); + } + parent = null; + } + + @Override + public long bytesRead() { + return delegate.bytesRead(); + } + + @Override + public boolean loadNextBatch() throws IOException { + return delegate.loadNextBatch(); + } + + @Override + public Set getDictionaryIds() { + return delegate.getDictionaryIds(); + } + + @Override + public Dictionary lookup(long id) { + return delegate.lookup(id); + } + + @Override + public Map getDictionaryVectors() throws IOException { + return delegate.getDictionaryVectors(); + } + + @Override + public VectorSchemaRoot getVectorSchemaRoot() throws IOException { + return delegate.getVectorSchemaRoot(); + } + + @Override + protected void closeReadSource() { + // Not actually called because we delegate all public methods + throw new AssertionError(); + } + + @Override + protected Schema readSchema() throws IOException { + // Not actually called because we delegate all public methods + throw new AssertionError(); + } +} diff --git a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/JniDriverTest.java b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/JniDriverTest.java index 94712c86f1..45368be733 100644 --- a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/JniDriverTest.java +++ b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/JniDriverTest.java @@ -144,7 +144,7 @@ void querySimple() throws Exception { } } - // Ensure strings with characters that differ between UTF-8 and Java's "modifiefd UTF-8" are + // Ensure strings with characters that differ between UTF-8 and Java's "modified UTF-8" are // properly serialized @Test void queryNonBmpUtf8() throws Exception { @@ -152,7 +152,9 @@ void queryNonBmpUtf8() throws Exception { JniDriver driver = new JniDriver(allocator); Map parameters = new HashMap<>(); JniDriver.PARAM_DRIVER.set(parameters, "adbc_driver_sqlite"); - String expected = "\uD83D\uDE00"; + String expected = "\uD83D\uDE00"; // U+1f600 Grinning face in big-endian + assertThat(expected.getBytes(StandardCharsets.UTF_8)) + .isEqualTo(new byte[] {(byte) 0xf0, (byte) 0x9f, (byte) 0x98, (byte) 0x80}); try (final AdbcDatabase db = driver.open(parameters); final AdbcConnection conn = db.connect(); @@ -186,6 +188,49 @@ void statementThrowsAfterClose() throws Exception { } } + @Test + void connectionClosesStatements() throws Exception { + try (final BufferAllocator allocator = new RootAllocator()) { + JniDriver driver = new JniDriver(allocator); + Map parameters = new HashMap<>(); + JniDriver.PARAM_DRIVER.set(parameters, "adbc_driver_sqlite"); + + try (final AdbcDatabase db = driver.open(parameters)) { + final AdbcConnection conn = db.connect(); + final AdbcStatement stmt = conn.createStatement(); + + conn.close(); + + assertThatThrownBy(() -> stmt.setSqlQuery("SELECT 1")) + .isInstanceOf(IllegalStateException.class) + .hasMessage("Native statement handle is closed"); + } + } + } + + @Test + void databaseClosesConnections() throws Exception { + try (final BufferAllocator allocator = new RootAllocator()) { + JniDriver driver = new JniDriver(allocator); + Map parameters = new HashMap<>(); + JniDriver.PARAM_DRIVER.set(parameters, "adbc_driver_sqlite"); + + final AdbcDatabase db = driver.open(parameters); + final AdbcConnection conn = db.connect(); + final AdbcStatement stmt = conn.createStatement(); + + db.close(); + + assertThatThrownBy(() -> stmt.setSqlQuery("SELECT 1")) + .isInstanceOf(IllegalStateException.class) + .hasMessage("Native statement handle is closed"); + + assertThatThrownBy(conn::commit) + .isInstanceOf(IllegalStateException.class) + .hasMessage("Native connection handle is closed"); + } + } + @Test void queryLarge() throws Exception { try (final BufferAllocator allocator = new RootAllocator()) { diff --git a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java index e102ebf221..9480d32545 100644 --- a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java +++ b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java @@ -18,6 +18,7 @@ import static org.assertj.core.api.Assertions.assertThat; +import java.lang.ref.WeakReference; import java.nio.ByteBuffer; import org.junit.jupiter.api.Test; @@ -119,4 +120,37 @@ void offsetSlow() throws Exception { assertThat(buf.remaining()).isEqualTo(3); assertThat(JniLoader.INSTANCE.internalGetByteBuffer(buf)).isEqualTo(new byte[] {1, 2, 3}); } + + @Test + void childReferencesCloses() throws Exception { + ChildReferences refs = new ChildReferences(); + var flag = new Closeable(); + refs.addReference(flag); + refs.close(); + assertThat(flag.closed).isTrue(); + } + + @Test + void childReferencesIsWeak() throws Exception { + ChildReferences refs = new ChildReferences(); + var flag = new Closeable(); + refs.addReference(flag); + var ref = new WeakReference<>(flag); + //noinspection UnusedAssignment + flag = null; + System.gc(); + System.gc(); + + assertThat(ref.get()).isNull(); + refs.close(); + } + + static final class Closeable implements AutoCloseable { + boolean closed = false; + + @Override + public void close() throws Exception { + closed = true; + } + } } From ea1254e317a9674b7997f4108ef3b90361778c4b Mon Sep 17 00:00:00 2001 From: David Li Date: Thu, 25 Jun 2026 14:52:29 +0900 Subject: [PATCH 2/3] nits --- .../adbc/driver/jni/impl/TiedArrowReader.java | 23 ++++++++++++------- .../arrow/adbc/driver/jni/JniDriverTest.java | 3 ++- .../arrow/adbc/driver/jni/impl/ImplTest.java | 11 +++++++-- 3 files changed, 26 insertions(+), 11 deletions(-) diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/TiedArrowReader.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/TiedArrowReader.java index db9006d37d..c3c1abea81 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/TiedArrowReader.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/TiedArrowReader.java @@ -45,20 +45,27 @@ class TiedArrowReader extends ArrowReader { @Override public void close(boolean closeReadSource) throws IOException { - delegate.close(closeReadSource); - if (parent != null) { - parent.getChildReferences().releaseReference(this); + try { + delegate.close(closeReadSource); + } finally { + // release even if we couldn't close properly + if (parent != null) { + parent.getChildReferences().releaseReference(this); + } + parent = null; } - parent = null; } @Override public void close() throws IOException { - delegate.close(); - if (parent != null) { - parent.getChildReferences().releaseReference(this); + try { + delegate.close(); + } finally { + if (parent != null) { + parent.getChildReferences().releaseReference(this); + } + parent = null; } - parent = null; } @Override diff --git a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/JniDriverTest.java b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/JniDriverTest.java index 45368be733..0605875a17 100644 --- a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/JniDriverTest.java +++ b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/JniDriverTest.java @@ -152,7 +152,8 @@ void queryNonBmpUtf8() throws Exception { JniDriver driver = new JniDriver(allocator); Map parameters = new HashMap<>(); JniDriver.PARAM_DRIVER.set(parameters, "adbc_driver_sqlite"); - String expected = "\uD83D\uDE00"; // U+1f600 Grinning face in big-endian + String expected = "\uD83D\uDE00"; // U+1f600 GRINNING FACE (big-endian UTF-16) + // Sanity check that this encodes to what we expect assertThat(expected.getBytes(StandardCharsets.UTF_8)) .isEqualTo(new byte[] {(byte) 0xf0, (byte) 0x9f, (byte) 0x98, (byte) 0x80}); diff --git a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java index 9480d32545..c52e8e7d6d 100644 --- a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java +++ b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java @@ -18,6 +18,7 @@ import static org.assertj.core.api.Assertions.assertThat; +import java.io.Closeable; import java.lang.ref.WeakReference; import java.nio.ByteBuffer; import org.junit.jupiter.api.Test; @@ -138,8 +139,14 @@ void childReferencesIsWeak() throws Exception { var ref = new WeakReference<>(flag); //noinspection UnusedAssignment flag = null; - System.gc(); - System.gc(); + + for (int i = 0; i < 50; i++) { + System.gc(); + if (ref.get() == null) { + break; + } + Thread.sleep(100); + } assertThat(ref.get()).isNull(); refs.close(); From 4f1fd242e08a90793b66d5070d613665e66221b2 Mon Sep 17 00:00:00 2001 From: David Li Date: Fri, 26 Jun 2026 07:56:50 +0900 Subject: [PATCH 3/3] fixes --- .../arrow/adbc/driver/jni/JniConnection.java | 11 ++++++----- .../arrow/adbc/driver/jni/JniDatabase.java | 8 ++++++-- .../arrow/adbc/driver/jni/JniStatement.java | 11 ++++++----- .../adbc/driver/jni/impl/ChildReferences.java | 19 +++++++++++++------ .../arrow/adbc/driver/jni/impl/ImplTest.java | 1 - 5 files changed, 31 insertions(+), 19 deletions(-) diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java index bf35d62f65..23b96ee853 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniConnection.java @@ -299,11 +299,12 @@ public void close() throws AdbcException { AutoCloseables.close(childReferences, handle); } catch (Exception e) { throw AdbcException.internal("[jni] failed to close connection").withCause(e); - } - final var parent = this.parent; - if (parent != null) { - parent.getChildReferences().releaseReference(this); - this.parent = null; + } finally { + final var parent = this.parent; + if (parent != null) { + parent.getChildReferences().releaseReference(this); + this.parent = null; + } } } diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniDatabase.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniDatabase.java index 482330fb02..1f5e10eb83 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniDatabase.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniDatabase.java @@ -50,8 +50,12 @@ public AdbcConnection connect() throws AdbcException { } @Override - public void close() throws Exception { - AutoCloseables.close(childReferences, handle); + public void close() throws AdbcException { + try { + AutoCloseables.close(childReferences, handle); + } catch (Exception e) { + throw AdbcException.internal("[jni] failed to close database").withCause(e); + } } @Override diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java index 8971336466..427d6b2837 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/JniStatement.java @@ -189,11 +189,12 @@ public void close() throws AdbcException { AutoCloseables.close(childReferences, handle, bindStream); } catch (Exception e) { throw AdbcException.internal("[jni] failed to close statement").withCause(e); - } - final var parent = this.parent; - if (parent != null) { - parent.getChildReferences().releaseReference(this); - this.parent = null; + } finally { + final var parent = this.parent; + if (parent != null) { + parent.getChildReferences().releaseReference(this); + this.parent = null; + } } } diff --git a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/ChildReferences.java b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/ChildReferences.java index af6ad9a512..0db99e58b9 100644 --- a/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/ChildReferences.java +++ b/java/driver/jni/src/main/java/org/apache/arrow/adbc/driver/jni/impl/ChildReferences.java @@ -28,20 +28,27 @@ * *

You are supposed to close statements before closing the connection (etc.). This class helps * track those references to prevent misuse at runtime. + * + *

This class is thread-safe. */ public final class ChildReferences implements AutoCloseable { private final Set openReferences; public ChildReferences() { - this.openReferences = Collections.newSetFromMap(new WeakHashMap<>()); + // TODO(lidavidm): we could use caffeine LoadingCache with weakKeys instead + this.openReferences = + Collections.synchronizedSet(Collections.newSetFromMap(new WeakHashMap<>())); } public void close() throws Exception { - try { - var closeables = new ArrayList<>(openReferences); - AutoCloseables.close(closeables); - } finally { - openReferences.clear(); + // synchronizedSet requires explicit synchronization for iteration + synchronized (openReferences) { + try { + var closeables = new ArrayList<>(openReferences); + AutoCloseables.close(closeables); + } finally { + openReferences.clear(); + } } } diff --git a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java index c52e8e7d6d..0d86646573 100644 --- a/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java +++ b/java/driver/jni/src/test/java/org/apache/arrow/adbc/driver/jni/impl/ImplTest.java @@ -18,7 +18,6 @@ import static org.assertj.core.api.Assertions.assertThat; -import java.io.Closeable; import java.lang.ref.WeakReference; import java.nio.ByteBuffer; import org.junit.jupiter.api.Test;