Skip to content

fix: JNI local reference cleanup in JVMClasses::with_env#4225

Merged
mbutrovich merged 2 commits intoapache:mainfrom
0lai0:feat-4172-JNIlocal
May 5, 2026
Merged

fix: JNI local reference cleanup in JVMClasses::with_env#4225
mbutrovich merged 2 commits intoapache:mainfrom
0lai0:feat-4172-JNIlocal

Conversation

@0lai0
Copy link
Copy Markdown
Contributor

@0lai0 0lai0 commented May 5, 2026

Which issue does this PR close?

Closes #4172

Rationale for this change

JVMClasses::with_env attaches native worker threads to the JVM and then runs JNI call-site closures directly. Executor worker threads can stay attached for the JVM lifetime, so JNI local references created inside those closures may accumulate instead of being released at each call boundary.

Wrapping each with_env closure in a JNI local frame gives all existing call sites a consistent push/pop boundary and prevents long-running executor threads from retaining local references across batches.

What changes are included in this PR?

  • Wrap JVMClasses::with_env closures with Env::with_local_frame.
    • Preserve the existing E: From<CometError> caller contract by adapting local-frame JNI errors internally.
    • Document that callers must not return JNI local references from with_env.
    • Audited existing callers: scan, shuffle scan, scalar subquery, memory pools, and parquet encryption convert JNI values to Rust-owned values before returning.

How are these changes tested?

  • cargo fmt -p datafusion-comet-jni-bridge
  • cargo check -p datafusion-comet-jni-bridge
  • cargo test -p datafusion-comet-jni-bridge
  • cargo check -p datafusion-comet --all-targets
  • cargo clippy -p datafusion-comet-jni-bridge --all-targets -- -D warnings
  • cargo clippy -p datafusion-comet --lib -- -D warnings
  • JAVA_TOOL_OPTIONS="-Xcheck:jni" ./mvnw test -Dtest=none -Dsuites="org.apache.comet.exec.CometNativeShuffleSuite" -Pspark-3.5

@mbutrovich mbutrovich self-requested a review May 5, 2026 12:19
@mbutrovich
Copy link
Copy Markdown
Contributor

Thanks @0lai0! Starting CI and will review today.

Copy link
Copy Markdown
Contributor

@mbutrovich mbutrovich left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks right to me, thanks for the quick fix @0lai0! Approved pending CI.

@mbutrovich mbutrovich merged commit 3fc301a into apache:main May 5, 2026
156 checks passed
@0lai0
Copy link
Copy Markdown
Contributor Author

0lai0 commented May 5, 2026

Thanks @mbutrovich for reviews. : )

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.

JNI local references accumulate across executor JVM lifetime in native call sites

2 participants