Inclusion of helpers for language-specific tweaks

[dnaka91]

Include some helpers that glue over the percularities of language-specific behavior in uniffi, specifically for the Kotlin/JVM target.

Background

I just recently moved a project at work over from plain JNI to using uniffi for a Rust/Android mixed project. It seemingly worked great at first, but it was only later that I discovered a few footguns here and there.

One of them was the attachment to the current thread for JVM targets, when spawning any kind of threads from the Rust side. In my case a multi-threaded Tokio runtime.

The docs about it were there, but not very prominently and I thought this could be something to be included directly into uniffi as utility functions to gloss over percularities of specific targets like JVM, Python, ...

Proposal

I propose to include some module into the uniffi crate, that would allow to easily attach the thread to the JVM which is what most people would want for better performance, I'd imagine.

This could be a function that is always present, regardless of the target, and only invoked when in a JVM environment.

Dependencies

I think it would be best to directly implement those APIs instead of relying on the jni crate. It brings lots of dependencies with it, and none of them are needed for the JNI APIs in question.

After some experimentation I now have a local minimal version that can attach the JVM to the current thread, only would need to copy it over to the uniffi crate.

Open questions

How to detect the target environment

The first thing that comes to mind are feature flags to enable a certain target environment. I'm using the cargo-ndk crate to build for Android and passing in a feature flag like --features uniffi/kotlin would probably be the easiest.

One danger is of course, when compiling with --all-features to compile for any other language it would at some point try to call some JNI functions that wouldn't be present.

An alternative could be to provide some extra config but invoking that in Cargo is more tedious.

Maybe there is a better way to do this?

Provide a target independent API or specific

Not sure what would be the best way. Either publicly expose those functions behind a feature flag or have an always present API that does its thing for Kotlin, but is a no-op for other languages. However, the latter feels more convenient.

For example:

// Assuming there is a `uniffi::kotlin` module, and the own crate has a `kotlin` feature flag.

#[cfg(feature = "kotlin")]
uniffi::kotlin::attach_current_thread();

// Assuming there is an always present `uniffi::util` module. Does its thing for Kotlin, and no-op for other languages.

uniffi::util::attach();

How to retrieve the JavaVM instance

Basically there are three choices. In any case, the VM instance is global and can be placed in a static variable for the remainder of the program. Only exception to that is probably if the Rust program starts the VM itself instead of being invoked as a library.

• Export a method called JNI_OnLoad, which will be invoked during System.loadLibray from Java. The instance of the VM is passed as first parameter.
  • However, the System.loadLibrary call has to be done manually once, as JNA doesn't trigger the JNI_OnLoad here (I guess because it itself is a shared lib that is having this init callback but then invokes the called library through regular C ABI without JNI involved).
• Export a custom init function that does the same, must be called explicitly as well, and instead get's a JNI env, which requires another call back through the FFI to get the VM.
• Use the JNI_GetCreatedJavaVMs function that can be called in any place, as long as in the context of a JVM. This will give a list of JVM instances but in most cases there should be only one so we can take the first in the list.
  • I personally am a fan of this one, because there is no extra step required that could be forgotten along the way.

Attach permanently or as daemon

