Skip to content

Commit

Permalink
refactor: Remove method get_global_jclass (#580)
Browse files Browse the repository at this point in the history
* refactor: Remove method get_global_jclass

The method uses unsafe code to cast a value to a static lifetime using
std::mem::transmute. But none of users of the method actually seems to
depend on this lifetime extention. So to me it seems better to just call
the underlying `find_class` method and delete the unsafe code.

* Apply suggestions from code review

Remove no longer needed comments.

Co-authored-by: Liang-Chi Hsieh <viirya@gmail.com>

---------

Co-authored-by: Liang-Chi Hsieh <viirya@gmail.com>
  • Loading branch information
eejbyfeldt and viirya authored Jun 21, 2024
1 parent 28309a4 commit 9f44a74
Show file tree
Hide file tree
Showing 5 changed files with 6 additions and 30 deletions.
4 changes: 1 addition & 3 deletions core/src/jvm_bridge/batch_iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.

use super::get_global_jclass;
use jni::{
errors::Result as JniResult,
objects::{JClass, JMethodID},
Expand All @@ -34,8 +33,7 @@ impl<'a> CometBatchIterator<'a> {
pub const JVM_CLASS: &'static str = "org/apache/comet/CometBatchIterator";

pub fn new(env: &mut JNIEnv<'a>) -> JniResult<CometBatchIterator<'a>> {
// Get the global class reference
let class = get_global_jclass(env, Self::JVM_CLASS)?;
let class = env.find_class(Self::JVM_CLASS)?;

Ok(CometBatchIterator {
class,
Expand Down
5 changes: 1 addition & 4 deletions core/src/jvm_bridge/comet_exec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ use jni::{
JNIEnv,
};

use super::get_global_jclass;

/// A struct that holds all the JNI methods and fields for JVM CometExec object.
pub struct CometExec<'a> {
pub class: JClass<'a>,
Expand Down Expand Up @@ -55,8 +53,7 @@ impl<'a> CometExec<'a> {
pub const JVM_CLASS: &'static str = "org/apache/spark/sql/comet/CometScalarSubquery";

pub fn new(env: &mut JNIEnv<'a>) -> JniResult<CometExec<'a>> {
// Get the global class reference
let class = get_global_jclass(env, Self::JVM_CLASS)?;
let class = env.find_class(Self::JVM_CLASS)?;

Ok(CometExec {
method_get_bool: env
Expand Down
5 changes: 1 addition & 4 deletions core/src/jvm_bridge/comet_metric_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ use jni::{
JNIEnv,
};

use super::get_global_jclass;

/// A struct that holds all the JNI methods and fields for JVM CometMetricNode class.
pub struct CometMetricNode<'a> {
pub class: JClass<'a>,
Expand All @@ -37,8 +35,7 @@ impl<'a> CometMetricNode<'a> {
pub const JVM_CLASS: &'static str = "org/apache/spark/sql/comet/CometMetricNode";

pub fn new(env: &mut JNIEnv<'a>) -> JniResult<CometMetricNode<'a>> {
// Get the global class reference
let class = get_global_jclass(env, Self::JVM_CLASS)?;
let class = env.find_class(Self::JVM_CLASS)?;

Ok(CometMetricNode {
method_get_child_node: env
Expand Down
4 changes: 1 addition & 3 deletions core/src/jvm_bridge/comet_task_memory_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ use jni::{
JNIEnv,
};

use crate::jvm_bridge::get_global_jclass;

/// A wrapper which delegate acquire/release memory calls to the
/// JVM side `CometTaskMemoryManager`.
#[derive(Debug)]
Expand All @@ -40,7 +38,7 @@ impl<'a> CometTaskMemoryManager<'a> {
pub const JVM_CLASS: &'static str = "org/apache/spark/CometTaskMemoryManager";

pub fn new(env: &mut JNIEnv<'a>) -> JniResult<CometTaskMemoryManager<'a>> {
let class = get_global_jclass(env, Self::JVM_CLASS)?;
let class = env.find_class(Self::JVM_CLASS)?;

let result = CometTaskMemoryManager {
class,
Expand Down
18 changes: 2 additions & 16 deletions core/src/jvm_bridge/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,8 @@
use crate::errors::CometResult;

use jni::{
errors::{Error, Result as JniResult},
objects::{JClass, JMethodID, JObject, JString, JThrowable, JValueGen, JValueOwned},
errors::Error,
objects::{JMethodID, JObject, JString, JThrowable, JValueGen, JValueOwned},
signature::ReturnType,
AttachGuard, JNIEnv,
};
Expand Down Expand Up @@ -176,20 +176,6 @@ pub(crate) use jni_new_string;
pub(crate) use jni_static_call;
pub(crate) use jvalues;

/// Gets a global reference to a Java class.
pub fn get_global_jclass(env: &mut JNIEnv, cls: &str) -> JniResult<JClass<'static>> {
let local_jclass = env.find_class(cls)?;
let global = env.new_global_ref::<JObject>(local_jclass.into())?;

// A hack to make the `JObject` static. This is safe because the global reference is never
// gc-ed by the JVM before dropping the global reference.
let global_obj = unsafe { std::mem::transmute::<_, JObject<'static>>(global.as_obj()) };
// Prevent the global reference from being dropped.
let _ = std::mem::ManuallyDrop::new(global);

Ok(JClass::from(global_obj))
}

mod comet_exec;
pub use comet_exec::*;
mod batch_iterator;
Expand Down

0 comments on commit 9f44a74

Please sign in to comment.