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
Which issue does this PR close?
Rationale for this change
SessionContext and DataFrame each held their native pointer in a plain
long nativeHandle, and every public method followed the pattern
The check and the use are separate operations on a non-volatile field, so a
close() on another thread can free the native Box in between: thread A reads
a live handle and enters JNI, thread B zeroes the field and Rust drops the
Box, and thread A dereferences freed memory. The == 0 guard is a TOCTOU, not
a 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 SessionContext and DataFrame are both Send + Sync.
Exactly four entry points take ownership and free: collectDataFrame,
executeStreamDataFrame, closeDataFrame, closeSessionContext. So only the
handle'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 NativeHandle owns the raw pointer and is the only thing
that reads or writes it:
pins drain, and surrender the pointer for freeing or consumption.
Every non-consuming method became acquire() / try / finally release().
DataFrame.collect and executeStream use claim(); both close() methods use
claimQuietly().
Two properties drove the design:
lets the eight set operations and join / joinOn pin two DataFrames at once
without a lock-ordering hazard. A ReentrantReadWriteLock (suggested on the
issue) would not: its readers queue behind a waiting writer, so a.union(b)
and b.union(a) on two threads with close()s interleaved could deadlock.
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 Arc refcounting listed as option 2 on the issue is not needed
given 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):
isolation: pin counting, acquire() after claim(), claim() blocking until
a latched pin releases, one winner among racing claimers, and interruption
during the drain being absorbed with the flag restored.
(held open by a TableProvider whose schema() parks on a latch), several
threads querying while another closes, and concurrent close() idempotence.
winner, close() waiting on an in-flight execution (latched scan),
concurrent non-consuming reads all succeeding, and opposing union orders
across 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, including
both integration-level close() tests, and pass again once restored.
The full suite is green (360 tests), as are spotless:check and cargo 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.md and
docs/source/user-guide/sessioncontext.md are updated accordingly:
object; it is a no-op after the first call, including concurrently.
instead 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.