Nested "critical sections" disallowed by borrow checker

[onsmith]

Description

I have a Rust library which uses this crate to access two Java arrays at once through nested critical sections. This behavior is explicitly permitted by the JNI spec:

Multiple pairs of GetPrimtiveArrayCritical and ReleasePrimitiveArrayCritical may be nested. For example:

  jint len = (*env)->GetArrayLength(env, arr1);
  jbyte *a1 = (*env)->GetPrimitiveArrayCritical(env, arr1, 0);
  jbyte *a2 = (*env)->GetPrimitiveArrayCritical(env, arr2, 0);
  /* We need to check in case the VM tried to make a copy. */
  if (a1 == NULL || a2 == NULL) {
    ... /* out of memory exception thrown */
  }
  memcpy(a1, a2, len);
  (*env)->ReleasePrimitiveArrayCritical(env, arr2, a2, 0);
  (*env)->ReleasePrimitiveArrayCritical(env, arr1, a1, 0);

As first predicted by @rib in this comment, the following minimal example no longer compiles as of version 0.21.0, due to the &mut reference to JNIEnv held by get_array_elements_critical():

// Prevent Rust from mangling the function name
#[no_mangle]
pub extern "system" fn Java_example_JNIExample_nestedCriticalSectionsDemo<'local>(
    mut env: JNIEnv<'local>,
    _class: JClass<'local>,
    src: JByteArray<'local>,
    dst: JByteArray<'local>,
) -> jlong {
    // ...

    {
        // critical section starts here
        let src_slice = unsafe { env.get_array_elements_critical(&src, ReleaseMode::NoCopyBack) };
        let dst_slice = unsafe { env.get_array_elements_critical(&dst, ReleaseMode::CopyBack) };

        // perform operations on src_slice and dst_slice, such as copy
        // note: this is for illustration, the real code handles unwrapping and index
        // bounds checking safely without panics
        dst_slice.unwrap().copy_from_slice(&src_slice.unwrap());

        // exit the critical section by dropping src_slice and dst_slice
    }

    // ...

    return 0;
}

Workaround

My current workaround is to unsafe_clone() the JNIEnv inside the critical section block and use the cloned object for one of the critical sections:

    {
        let mut env2 = unsafe { env.unsafe_clone() };

        // critical section starts here
        let src_slice = unsafe { env.get_array_elements_critical(&src, ReleaseMode::NoCopyBack) };
        let dst_slice = unsafe { env2.get_array_elements_critical(&dst, ReleaseMode::CopyBack) };

        // ...
    }

Questions

1. Is the workaround described above the safest approach for achieving the desired behavior with version 0.21.0 of this crate?
2. Should guidance about how to handle this situation be added to the documentation or migration guide? This is a valid use case permitted by the JNI spec that needs special consideration for folks migrating to 0.21.0.
3. In the longer term, should developers using this crate anticipate a future change to the API for functions like get_array_elements_critical() to make situations like this one work straight out of the box?

Thank you!

[rib]

Right, it's been really fiddly trying to get some semblance of safety into this corner of the API and the pass we made over this for 0.21 did have the unfortunate side effect of making nested critical sections difficult.

Yeah your workaround makes sense. It's a risky footgun in the sense that you can now freely use JNI which is not safe to do inside a critical section, but if you're only using it to create a nested critical section then hopefully that's OK. I don't think there's really any safer workaround currently.

This exception to the rule that you can't use JNI within a critical section but you can create a nested critical section is fiddly to ensure.

One option could be that we just remove the mutable borrow and extend the long list of safety issues you're responsible for, but my general hope was that we might be able to reduce the number of safety issues that need to be dealt with manually.

Maybe we could actually add an API to AutoElementsCritical (which is holding a mutable borrow of the env) that would let you create a nested critical section from an existing AutoElementsCritical?

It's maybe a bit funky that you'd have two APIs for essentially the same thing, one via JNIEnv and then one via AutoElementsCritical for nesting, but that could let us enforce this awkward rule via the type system fairly neatly I suppose.

In terms of longer term changes - I think it's likely that this corner of the API should evolve still, while there is currently a fairly unwieldy list of documented safety issues that could be good to try and reduce.

With the 0.21 release we had already accumulated quite a few breaking changes in the name of improving safety and I wasn't looking to boil the ocean trying to address every issue that we found when looking at these APIs, and dragging out the release too long.

There was some discussion that we could consider introducing a closure for the critical section to help us tackle some of the other safety issues here (ref: https://github.com/jni-rs/jni-rs/pull/400#issuecomment-1404319714) which would probably be the most significant change I can think of. This isn't being actively looked at yet though, and isn't really a priority for me atm.

If we were to add some way to nest a critical section via AutoElementsCritical then I guess that would essentially remain compatible with any change to introduce a closure for the critical section.

I think I'd like to put back the mutable borrow of the array itself too if we can implement AsMut<JPrimitiveArray> via a GlobalRef that doesn't erase the type of reference it holds.

Happy to hear if you have other thoughts / ideas for iterating this API based on your usage?