fix: serialise native handle lifetime against concurrent close - #115
Open
andygrove wants to merge 1 commit into
Open
fix: serialise native handle lifetime against concurrent close#115andygrove wants to merge 1 commit into
andygrove wants to merge 1 commit into
Conversation
SessionContext and DataFrame each held a raw `long nativeHandle` guarded by `if (nativeHandle == 0) throw`. The check and the JNI call were separate operations, so a concurrent close() could free the native Box between them and leave the other thread dereferencing freed memory. Introduce a package-private NativeHandle that owns the pointer. Every call pins the handle for its duration (acquire/release); close() and the consuming operations claim it, waiting for in-flight calls to drain before handing the pointer back to Rust to be freed. acquire() never blocks, which is what lets the two-handle set operations and joins pin both DataFrames without a lock-ordering hazard. A read/write lock would not: its readers queue behind a waiting writer, so opposing pin orders with closes interleaved could deadlock. claim() is the only blocking operation and always makes progress, since a thread holding a pin is inside a native call and never claims. No native change is needed. Every non-consuming JNI entry point already takes a shared reference, and DataFusion's SessionContext and DataFrame are Send + Sync, so only the handle's lifetime had to be serialised. Closes apache#40
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.
Which issue does this PR close?
Rationale for this change
SessionContextandDataFrameeach held their native pointer in a plainlong nativeHandle, and every public method followed the patternThe check and the use are separate operations on a non-volatile field, so a
close()on another thread can free the nativeBoxin between: thread A readsa live handle and enters JNI, thread B zeroes the field and Rust drops the
Box, and thread A dereferences freed memory. The== 0guard is a TOCTOU, nota fix.
The docs said "not thread-safe", which made this defensible but not
comfortable — the first multi-threaded user (a server, a Flink or Spark
integration) would have hit it in production rather than in dev. It is also a
shape that recurs every time a new long-lived handle is added.
Reading the native side showed the problem is entirely Java-side. Every
non-consuming JNI entry point already takes a shared reference —
&*(handle as *const SessionContext)or&*(handle as *const DataFrame)}.clone()— and DataFusion's
SessionContextandDataFrameare bothSend + Sync.Exactly four entry points take ownership and free:
collectDataFrame,executeStreamDataFrame,closeDataFrame,closeSessionContext. So only thehandle's lifetime needed serialising; the operations themselves are already
safe to overlap, and no Rust change is required.
What changes are included in this PR?
A new package-private
NativeHandleowns the raw pointer and is the only thingthat reads or writes it:
acquire()/release()pin the handle for the duration of one native call.claim()/claimQuietly()take exclusive ownership, block until in-flightpins drain, and surrender the pointer for freeing or consumption.
Every non-consuming method became
acquire()/ try /finally release().DataFrame.collectandexecuteStreamuseclaim(); bothclose()methods useclaimQuietly().Two properties drove the design:
acquire()never blocks — it pins immediately or throws. That is whatlets the eight set operations and
join/joinOnpin two DataFrames at oncewithout a lock-ordering hazard. A
ReentrantReadWriteLock(suggested on theissue) would not: its readers queue behind a waiting writer, so
a.union(b)and
b.union(a)on two threads withclose()s interleaved could deadlock.claim()is the only blocking operation, and it always makes progress —a thread holding a pin is inside a native call and never claims, so the pins
it waits on are guaranteed to be released.
The monitor is held only for bookkeeping, never across a JNI call, so
independent operations on the same object still run concurrently.
The Rust-side
Arcrefcounting listed as option 2 on the issue is not neededgiven the above, and would have meant touching every JNI entry point.
Argument validation stays in its existing position relative to the closed check
within each method, so no currently-observable exception changes.
Are these changes tested?
Yes, three new suites (18 tests):
NativeHandleTest— pure JVM, exercises the lifetime state machine inisolation: pin counting,
acquire()afterclaim(),claim()blocking untila latched pin releases, one winner among racing claimers, and interruption
during the drain being absorbed with the flag restored.
SessionContextConcurrencyTest—close()waiting on an in-flight call(held open by a
TableProviderwhoseschema()parks on a latch), severalthreads querying while another closes, and concurrent
close()idempotence.DataFrameConcurrencyTest— racingcollect()resolving to exactly onewinner,
close()waiting on an in-flight execution (latchedscan),concurrent non-consuming reads all succeeding, and opposing
unionordersacross threads not deadlocking.
I verified the tests actually detect the bug by temporarily removing the drain
loop from
NativeHandle: exactly the four drain-dependent tests fail, includingboth integration-level
close()tests, and pass again once restored.The full suite is green (360 tests), as are
spotless:checkandcargo fmt.No Rust files are touched.
Are there any user-facing changes?
No API changes, but the documented contract is stronger, and the Javadoc plus
README.md,docs/source/user-guide/quickstart.mdanddocs/source/user-guide/sessioncontext.mdare updated accordingly:close()blocks until in-flight calls return, then releases the nativeobject; it is a no-op after the first call, including concurrently.
close()throwsIllegalStateExceptioninstead of producing a use-after-free.
The docs are explicit that this covers the handle's lifetime and not the
ordering of overlapping operations — registering a table concurrently with a
query that reads it still races in the ordinary way.