Skip to content

fix: serialise native handle lifetime against concurrent close - #115

Open
andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:fix/handle-close-race
Open

fix: serialise native handle lifetime against concurrent close#115
andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:fix/handle-close-race

Conversation

@andygrove

Copy link
Copy Markdown
Member

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

if (nativeHandle == 0) throw new IllegalStateException(...);
someNativeCall(nativeHandle, ...);

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:

  • acquire() / release() pin the handle for the duration of one native call.
  • claim() / claimQuietly() take exclusive ownership, block until in-flight
    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:

  • acquire() never blocks — it pins immediately or throws. That is what
    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.
  • 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 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):

  • NativeHandleTest — pure JVM, exercises the lifetime state machine in
    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.
  • SessionContextConcurrencyTestclose() waiting on an in-flight call
    (held open by a TableProvider whose schema() parks on a latch), several
    threads querying while another closes, and concurrent close() idempotence.
  • DataFrameConcurrencyTest — racing collect() resolving to exactly one
    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:

  • Both classes are now safe to share between threads.
  • close() blocks until in-flight calls return, then releases the native
    object; it is a no-op after the first call, including concurrently.
  • A call that loses the race to a close() throws IllegalStateException
    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.

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bug: SessionContext.close() / DataFrame.close() race with concurrent JNI calls (use-after-free)

1 participant