fix(java): hold Dataset read lock when passing Dataset into native calls - #8575
Open
zhangyue19921010 wants to merge 2 commits into
Open
fix(java): hold Dataset read lock when passing Dataset into native calls#8575zhangyue19921010 wants to merge 2 commits into
zhangyue19921010 wants to merge 2 commits into
Conversation
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".
…ad 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.
Contributor
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The remaining public JNI bypass is fixed: commit compaction preserves its public signature behind the dataset read-lock boundary and keeps the raw native entry point private. The change now covers the reported close/use race at the Java/native boundary.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Datasetguards its native handle with aReentrantReadWriteLock: its own methods take the read lock, andclose()takes the write lock. However, entry points outside the class pass theDatasetobject into JNI without acquiring this lock:Fragment:newScan/deleteRows/countRows/mergeColumns/updateColumnsLanceScanner.create/AsyncScanner.createSqlQuery.intoBatchRecordsCommitBuilder.execute(dataset branch)Compaction.planCompaction/commitCompaction,CompactionTask.executeVectorTrainer.trainIvfCentroids/trainPqCodebookDatasetDeltaBuilder.buildShardWriter.create,LsmScanner.fromSnapshots,LsmPointLookupPlanner/LsmVectorSearchPlannerconstructorsSo
Dataset.close()is not mutually exclusive with these calls. A concurrentclose()hits a flaw in jni-rstake_rust_field: if another thread holds the handle mutex,try_lockfails and the boxedBlockingDatasetis freed on the early-return path — a native use-after-free. Under contention (e.g. 16 threadsFragment.countRows()vsclose()) this aborts the JVM (SIGABRT, malloc corruption); it also leaves both the Rust field and the Java handle in a corrupt state, crashing later unrelated calls. Any JVM sharing oneDatasetacross threads with a close/eviction path (e.g. Spark executors) is affected.Fix
Follow the existing locking mechanism instead of introducing a new one:
Dataset.acquireReadLock(): acquires the shared read lock and rejects closed datasets, mirroring thereadLock+handle != 0checkDataset's own methods already use.Datasetintry (LockManager.ReadLock readLock = dataset.acquireReadLock()) { ... }.close()(write lock) now waits for in-flight native calls; callers racing withclose()either complete safely or fail cleanly withIllegalArgumentException("Dataset is closed"). The read lock is reentrant, so nested paths likeDataset.newScan→LanceScanner.createare safe.