Merge pull request #3 from wuwbobo2021/receiver

Fixes possible crash on thread exit

Author: Dirbaio

java-spaghetti/src/vm.rs

@@ -1,3 +1,4 @@
+use std::cell::{Cell, OnceCell};
 use std::ptr::null_mut;
 
 use jni_sys::*;
@@ -37,17 +38,57 @@ impl VM {
         F: for<'env> FnOnce(Env<'env>) -> R,
     {
         let mut env = null_mut();
-        match unsafe { ((**self.0).v1_2.GetEnv)(self.0, &mut env, JNI_VERSION_1_2) } {
-            JNI_OK => callback(unsafe { Env::from_raw(env as _) }),
-            JNI_EDETACHED => match unsafe { ((**self.0).v1_2.AttachCurrentThread)(self.0, &mut env, null_mut()) } {
-                JNI_OK => callback(unsafe { Env::from_raw(env as _) }),
-                unexpected => panic!("AttachCurrentThread returned unknown error: {}", unexpected),
-            },
+        let just_attached = match unsafe { ((**self.0).v1_2.GetEnv)(self.0, &mut env, JNI_VERSION_1_2) } {
+            JNI_OK => false,
+            JNI_EDETACHED => {
+                let ret = unsafe { ((**self.0).v1_2.AttachCurrentThread)(self.0, &mut env, null_mut()) };
+                if ret != JNI_OK {
+                    panic!("AttachCurrentThread returned unknown error: {}", ret)
+                }
+                if !get_thread_exit_flag() {
+                    set_thread_attach_flag(self.0);
+                }
+                true
+            }
             JNI_EVERSION => panic!("GetEnv returned JNI_EVERSION"),
             unexpected => panic!("GetEnv returned unknown error: {}", unexpected),
+        };
+
+        let result = callback(unsafe { Env::from_raw(env as _) });
+
+        if just_attached && get_thread_exit_flag() {
+            // this is needed in case of `with_env` is used on dropping some thread-local instance.
+            unsafe { ((**self.0).v1_2.DetachCurrentThread)(self.0) };
         }
+
+        result
     }
 }
 
 unsafe impl Send for VM {}
 unsafe impl Sync for VM {}
+
+thread_local! {
+    static THREAD_ATTACH_FLAG: Cell<Option<AttachFlag>> = const { Cell::new(None) };
+    static THREAD_EXIT_FLAG: OnceCell<()> = const { OnceCell::new() };
+}
+
+struct AttachFlag {
+    raw_vm: *mut JavaVM,
+}
+
+impl Drop for AttachFlag {
+    fn drop(&mut self) {
+        // avoids the fatal error "Native thread exiting without having called DetachCurrentThread"
+        unsafe { ((**self.raw_vm).v1_2.DetachCurrentThread)(self.raw_vm) };
+        let _ = THREAD_EXIT_FLAG.with(|flag| flag.set(()));
+    }
+}
+
+fn set_thread_attach_flag(raw_vm: *mut JavaVM) {
+    THREAD_ATTACH_FLAG.replace(Some(AttachFlag { raw_vm }));
+}
+
+fn get_thread_exit_flag() -> bool {
+    THREAD_EXIT_FLAG.with(|flag| flag.get().is_some())
+}