There are two ways of attaching to a thread, as non-daemon or as daemon. Permanently is just a term from the jni crate which means, the thread is kept attached until it shuts down (by using a thread local attach guard that'll call detach_current_thread on drop).

The daemon variant has the only difference that it doesn't prevent the JVM from shutting down. The non-daemon variant will keep the JVM running until all attached threads shut down themselves.

When setting up my tokio runtime I first thought the non-daemon variant is the right choice, and it's mentioned in the docs as well. But as the threads are kept alive as long as the runtime exists, it'll block the JVM from shutting down if the runtime is put into a global static variable.

The workaround is to have a custom unload function that will shut down the tokio runtime. However, this is tricky in Android as there is no global shutdown where it could be invoked, and according to Android docs the process is killed in the end anyway, meaning blocking the JVM from shutting down is likely not relevant in this context.

The daemon variant of course doesn't have this problem, but it's unclear if this might cause some issues in tokio if the JVM shuts down threads without its approval. Well, the whole application is shutting down so it shouldn't matter much. But is there a chance some shutdown step might fail due to the threads being stopped this way?

So long story short, it's probably best to expose both ways of attaching to the JVM, with some good docs, and let the user choose based on their requirements.

[mhammond]

That's a very elaborate proposal, but it's difficult to comment without more specifics.

Ultimately though, it sounds like calling back into Kotlin from Rust created threads are causing you problems, right? @bendk don't we have some comments or something around this in app-services?

[bendk]

That's a very elaborate proposal, but it's difficult to comment without more specifics.

Ultimately though, it sounds like calling back into Kotlin from Rust created threads are causing you problems, right? @bendk don't we have some comments or something around this in app-services?

I remember seeing something about this at one point, but I don't remember writing comments about it. I could definitely be forgetting though. @dnaka91 what happens when you try to call a callback from a tokio-spawned thread? Does it just crash?

[mhammond]

Ah, I think I was referring to https://github.com/mozilla/application-services/blob/main/docs/android-faqs.md#what-challenges-exist-when-calling-back-into-kotlin-from-rust - which does describe the problem but doesn't really offer any solutions.

[dnaka91]

I already fixed the problem on my side, and what I mean is the performance issue mentioned in this PR.

It is just a bit hard to discover and I think it's worth a new sample in the repo to show a setup with Tokio and a multi-thread runtime.

What I'm proposing is, to include that attach_current_thread functionality behind a general function (which could potentially include other tweaks).

A benefit would also be, that this can be implemented with zero dependencies actually, as the docs use jni, which brings in several dependencies just for this one workaround of using JNA.

[bendk]

I have to admit, I don't understand the JVM issues here. However, I think we could have a general solution that looked something like this:

• Add uniffi::attach_current_thread to the UniFFI API
• Generate a FFI function that looks like extern "C" pub fn add_attach_current_thread_callback(callback: extern "C" fn(context: void*))
  • The generated bindings would call this function to add a custom callback on startup.
  • Kotlin/Java can use this to attach the thread to the JVM, other languages could also maybe use it. I think we would need to do something like this on Python, if ctypes wasn't implicitly handling it.
  • When users call attach_current_thread, UniFFI will ensure the callback is called.
  • I guess we should allow multiple callbacks

Would this work for you?

[dnaka91]

Basically, but I wonder if JNA even gives you the option to attach the thread to the JVM 🤔.

Ideally JNA would detach a thread so soon again and simply keep it attached but I guess they have their reason.

So the flow would be like:

• Uniffi registers a callback on startup in Kotlin, ending in some call through JNA to register it on Rust side.
• Rust calls uniffi::attach_current_thread which would then iterate over the registered callacks and invoke them, causing a callback through JNA back to Kotlin.

That sounds like quite a lot of passing forth and back between Rust, JNA and Kotlin.

Ideally this would just do a single call all in Rust by directly doing FFI calls into JNI to attach the current thread. As I mentioned, you can get a list of VM instances through JNI, then use the first one to attach the thread.

---

Another thing to consider is to permanently attach or attach as daemon.

This is, again, very JVM specific. I extracted my JNI calls into a minimal setup that would fit well into Uniffi. While trying it out, I noticed that in regular non-Android JVMs you could accidentally make the JVM hang.

The reason here is, that permanently attaching a thread prevents the JVM from shutting down. Usually this wouldn't be a problem, but if you ever put a thread handle into a static that would keep it around forever (like a tokio runtime, which keeps the thread pool alive until dropped), then the JVM waits for it forever.

The obvious workaround is to attach as daemon, which simply doesn't prevent the JVM to shut down. That sounds good for most cases, as long as the Rust code is loaded as shared lib anyway. Meaning when the JVM shuts down, we are likely happy with shutting everything else down too.

But I don't know how tokio behaves if you steal its threads, even in a shutdown step. And maybe there are some corner cases where you want to attach permanently.

So long story short this is eventually something that should be exposed to the user for choosing.

[bendk]

Ideally this would just do a single call all in Rust by directly doing FFI calls into JNI to attach the current thread. As I mentioned, you can get a list of VM instances through JNI, then use the first one to attach the thread.

I'm understanding better now. So this means the Rust code needs a dependency on the jni crate and it needs to make calls into there.

So long story short this is eventually something that should be exposed to the user for choosing.

Given that users will need to configure the behavior themselves, I don't think we should have a generic "attach_to_thread" function. It seems better to have a JVM-specific one.

This could live in the uniffi::kotlin module, behind a feature flag. However, I'm kind of thinking it could be its own crate. AFAIK, this is a niche issue and most Kotlin users don't need to think about it or pull in the jni dependency.

The other reason to do this is that it simplifies the dependencies management for the Mozilla team. We currently vendor in UniFFI the our monorepo, which requires auditing all of it's dependencies and also vendoring them in -- even if they're behind a feature flag. If this was in a separate crate, then we don't need to deal with that.

[dnaka91]

I'm understanding better now. So this means the Rust code needs a dependency on the jni crate and it needs to make calls into there.

Actually I was hoping to implement it directly (well, I already did). The jni crate brings a variety of dependencies and the part of the JNI API we need doesn't require any of those. So it feels simpler to directly do a few very limited unsafe calls.

Let me create a draft PR of my local testing setup that I did so far, maybe it gives a better understanding of how much additional code this means.