Struggling with API changes from 0.20.0 to 0.20.1

[amunra]

Hello,

I'm migrating from 0.20.0 to 0.21.0 and the newly introduced API changes lead me to introduce more unsafe code to get things working and generally feel less ergonomic.

In other words, I'm probably doing something wrong and I'd appreciate some assistance :-).

For reference, here's a link to the code I'm discussing: https://github.com/questdb/questdb/blob/43a91d56453946ac4e4a70ec495306441f9cdafa/core/src/main/rust/questdb-jni/src/log.rs#L45

0.20.0 API

I'm caching class and method lookups so these can be used across threads to call from Rust back into Java: I'm bridging Rust log to our internal logging framework and I don't want to perform these lookups repeatedly.

During initialization, I have some code that looks up a class and stores it as a GlobalRef so I can re-use it later from another thread:

let log_factory_class = env.find_class("io/questdb/log/LogFactory")?;
let log_factory_class = env.new_global_ref(log_factory_class)?;

Later, from another thread, I'm using the looked-up information to call a static method.
Note that the thread in question is a thread that was started by Rust (not from a JNI call), so I'm also using AutoLocal ensure my local references are properly cleaned up.

let j_name = jenv.new_string(name)
    .expect("could not construct logger name string");
let j_name = jenv.auto_local(j_name);
let obj = jenv.call_static_method_unchecked(
    call_state.log_factory_class.as_obj(),
    call_state.log_factory_get_log_meth,
    ReturnType::Object,
    &[JValue::Object(j_name.as_obj()).into()])
    .expect("io.questdb.log.LogFactory::getLog(String) call failed")
    .l()
    .expect("io.questdb.log.LogFactory::getLog(String) didn't return an object");

0.21.0 API

The call_static_method_unchecked signature has changed and I need to dance around it with unsafe code for both JClass (from a GlobalRef) and JString (from an AutoLocal<JString>).

let j_name = jenv.new_string(name)
    .expect("could not construct logger name string");
let j_name = jenv.auto_local(j_name);

// Note: `call_state.log_factory_class.as_obj()` returns a `&JObject` whilst I need a JObject.
let j_class_obj: JObject = unsafe { JObject::from_raw(call_state.log_factory_class.as_obj().as_raw()) };
let j_class = JClass::from(j_class_obj);
let j_name_raw: jobject = j_name.as_raw();
let j_name_obj: JObject = unsafe { JObject::from_raw(j_name_raw) };
let j_name_value = JValue::Object(&j_name_obj);
let obj = unsafe { jenv.call_static_method_unchecked(
    j_class,
    call_state.log_factory_get_log_meth,
    ReturnType::Object,
    &[j_name_value.as_jni()])}
    .expect("io.questdb.log.LogFactory::getLog(String) call failed")
    .l()
    .expect("io.questdb.log.LogFactory::getLog(String) didn't return an object");

This is significantly more complicated and requires unsafe in a few places.

My best guess is that I mis-understood how to use the API somewhere and I'm doing something wrong.

Any advice here is wecome:
What's the proper way of caching class/method lookups and using AutoLocal from another (non-java) thread?

Thanks!

[fpoli]

Have you tried directly passing a &call_state.log_factory_class as first parameter to jenv.call_static_method_unchecked(..) instead of going through .as_obj().as_raw()...? From the documentation it seems that a &GlobalRef should be accepted because of this impl<'local, 'obj_ref> Desc<'local, JClass<'static>> for &'obj_ref GlobalRef.

[amunra]

That worked, thanks!

[fpoli]

Also, let j_name_value = JValue::Object(j_name.as_ref()) should work because of fn as_ref(&self) -> &JObject<'local>.

[amunra]

Also worked, thank you!

[rib]

This is significantly more complicated and requires unsafe in a few places

Yeah, unfortunately we've repeatedly found that the JNI API has exposed unsafe APIs as if they are safe and 0.21 is just being a bit more honest about the safety. So actually these things were unsafe in 0.20 but it wasn't obvious. In some cases that does make things less ergonomic but I think it's probably better to know when things are unsafe than not. This way it's also clearer where more work is needed to provide a genuinely safe API.

Note that the thread in question is a thread that was started by Rust (not from a JNI call), so I'm also using AutoLocal ensure my local references are properly cleaned up.

One API that got some improvements in 0.21 was JNIEnv::with_local_frame() which you may want to investigate using instead of creating explicit AutoLocal wrappers. Potentially this could be an ergonomic improvement for your code.

In general I think you should only really need to use AutoLocal if you have a situation where you need to make sure that you delete certain local references before exiting your current local frame (such as a loop that generates a huge number of temporary references).

[amunra]

Very nice! with_local_frame is pretty neat and simplifies my code a fair bit.

If I can waste a bit more of your generous time, I had a few questions:

