From 9f9a920ab1a200d4ee5f44b68d1bc1477cf1dcdf Mon Sep 17 00:00:00 2001 From: Yue Zhang Date: Mon, 17 Aug 2026 04:48:38 -0400 Subject: [PATCH 1/2] fix(java): hold Dataset read lock when passing Dataset into native calls Dataset guards its native handle with a ReentrantReadWriteLock: reads take the read lock and close() takes the write lock. However, entry points that pass the Dataset object into JNI from outside the class (Fragment scans and row operations, LanceScanner/AsyncScanner creation, SqlQuery, CommitBuilder, Compaction, CompactionTask, VectorTrainer, DatasetDeltaBuilder, and the memwal ShardWriter/LsmScanner/planner constructors) invoked native methods without acquiring that lock, so close() was not mutually exclusive with them. A concurrent close() then hits the jni-rs take_rust_field flaw: when another thread holds the handle mutex, take_rust_field frees the boxed BlockingDataset on its try_lock early-return, causing a native use-after-free (JVM SIGABRT under contention) and leaving both the Rust field and the Java handle in a corrupt state that crashes later, unrelated calls. Add Dataset.acquireReadLock(), which pins the native handle and rejects closed datasets, and wrap every external native call that borrows the Dataset in it. Concurrent callers now either complete before close() or fail cleanly with "Dataset is closed". --- .../main/java/org/lance/CommitBuilder.java | 36 +++++----- java/src/main/java/org/lance/Dataset.java | 20 ++++++ java/src/main/java/org/lance/Fragment.java | 20 ++++-- java/src/main/java/org/lance/SqlQuery.java | 3 +- .../java/org/lance/compaction/Compaction.java | 67 ++++++++++--------- .../org/lance/compaction/CompactionTask.java | 37 +++++----- .../org/lance/delta/DatasetDeltaBuilder.java | 5 +- .../org/lance/index/vector/VectorTrainer.java | 9 ++- .../main/java/org/lance/ipc/AsyncScanner.java | 59 ++++++++-------- .../main/java/org/lance/ipc/LanceScanner.java | 61 +++++++++-------- .../lance/memwal/LsmPointLookupPlanner.java | 4 +- .../java/org/lance/memwal/LsmScanner.java | 8 ++- .../lance/memwal/LsmVectorSearchPlanner.java | 16 +++-- .../java/org/lance/memwal/ShardWriter.java | 8 ++- .../src/test/java/org/lance/FragmentTest.java | 53 +++++++++++++++ 15 files changed, 260 insertions(+), 146 deletions(-) diff --git a/java/src/main/java/org/lance/CommitBuilder.java b/java/src/main/java/org/lance/CommitBuilder.java index fdb95828942..b0e29d45ac2 100644 --- a/java/src/main/java/org/lance/CommitBuilder.java +++ b/java/src/main/java/org/lance/CommitBuilder.java @@ -279,23 +279,25 @@ public CommitBuilder commitTimeout(Duration timeout) { public Dataset execute(Transaction transaction) { Preconditions.checkNotNull(transaction, "Transaction must not be null"); if (dataset != null) { - Dataset result = - nativeCommitToDataset( - dataset, - transaction, - detached, - enableV2ManifestPaths, - writeParams, - useStableRowIds, - storageFormat, - maxRetries, - skipAutoCleanup, - namespaceClient, - tableId, - namespaceClientManagedVersioning, - commitTimeoutNanos); - result.setAllocator(dataset.allocator()); - return result; + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + Dataset result = + nativeCommitToDataset( + dataset, + transaction, + detached, + enableV2ManifestPaths, + writeParams, + useStableRowIds, + storageFormat, + maxRetries, + skipAutoCleanup, + namespaceClient, + tableId, + namespaceClientManagedVersioning, + commitTimeoutNanos); + result.setAllocator(dataset.allocator()); + return result; + } } if (uri != null) { Dataset result = diff --git a/java/src/main/java/org/lance/Dataset.java b/java/src/main/java/org/lance/Dataset.java index 1ce7d63c8b3..0666ff08c72 100644 --- a/java/src/main/java/org/lance/Dataset.java +++ b/java/src/main/java/org/lance/Dataset.java @@ -1736,6 +1736,26 @@ private void updateToNewDataset(Dataset newDataset) { newDataset.nativeDatasetHandle = 0; } + /** + * Acquires a shared read lock that pins the native dataset handle, blocking a concurrent {@link + * #close()} until the lock is released. + * + *

Any code that passes this {@link Dataset} into a native method must hold this lock for the + * whole native call; otherwise {@code close()} can release the native dataset mid-call and crash + * the JVM. The lock is reentrant and intended for try-with-resources use. + * + * @return the acquired read lock + * @throws IllegalArgumentException if the dataset is already closed + */ + public LockManager.ReadLock acquireReadLock() { + LockManager.ReadLock readLock = lockManager.acquireReadLock(); + if (nativeDatasetHandle == 0) { + readLock.close(); + throw new IllegalArgumentException("Dataset is closed"); + } + return readLock; + } + /** * Closes this dataset and releases any system resources associated with it. If the dataset is * already closed, then invoking this method has no effect. diff --git a/java/src/main/java/org/lance/Fragment.java b/java/src/main/java/org/lance/Fragment.java index 0040fdbf869..dfb4652cf58 100644 --- a/java/src/main/java/org/lance/Fragment.java +++ b/java/src/main/java/org/lance/Fragment.java @@ -113,7 +113,9 @@ public LanceScanner newScan(ScanOptions options) { * returns a new fragment with the updated deletion vector. */ public FragmentMetadata deleteRows(List rowIndexes) { - return nativeDeleteRows(dataset, fragmentMetadata.getId(), rowIndexes); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativeDeleteRows(dataset, fragmentMetadata.getId(), rowIndexes); + } } private static native FragmentMetadata nativeDeleteRows( @@ -129,7 +131,9 @@ public int getId() { * @return row counts in this Fragment */ public int countRows() { - return countRowsNative(dataset, fragmentMetadata.getId()); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return countRowsNative(dataset, fragmentMetadata.getId()); + } } /** @@ -153,8 +157,10 @@ public int countRows() { * @return the fragment metadata and new schema. */ public FragmentMergeResult mergeColumns(ArrowArrayStream stream, String leftOn, String rightOn) { - return nativeMergeColumns( - dataset, fragmentMetadata.getId(), stream.memoryAddress(), leftOn, rightOn); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativeMergeColumns( + dataset, fragmentMetadata.getId(), stream.memoryAddress(), leftOn, rightOn); + } } private native FragmentMergeResult nativeMergeColumns( @@ -186,8 +192,10 @@ private native FragmentMergeResult nativeMergeColumns( */ public FragmentUpdateResult updateColumns( ArrowArrayStream stream, String leftOn, String rightOn) { - return nativeUpdateColumns( - dataset, fragmentMetadata.getId(), stream.memoryAddress(), leftOn, rightOn); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativeUpdateColumns( + dataset, fragmentMetadata.getId(), stream.memoryAddress(), leftOn, rightOn); + } } public FragmentUpdateResult updateColumns(ArrowArrayStream stream) { diff --git a/java/src/main/java/org/lance/SqlQuery.java b/java/src/main/java/org/lance/SqlQuery.java index cce6d939222..cb149a7fa6f 100644 --- a/java/src/main/java/org/lance/SqlQuery.java +++ b/java/src/main/java/org/lance/SqlQuery.java @@ -51,7 +51,8 @@ public SqlQuery withRowAddr(boolean withAddr) { } public ArrowReader intoBatchRecords() throws IOException { - try (ArrowArrayStream s = ArrowArrayStream.allocateNew(dataset.allocator())) { + try (LockManager.ReadLock readLock = dataset.acquireReadLock(); + ArrowArrayStream s = ArrowArrayStream.allocateNew(dataset.allocator())) { intoBatchRecords( dataset, sql, Optional.ofNullable(table), withRowId, withRowAddr, s.memoryAddress()); return Data.importArrayStream(dataset.allocator(), s); diff --git a/java/src/main/java/org/lance/compaction/Compaction.java b/java/src/main/java/org/lance/compaction/Compaction.java index 4e231142f6c..6c91969ee46 100644 --- a/java/src/main/java/org/lance/compaction/Compaction.java +++ b/java/src/main/java/org/lance/compaction/Compaction.java @@ -15,6 +15,7 @@ import org.lance.Dataset; import org.lance.JniLoader; +import org.lance.LockManager; import com.google.common.base.Preconditions; @@ -32,21 +33,23 @@ public static CompactionPlan planCompaction( Preconditions.checkNotNull(dataset); Preconditions.checkNotNull(compactionOptions); - return nativePlanCompaction( - dataset, - compactionOptions.getTargetRowsPerFragment(), - compactionOptions.getMaxRowsPerGroup(), - compactionOptions.getMaxBytesPerFile(), - compactionOptions.getMaterializeDeletions(), - compactionOptions.getMaterializeDeletionsThreshold(), - compactionOptions.getNumThreads(), - compactionOptions.getBatchSize(), - compactionOptions.getDeferIndexRemap(), - compactionOptions.getCompactionMode(), - compactionOptions.getBinaryCopyReadBatchBytes(), - compactionOptions.getMaxSourceFragments(), - compactionOptions.getMaxSourceRows(), - compactionOptions.getMaxSourceBytes()); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativePlanCompaction( + dataset, + compactionOptions.getTargetRowsPerFragment(), + compactionOptions.getMaxRowsPerGroup(), + compactionOptions.getMaxBytesPerFile(), + compactionOptions.getMaterializeDeletions(), + compactionOptions.getMaterializeDeletionsThreshold(), + compactionOptions.getNumThreads(), + compactionOptions.getBatchSize(), + compactionOptions.getDeferIndexRemap(), + compactionOptions.getCompactionMode(), + compactionOptions.getBinaryCopyReadBatchBytes(), + compactionOptions.getMaxSourceFragments(), + compactionOptions.getMaxSourceRows(), + compactionOptions.getMaxSourceBytes()); + } } public static CompactionMetrics commitCompaction( @@ -54,22 +57,24 @@ public static CompactionMetrics commitCompaction( Preconditions.checkNotNull(dataset); Preconditions.checkNotNull(rewriteResults); Preconditions.checkNotNull(compactionOptions); - return nativeCommitCompaction( - dataset, - rewriteResults, - compactionOptions.getTargetRowsPerFragment(), - compactionOptions.getMaxRowsPerGroup(), - compactionOptions.getMaxBytesPerFile(), - compactionOptions.getMaterializeDeletions(), - compactionOptions.getMaterializeDeletionsThreshold(), - compactionOptions.getNumThreads(), - compactionOptions.getBatchSize(), - compactionOptions.getDeferIndexRemap(), - compactionOptions.getCompactionMode(), - compactionOptions.getBinaryCopyReadBatchBytes(), - compactionOptions.getMaxSourceFragments(), - compactionOptions.getMaxSourceRows(), - compactionOptions.getMaxSourceBytes()); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativeCommitCompaction( + dataset, + rewriteResults, + compactionOptions.getTargetRowsPerFragment(), + compactionOptions.getMaxRowsPerGroup(), + compactionOptions.getMaxBytesPerFile(), + compactionOptions.getMaterializeDeletions(), + compactionOptions.getMaterializeDeletionsThreshold(), + compactionOptions.getNumThreads(), + compactionOptions.getBatchSize(), + compactionOptions.getDeferIndexRemap(), + compactionOptions.getCompactionMode(), + compactionOptions.getBinaryCopyReadBatchBytes(), + compactionOptions.getMaxSourceFragments(), + compactionOptions.getMaxSourceRows(), + compactionOptions.getMaxSourceBytes()); + } } public static native CompactionMetrics nativeCommitCompaction( diff --git a/java/src/main/java/org/lance/compaction/CompactionTask.java b/java/src/main/java/org/lance/compaction/CompactionTask.java index 63c4a7043f4..c8ca3e9bbb9 100644 --- a/java/src/main/java/org/lance/compaction/CompactionTask.java +++ b/java/src/main/java/org/lance/compaction/CompactionTask.java @@ -14,6 +14,7 @@ package org.lance.compaction; import org.lance.Dataset; +import org.lance.LockManager; import com.google.common.base.MoreObjects; @@ -42,23 +43,25 @@ public String toString() { } public RewriteResult execute(Dataset dataset) { - return nativeExecute( - dataset, - taskData, - readVersion, - compactionOptions.getTargetRowsPerFragment(), - compactionOptions.getMaxRowsPerGroup(), - compactionOptions.getMaxBytesPerFile(), - compactionOptions.getMaterializeDeletions(), - compactionOptions.getMaterializeDeletionsThreshold(), - compactionOptions.getNumThreads(), - compactionOptions.getBatchSize(), - compactionOptions.getDeferIndexRemap(), - compactionOptions.getCompactionMode(), - compactionOptions.getBinaryCopyReadBatchBytes(), - compactionOptions.getMaxSourceFragments(), - compactionOptions.getMaxSourceRows(), - compactionOptions.getMaxSourceBytes()); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativeExecute( + dataset, + taskData, + readVersion, + compactionOptions.getTargetRowsPerFragment(), + compactionOptions.getMaxRowsPerGroup(), + compactionOptions.getMaxBytesPerFile(), + compactionOptions.getMaterializeDeletions(), + compactionOptions.getMaterializeDeletionsThreshold(), + compactionOptions.getNumThreads(), + compactionOptions.getBatchSize(), + compactionOptions.getDeferIndexRemap(), + compactionOptions.getCompactionMode(), + compactionOptions.getBinaryCopyReadBatchBytes(), + compactionOptions.getMaxSourceFragments(), + compactionOptions.getMaxSourceRows(), + compactionOptions.getMaxSourceBytes()); + } } private native RewriteResult nativeExecute( diff --git a/java/src/main/java/org/lance/delta/DatasetDeltaBuilder.java b/java/src/main/java/org/lance/delta/DatasetDeltaBuilder.java index 9084da2ab9c..9813b8aeba1 100755 --- a/java/src/main/java/org/lance/delta/DatasetDeltaBuilder.java +++ b/java/src/main/java/org/lance/delta/DatasetDeltaBuilder.java @@ -15,6 +15,7 @@ import org.lance.Dataset; import org.lance.JniLoader; +import org.lance.LockManager; import java.util.Optional; @@ -71,7 +72,9 @@ public DatasetDeltaBuilder withEndVersion(long version) { /** Build the DatasetDelta after validating builder state. */ public DatasetDelta build() { - return nativeBuild(dataset, comparedAgainst, beginVersion, endVersion); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativeBuild(dataset, comparedAgainst, beginVersion, endVersion); + } } private static native DatasetDelta nativeBuild( diff --git a/java/src/main/java/org/lance/index/vector/VectorTrainer.java b/java/src/main/java/org/lance/index/vector/VectorTrainer.java index 9514c356fe7..640398ffcf6 100755 --- a/java/src/main/java/org/lance/index/vector/VectorTrainer.java +++ b/java/src/main/java/org/lance/index/vector/VectorTrainer.java @@ -15,6 +15,7 @@ import org.lance.Dataset; import org.lance.JniLoader; +import org.lance.LockManager; import org.lance.index.DistanceType; import org.apache.arrow.util.Preconditions; @@ -64,7 +65,9 @@ public static float[] trainIvfCentroids( column != null && !column.isEmpty(), "column cannot be null or empty"); Preconditions.checkArgument(params != null, "params cannot be null"); Preconditions.checkArgument(distanceType != null, "distanceType cannot be null"); - return nativeTrainIvfCentroids(dataset, column, params, distanceType.toString()); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativeTrainIvfCentroids(dataset, column, params, distanceType.toString()); + } } /** @@ -98,7 +101,9 @@ public static float[] trainPqCodebook( column != null && !column.isEmpty(), "column cannot be null or empty"); Preconditions.checkArgument(params != null, "params cannot be null"); Preconditions.checkArgument(distanceType != null, "distanceType cannot be null"); - return nativeTrainPqCodebook(dataset, column, params, distanceType.toString()); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + return nativeTrainPqCodebook(dataset, column, params, distanceType.toString()); + } } private static native float[] nativeTrainIvfCentroids( diff --git a/java/src/main/java/org/lance/ipc/AsyncScanner.java b/java/src/main/java/org/lance/ipc/AsyncScanner.java index b60d91c25f3..193622f51bc 100644 --- a/java/src/main/java/org/lance/ipc/AsyncScanner.java +++ b/java/src/main/java/org/lance/ipc/AsyncScanner.java @@ -62,34 +62,37 @@ public static AsyncScanner create( Preconditions.checkNotNull(dataset); Preconditions.checkNotNull(options); Preconditions.checkNotNull(allocator); - AsyncScanner scanner = - createAsyncScanner( - dataset, - options.getFragmentIds(), - options.getColumns(), - options.getSubstraitFilter(), - options.getFilter(), - options.getBatchSize(), - options.getBatchSizeBytes(), - options.getIoBufferSize(), - options.getLimit(), - options.getOffset(), - options.getNearest(), - options.getFullTextQuery(), - options.isPrefilter(), - options.isWithRowId(), - options.isWithRowAddress(), - options.getBatchReadahead(), - options.getFragmentReadahead(), - options.isScanInOrder(), - options.getLateMaterialization(), - options.getColumnOrderings(), - options.isUseScalarIndex(), - options.isFastSearch(), - options.getSubstraitAggregate(), - options.isIncludeDeletedRows(), - options.isStrictBatchSize(), - options.isDisableScoringAutoprojection()); + AsyncScanner scanner; + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + scanner = + createAsyncScanner( + dataset, + options.getFragmentIds(), + options.getColumns(), + options.getSubstraitFilter(), + options.getFilter(), + options.getBatchSize(), + options.getBatchSizeBytes(), + options.getIoBufferSize(), + options.getLimit(), + options.getOffset(), + options.getNearest(), + options.getFullTextQuery(), + options.isPrefilter(), + options.isWithRowId(), + options.isWithRowAddress(), + options.getBatchReadahead(), + options.getFragmentReadahead(), + options.isScanInOrder(), + options.getLateMaterialization(), + options.getColumnOrderings(), + options.isUseScalarIndex(), + options.isFastSearch(), + options.getSubstraitAggregate(), + options.isIncludeDeletedRows(), + options.isStrictBatchSize(), + options.isDisableScoringAutoprojection()); + } scanner.allocator = allocator; return scanner; } diff --git a/java/src/main/java/org/lance/ipc/LanceScanner.java b/java/src/main/java/org/lance/ipc/LanceScanner.java index 4cf50f6bdcf..d7b0a22983a 100644 --- a/java/src/main/java/org/lance/ipc/LanceScanner.java +++ b/java/src/main/java/org/lance/ipc/LanceScanner.java @@ -58,35 +58,38 @@ public static LanceScanner create( Preconditions.checkNotNull(dataset); Preconditions.checkNotNull(options); Preconditions.checkNotNull(allocator); - LanceScanner scanner = - createScanner( - dataset, - options.getFragmentIds(), - options.getColumns(), - options.getSubstraitFilter(), - options.getFilter(), - options.getBatchSize(), - options.getBatchSizeBytes(), - options.getIoBufferSize(), - options.getLimit(), - options.getOffset(), - options.getNearest(), - options.getFullTextQuery(), - options.isPrefilter(), - options.isWithRowId(), - options.isWithRowAddress(), - options.getBatchReadahead(), - options.getFragmentReadahead(), - options.isScanInOrder(), - options.getLateMaterialization(), - options.getColumnOrderings(), - options.isUseScalarIndex(), - options.isFastSearch(), - options.getSubstraitAggregate(), - options.isCollectStats(), - options.isIncludeDeletedRows(), - options.isStrictBatchSize(), - options.isDisableScoringAutoprojection()); + LanceScanner scanner; + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + scanner = + createScanner( + dataset, + options.getFragmentIds(), + options.getColumns(), + options.getSubstraitFilter(), + options.getFilter(), + options.getBatchSize(), + options.getBatchSizeBytes(), + options.getIoBufferSize(), + options.getLimit(), + options.getOffset(), + options.getNearest(), + options.getFullTextQuery(), + options.isPrefilter(), + options.isWithRowId(), + options.isWithRowAddress(), + options.getBatchReadahead(), + options.getFragmentReadahead(), + options.isScanInOrder(), + options.getLateMaterialization(), + options.getColumnOrderings(), + options.isUseScalarIndex(), + options.isFastSearch(), + options.getSubstraitAggregate(), + options.isCollectStats(), + options.isIncludeDeletedRows(), + options.isStrictBatchSize(), + options.isDisableScoringAutoprojection()); + } scanner.allocator = allocator; scanner.dataset = dataset; scanner.options = options; diff --git a/java/src/main/java/org/lance/memwal/LsmPointLookupPlanner.java b/java/src/main/java/org/lance/memwal/LsmPointLookupPlanner.java index d58e4e81d57..941a68a1606 100644 --- a/java/src/main/java/org/lance/memwal/LsmPointLookupPlanner.java +++ b/java/src/main/java/org/lance/memwal/LsmPointLookupPlanner.java @@ -60,7 +60,9 @@ public LsmPointLookupPlanner( Preconditions.checkNotNull(dataset, "dataset must not be null"); Preconditions.checkNotNull(shardSnapshots, "shardSnapshots must not be null"); this.allocator = dataset.allocator(); - nativeCreate(dataset, shardSnapshots, Optional.ofNullable(pkColumns)); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + nativeCreate(dataset, shardSnapshots, Optional.ofNullable(pkColumns)); + } } private native void nativeCreate( diff --git a/java/src/main/java/org/lance/memwal/LsmScanner.java b/java/src/main/java/org/lance/memwal/LsmScanner.java index 509e70f7b0a..308d8009d71 100644 --- a/java/src/main/java/org/lance/memwal/LsmScanner.java +++ b/java/src/main/java/org/lance/memwal/LsmScanner.java @@ -61,9 +61,11 @@ private LsmScanner() {} public static LsmScanner fromSnapshots(Dataset dataset, List shardSnapshots) { Preconditions.checkNotNull(dataset, "dataset must not be null"); Preconditions.checkNotNull(shardSnapshots, "shardSnapshots must not be null"); - LsmScanner scanner = createFromSnapshots(dataset, shardSnapshots); - scanner.allocator = dataset.allocator(); - return scanner; + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + LsmScanner scanner = createFromSnapshots(dataset, shardSnapshots); + scanner.allocator = dataset.allocator(); + return scanner; + } } static native LsmScanner createFromSnapshots(Dataset dataset, List shardSnapshots); diff --git a/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java b/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java index b17bb8e37a0..b91e5ae4a7d 100644 --- a/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java +++ b/java/src/main/java/org/lance/memwal/LsmVectorSearchPlanner.java @@ -90,13 +90,15 @@ public LsmVectorSearchPlanner( Preconditions.checkNotNull(shardSnapshots, "shardSnapshots must not be null"); Preconditions.checkNotNull(vectorColumn, "vectorColumn must not be null"); this.allocator = dataset.allocator(); - nativeCreate( - dataset, - shardSnapshots, - vectorColumn, - Optional.ofNullable(pkColumns), - Optional.ofNullable(distanceType), - Optional.ofNullable(filter)); + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + nativeCreate( + dataset, + shardSnapshots, + vectorColumn, + Optional.ofNullable(pkColumns), + Optional.ofNullable(distanceType), + Optional.ofNullable(filter)); + } } private native void nativeCreate( diff --git a/java/src/main/java/org/lance/memwal/ShardWriter.java b/java/src/main/java/org/lance/memwal/ShardWriter.java index da4c621e934..290b4c8f3c2 100644 --- a/java/src/main/java/org/lance/memwal/ShardWriter.java +++ b/java/src/main/java/org/lance/memwal/ShardWriter.java @@ -65,9 +65,11 @@ private ShardWriter() {} public static ShardWriter create(Dataset dataset, String shardId, ShardWriterConfig config) { Preconditions.checkNotNull(dataset, "dataset must not be null"); Preconditions.checkNotNull(shardId, "shardId must not be null"); - ShardWriter writer = createNative(dataset, shardId, config); - writer.allocator = dataset.allocator(); - return writer; + try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { + ShardWriter writer = createNative(dataset, shardId, config); + writer.allocator = dataset.allocator(); + return writer; + } } static native ShardWriter createNative(Dataset dataset, String shardId, ShardWriterConfig config); diff --git a/java/src/test/java/org/lance/FragmentTest.java b/java/src/test/java/org/lance/FragmentTest.java index 3f30fa49863..6bbe4a39231 100644 --- a/java/src/test/java/org/lance/FragmentTest.java +++ b/java/src/test/java/org/lance/FragmentTest.java @@ -39,6 +39,11 @@ import java.util.Collections; import java.util.List; import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; import static org.junit.jupiter.api.Assertions.assertArrayEquals; @@ -514,4 +519,52 @@ void testFragmentStatisticsOnEmptyDataset(@TempDir Path tempDir) { } } } + + @Test + void testCountRowsConcurrentWithClose(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("count_rows_close_race").toString(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + FragmentMetadata fragmentMeta = testDataset.createNewFragment(100); + FragmentOperation.Append appendOp = new FragmentOperation.Append(Arrays.asList(fragmentMeta)); + Dataset dataset = Dataset.commit(allocator, datasetPath, appendOp, Optional.of(1L)); + Fragment fragment = dataset.getFragments().get(0); + + int threadCount = 8; + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + try { + CountDownLatch start = new CountDownLatch(1); + List> futures = new ArrayList<>(); + for (int i = 0; i < threadCount; i++) { + futures.add( + executor.submit( + () -> { + start.await(); + // Hammer countRows until close() wins the race. The only acceptable + // failure is the "Dataset is closed" rejection; anything else (a native + // crash or "Null pointer in rust value from Java") means the native + // handle was released while still in use. + while (true) { + try { + fragment.countRows(); + } catch (IllegalArgumentException e) { + assertEquals("Dataset is closed", e.getMessage()); + return null; + } + } + })); + } + start.countDown(); + Thread.sleep(50); + dataset.close(); + for (Future future : futures) { + future.get(30, TimeUnit.SECONDS); + } + } finally { + executor.shutdownNow(); + } + } + } } From 0cf78e6be2904c0aeb2056ec14b7605df44ee995 Mon Sep 17 00:00:00 2001 From: Yue Zhang Date: Mon, 17 Aug 2026 05:29:25 -0400 Subject: [PATCH 2/2] fix(java): keep raw commit-compaction JNI entry behind the Dataset read lock Compaction.nativeCommitCompaction was a public raw native method, so callers could invoke it directly and bypass the read lock, leaving the close/use race reachable. Preserve the public signature as a Java wrapper that acquires the dataset read lock and delegates to a private renamed native method (commitCompactionNative); commitCompaction now goes through the wrapper instead of locking itself. --- java/lance-jni/src/optimize.rs | 2 +- .../java/org/lance/compaction/Compaction.java | 69 +++++++++++++++---- 2 files changed, 55 insertions(+), 16 deletions(-) diff --git a/java/lance-jni/src/optimize.rs b/java/lance-jni/src/optimize.rs index 7ebc0e7b095..d6117238980 100644 --- a/java/lance-jni/src/optimize.rs +++ b/java/lance-jni/src/optimize.rs @@ -122,7 +122,7 @@ fn inner_plan_compaction<'local>( } #[unsafe(no_mangle)] -pub extern "system" fn Java_org_lance_compaction_Compaction_nativeCommitCompaction<'local>( +pub extern "system" fn Java_org_lance_compaction_Compaction_commitCompactionNative<'local>( mut env: JNIEnv<'local>, _obj: JObject, java_dataset: JObject, // Dataset diff --git a/java/src/main/java/org/lance/compaction/Compaction.java b/java/src/main/java/org/lance/compaction/Compaction.java index 6c91969ee46..4ac44b2a2b4 100644 --- a/java/src/main/java/org/lance/compaction/Compaction.java +++ b/java/src/main/java/org/lance/compaction/Compaction.java @@ -57,27 +57,66 @@ public static CompactionMetrics commitCompaction( Preconditions.checkNotNull(dataset); Preconditions.checkNotNull(rewriteResults); Preconditions.checkNotNull(compactionOptions); + return nativeCommitCompaction( + dataset, + rewriteResults, + compactionOptions.getTargetRowsPerFragment(), + compactionOptions.getMaxRowsPerGroup(), + compactionOptions.getMaxBytesPerFile(), + compactionOptions.getMaterializeDeletions(), + compactionOptions.getMaterializeDeletionsThreshold(), + compactionOptions.getNumThreads(), + compactionOptions.getBatchSize(), + compactionOptions.getDeferIndexRemap(), + compactionOptions.getCompactionMode(), + compactionOptions.getBinaryCopyReadBatchBytes(), + compactionOptions.getMaxSourceFragments(), + compactionOptions.getMaxSourceRows(), + compactionOptions.getMaxSourceBytes()); + } + + /** + * Java wrapper around the raw commit-compaction JNI call. It acquires the dataset read lock so + * the native call cannot race with {@link Dataset#close()}; keep the raw native method private so + * no caller can bypass this lock. + */ + public static CompactionMetrics nativeCommitCompaction( + Dataset dataset, + List rewriteResults, + Optional targetRowsPerFragment, + Optional maxRowsPerGroup, + Optional maxBytesPerFile, + Optional materializeDeletions, + Optional materializeDeletionsThreshold, + Optional numThreads, + Optional batchSize, + Optional deferIndexRemap, + Optional compactionMode, + Optional binaryCopyReadBatchBytes, + Optional maxSourceFragments, + Optional maxSourceRows, + Optional maxSourceBytes) { try (LockManager.ReadLock readLock = dataset.acquireReadLock()) { - return nativeCommitCompaction( + return commitCompactionNative( dataset, rewriteResults, - compactionOptions.getTargetRowsPerFragment(), - compactionOptions.getMaxRowsPerGroup(), - compactionOptions.getMaxBytesPerFile(), - compactionOptions.getMaterializeDeletions(), - compactionOptions.getMaterializeDeletionsThreshold(), - compactionOptions.getNumThreads(), - compactionOptions.getBatchSize(), - compactionOptions.getDeferIndexRemap(), - compactionOptions.getCompactionMode(), - compactionOptions.getBinaryCopyReadBatchBytes(), - compactionOptions.getMaxSourceFragments(), - compactionOptions.getMaxSourceRows(), - compactionOptions.getMaxSourceBytes()); + targetRowsPerFragment, + maxRowsPerGroup, + maxBytesPerFile, + materializeDeletions, + materializeDeletionsThreshold, + numThreads, + batchSize, + deferIndexRemap, + compactionMode, + binaryCopyReadBatchBytes, + maxSourceFragments, + maxSourceRows, + maxSourceBytes); } } - public static native CompactionMetrics nativeCommitCompaction( + private static native CompactionMetrics commitCompactionNative( Dataset dataset, List rewriteResults, Optional targetRowsPerFragment,