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/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..4ac44b2a2b4 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( @@ -72,7 +75,48 @@ public static CompactionMetrics commitCompaction( compactionOptions.getMaxSourceBytes()); } - public static native CompactionMetrics nativeCommitCompaction( + /** + * 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 commitCompactionNative( + dataset, + rewriteResults, + targetRowsPerFragment, + maxRowsPerGroup, + maxBytesPerFile, + materializeDeletions, + materializeDeletionsThreshold, + numThreads, + batchSize, + deferIndexRemap, + compactionMode, + binaryCopyReadBatchBytes, + maxSourceFragments, + maxSourceRows, + maxSourceBytes); + } + } + + private static native CompactionMetrics commitCompactionNative( Dataset dataset, List rewriteResults, Optional targetRowsPerFragment, 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(); + } + } + } }