(1) Out of curiosity, what is the main motivation behind changing JNIEnv to require mut?

My source of confusion here is that vm.attach_current_thread_permanently() allows - at any stage - to get a new equivalent JNIEnv object, kinda invalidating the mut requirement.

(2) Related to (1), I've got a little helper function to get a new JNIEnv for the current (non Java) thread.

static mut JAVA_VM: Option<JavaVM> = None;

fn get_jenv() -> Option<jni::JNIEnv<'static>> {
    if let Some(vm) = unsafe { &JAVA_VM } {
        let j_env = vm.attach_current_thread_permanently()
            .expect("failed to attach jni env");
        Some(j_env)
    }
    else {
        None
    }
}

Would it be sensible to cache the JNIEnv in a thread local and avoid internally calling into AttachCurrentThread repeatedly -- I'm a bit fuzzy on the safety of doing this (context: JAVA_VM is only set once at startup immediately after loading the jni dynamic library with System.load).

(3) I've ended up implementing a pretty ugly macro for cases where a jni::errors::Result error was triggered by a Java exception.

#[macro_export]
macro_rules! unwrap_or_throw {
    ($env:expr, $result:expr, $fallback_java_exception_class:literal $(, $return_sentinel:literal)?) => {
        match $result {
            Ok(result) => result,
            Err(e) => {
                let throwable = $env.exception_occurred().unwrap();
                if throwable.is_null() {
                    $env.throw_new($fallback_java_exception_class, &format!("{}", e))
                        .expect("failed to throw java exception");
                }
                return $($return_sentinel)?;
            }
        }
    };
    ($env:expr, $result:expr) => {
        unwrap_or_throw!($env, $result, "java/lang/RuntimeException")
    };
}

Usage example, with called function returning a Result and macro converting it to an exception:

#[no_mangle]
pub extern "system" fn Java_io_questdb_std_Os_initQuestdbJni(mut env: JNIEnv, _class: JClass) {
    let vm = unwrap_or_throw!(env, env.get_java_vm());
    unsafe { JAVA_VM = Some(vm); }
}

Is there a better way to figure out if a jni error result carries an exception with it and generally bridge jni crate errors to java exceptions?

(4) Shameless plug

I've also started a Maven plugin to build Rust code within a Java project: https://github.com/questdb/rust-maven-plugin. The use case is obviously to make it easier to use the jni crate. It fits our own use case, but I expect it evolve much beyond our needs.

[fpoli]

My source of confusion here is that vm.attach_current_thread_permanently() allows - at any stage - to get a new equivalent JNIEnv object, kinda invalidating the mut requirement.

That's indeed suspicious. Maybe the refactoring to 0.21.0 forgot about that method?

[rib]

Out of curiosity, what is the main motivation behind changing JNIEnv to require mut

The motivation here was generally related to making JNI references !Copy and making it much harder to access invalid Java reference pointers after they are freed (#398 and #392)

In particular see #392 (comment)

This is mainly to deal with situations where you might end up with more than one JNIEnv in scope because you have used one to call with_local_frame() and you are then passed a new reference to a JNIEnv that is associated with your new local frame.

In this situation we want to ensure that you don't try to continue using the outer env to create local references that have a Rust lifetime that ties the references to the wrong (outer) local frame.

The general rule is that any method that can create a Java reference needs to take a mutable reference to the environment and then because a new with_local_frame scope borrows an environment mutably you will be blocked from inadvertently associating local references with the lifetime of an outer environment.

Using vm.attach_current_thread_permanently() can currently sidestep that protection, certainly, but to really shoot yourself in the foot I think you'd have to create multiple JNIEnv references in an outer scope before using something like with_local_frame to create an inner frame/scope and then you would have to use the redundant outer env that you didn't borrow mutably to then create a type like JObject that would be associated with the lifetime of the outer environment. So although that's possible with the current API it's still hopefully harder to shoot yourself in the foot than it was the the 0.20 API.

[rib]

My source of confusion here is that vm.attach_current_thread_permanently() allows - at any stage - to get a new equivalent JNIEnv object, kinda invalidating the mut requirement.

That's indeed suspicious. Maybe the refactoring to 0.21.0 forgot about that method?

We can't really claim to have done a holistic refactor during 0.21. Although it looks like quite a lot of changes, they mainly came about from working through a few specific issues (mainly around !Copy reference types, adding wrapper types for array references). So the changes for 0.21 only really made incremental progress in terms of making the jni crate safer. Certain corners of the API (like thread attachment) didn't really get any significant attention in the last release.

Although I think we did talk briefly about bigger changes to how an env can be acquired - I would have also been very cautious about biting off more than we can chew - especially considering how long the jni project was inactive for.

I expect that there's room for improvement around the thread attachment APIs (e.g. "daemon" based attachments are quite a bad fit for Rust safety).

One of the ideas I'd like to explore further is whether we could have separate types to differentiate a transparent wrapper around a raw JNIRawEnv pointer from some kind of JNIEnv that can be associated with a minimal internal cache. I'm interested in seeing if we could have a stateful env type that would give us a place to cache method IDs for common types like lists or maps to be able to have more efficient and ergonomic proxy APIs.

Investigating that further could potentially help tighten up this kind of loop hole where code can effectively create an env out of thin air and could additionally lead to nicer ergonomics too.

This kind of change would also need to tie into all the various ways that you can access an env reference, including these attachment APIs but also potentially a catch_unwind utility like this for native methods: #432

[argv-minus-one]

(1) Out of curiosity, what is the main motivation behind changing JNIEnv to require mut?

My source of confusion here is that vm.attach_current_thread_permanently() allows - at any stage - to get a new equivalent JNIEnv object, kinda invalidating the mut requirement.

Yes, the JNIEnv-returning methods of JavaVM are a remaining soundness hole in this library. To make them sound, they would have to either be unsafe or fail if the thread is already attached. @rib, do you think we should do that?

(2) Related to (1), I've got a little helper function to get a new JNIEnv for the current (non Java) thread.

static mut JAVA_VM: Option<JavaVM> = None;

fn get_jenv() -> Option<jni::JNIEnv<'static>> {
    if let Some(vm) = unsafe { &JAVA_VM } {
        let j_env = vm.attach_current_thread_permanently()
            .expect("failed to attach jni env");
        Some(j_env)
    }
    else {
        None
    }
}

Would it be sensible to cache the JNIEnv in a thread local and avoid internally calling into AttachCurrentThread repeatedly -- I'm a bit fuzzy on the safety of doing this (context: JAVA_VM is only set once at startup immediately after loading the jni dynamic library with System.load).

I don't think that's normally possible, since JNIEnv has a lifetime parameter. It isn't possible to insert anything non-'static into a thread-local variable, is it?

The JNI spec kinda-sorta implies that you're allowed to keep a JNIEnv. There is a safety issue, though: if the thread is detached, the JNIEnv becomes invalid, and using it after that causes undefined behavior.

You might want to cache the JavaVM instead and use JavaVM::get_env when needed. This will avoid the above problems.

Is there a better way to figure out if a jni error result carries an exception with it and generally bridge jni crate errors to java exceptions?

To see if the result carries a Java exception, check if it's Err(jni::errors::Error::JavaException).

We don't yet have anything for converting the other jni-rs errors into Java exceptions. I worked on that for #432, but I'm waiting to hear back from @jrobsonchase on a question I posted there.

[amunra]

Thanks for the insightful details.

[rib]

Yes, the JNIEnv-returning methods of JavaVM are a remaining soundness hole in this library. To make them sound, they would have to either be unsafe or fail if the thread is already attached. @rib, do you think we should do that?

apologies for the delay replying here - I might be a bit busy to follow up on this closely for a while, sorry.

I could have sworn we had some discussion around this before but wasn't able to dig it up.

Making them unsafe feels a bit frustrating but yeah it's an option.

If we're going to break the thread attachment APIs, I think I'd like to at least do a minimal exploration of my idea to introduce a stateful JNIEnv in some form - from the pov of trying to avoid multiple breakages across the same APIs.

Failing if the thread is already attached initially seems like that could be a big pain but maybe in combination with an unsafe alternative it would be OK.

I can imagine that some crates might want to hide the details of JNI from their users (apart from possibly being initialized with a JVM pointer), and so then there could be various APIs that want to use JNI internally that haven't been directly passed a JNIEnv. The implementation of these APIs would need a way to access a JNIEnv - even though it's possible that the thread has been attached to the JVM already.

It's perhaps reasonable for that use case that it would require an unsafe API where you just have to promise that there's no other JNIEnv in scope.  This kind of crate would use this unsafe API at its frontend and would then pass the acquired JNIEnv wherever it's needed internally.

On a related note also see this issue with thread attachment that was filed recently: #441

[fpoli]

Related to the "how to use JNI 0.21" discussion, I recently noticed that nikomatsakis is working on duchess; a codegen to call Java code from Rust. They explain some ideas of the implementation here. The abstraction that they build around the JVM looks intersting; probably something to consider when designing future changes to the API of jni-rs.

[rib]

Thanks, yeah I started having a nose through duchess the other day. One thing that stood out as surprising was that it currently looks like it's going to lean heavily on AutoLocal (essentially renamed as Local) which didn't sound ideal.

I don't think Java local references are really designed / intended to be explicitly deleted as a matter of course vs letting the local frame unwind and release local references. If AutoLocal gets used for almost all local references then I'd expect that will lead to a lot of redundant bookkeeping to explicitly remove local references from an internal data structure that's quite likely not optimized for that usage model.

[argv-minus-one]

For reference, here are the results of the benchmark I did last year.