Invalidate Callable when JVM is closing.

Author: CedNaru

src/jvm_wrapper/jvm_instance_wrapper.h

@@ -68,7 +68,13 @@ public:
 template<class Derived, const char* FqName>
 class JvmInstanceWrapper {
 protected:
-    bool is_weak = false;
+    enum RefType {
+        STRONG,
+        WEAK,
+        DEAD
+    };
+
+    RefType ref_type = STRONG;
     jni::JObject wrapped;
 
     explicit JvmInstanceWrapper(jni::Env& p_env, jni::JObject p_wrapped);
@@ -79,6 +85,7 @@ class JvmInstanceWrapper {
     const jni::JObject& get_wrapped() const;
     void swap_to_strong_unsafe(jni::Env& p_env);
     void swap_to_weak_unsafe(jni::Env& p_env);
+    void invalidate_ref(jni::Env& p_env);
 
     static bool initialize(jni::Env& p_env, ClassLoader* class_loader);
     static Derived* create_instance(jni::Env& p_env, ClassLoader* class_loader);
@@ -119,16 +126,12 @@ void JvmInstanceWrapper<Derived, FqName>::finalize(jni::Env& p_env, ClassLoader*
 template<class Derived, const char* FqName>
 JvmInstanceWrapper<Derived, FqName>::~JvmInstanceWrapper() {
     jni::Env env {jni::Jvm::current_env()};
-    if (is_weak) {
-        wrapped.delete_weak_ref(env);
-    } else {
-        wrapped.delete_global_ref(env);
-    }
+    invalidate_ref(env);
 }
 
 template<class Derived, const char* FqName>
 bool JvmInstanceWrapper<Derived, FqName>::is_ref_weak() const {
-    return is_weak;
+    return ref_type == WEAK;
 }
 
 template<class Derived, const char* FqName>
@@ -137,7 +140,7 @@ void JvmInstanceWrapper<Derived, FqName>::swap_to_strong_unsafe(jni::Env& p_env)
     jni::JObject new_ref = wrapped.new_global_ref<jni::JObject>(p_env);
     wrapped.delete_weak_ref(p_env);
     wrapped = new_ref;
-    is_weak = false;
+    ref_type = STRONG;
 }
 
 template<class Derived, const char* FqName>
@@ -146,12 +149,27 @@ void JvmInstanceWrapper<Derived, FqName>::swap_to_weak_unsafe(jni::Env& p_env) {
     jni::JObject new_ref = wrapped.new_weak_ref<jni::JObject>(p_env);
     wrapped.delete_global_ref(p_env);
     wrapped = new_ref;
-    is_weak = true;
+    ref_type = WEAK;
+}
+
+template<class Derived, const char* FqName>
+void JvmInstanceWrapper<Derived, FqName>::invalidate_ref(jni::Env& p_env) {
+    switch (ref_type) {
+        case STRONG:
+            wrapped.delete_global_ref(p_env);
+            break;
+        case WEAK:
+            wrapped.delete_weak_ref(p_env);
+            break;
+        default:
+            break;
+    }
+    ref_type = DEAD;
 }
 
 template<class Derived, const char* FqName>
 const jni::JObject& JvmInstanceWrapper<Derived, FqName>::get_wrapped() const {
     return wrapped;
 }
 
-#endif// GODOT_JVM_JVM_INSTANCE_WRAPPER_H
+#endif // GODOT_JVM_JVM_INSTANCE_WRAPPER_H

src/jvm_wrapper/kotlin_callable_custom.cpp

@@ -1,6 +1,7 @@
 #include "kotlin_callable_custom.h"
 
 #include "gd_kotlin.h"
+#include "jvm_wrapper/memory/memory_manager.h"
 #include "jvm_wrapper/memory/transfer_context.h"
 
 void LambdaContainer::invoke(jni::Env& p_env, const Variant** p_args, int args_count, Variant& r_ret) const {
@@ -82,4 +83,11 @@ bool JvmCallableCustom::compare_less(const CallableCustom* p_a, const CallableCu
 }
 
 JvmCallableCustom::JvmCallableCustom(jni::Env& p_env, jni::JObject p_wrapped, Variant::Type return_type, int p_hash_code, bool p_has_on_destroy) :
-  lambda(p_env, p_wrapped, return_type, p_hash_code, p_has_on_destroy) {}
+  lambda(p_env, p_wrapped, return_type, p_hash_code, p_has_on_destroy) {
+    MemoryManager::get_instance().add_callable(&self_list);
+}
+
+void JvmCallableCustom::invalidate() {
+    jni::Env env {jni::Jvm::current_env()};
+    lambda.invalidate_ref(env);
+}

src/jvm_wrapper/kotlin_callable_custom.h

@@ -1,8 +1,8 @@
 #ifndef GODOT_JVM_KOTLIN_CALLABLE_CUSTOM_H
 #define GODOT_JVM_KOTLIN_CALLABLE_CUSTOM_H
 
-#include "jvm_wrapper/jvm_instance_wrapper.h"
 #include "core/variant/callable.h"
+#include "jvm_wrapper/jvm_instance_wrapper.h"
 
 JVM_INSTANCE_WRAPPER(LambdaContainer, "godot.core.LambdaContainer") {
     JVM_CLASS(LambdaContainer)
@@ -44,15 +44,16 @@ class JvmCallableCustom : public CallableCustom {
     CompareLessFunc get_compare_less_func() const override;
     ObjectID get_object() const override;
 
+    void invalidate();
+
     JvmCallableCustom(jni::Env& p_env, jni::JObject p_wrapped, Variant::Type return_type, int p_hash_code, bool p_has_on_destroy);
-    ~JvmCallableCustom() = default;
 
 private:
     LambdaContainer lambda;
+    SelfList<JvmCallableCustom> self_list {this};
 
     static bool compare_equal(const CallableCustom* p_a, const CallableCustom* p_b);
     static bool compare_less(const CallableCustom* p_a, const CallableCustom* p_b);
 };
 
-
-#endif //GODOT_JVM_KOTLIN_CALLABLE_CUSTOM_H
+#endif // GODOT_JVM_KOTLIN_CALLABLE_CUSTOM_H

src/jvm_wrapper/memory/memory_manager.cpp

@@ -141,6 +141,16 @@ void MemoryManager::sync_memory(jni::Env& p_env) {
 
 void MemoryManager::clean_up(jni::Env& p_env) {
     JVM_LOG_VERBOSE("Cleaning JVM Memory...");
+
+    //Iterate over all remaining custom callables and invalidate their JNI reference before the JVM shut down.
+    SelfList<JvmCallableCustom> *it = callables.first();
+    while (it) {
+        JvmCallableCustom *obj = it->self();
+        obj->invalidate();
+        it = it->next();
+    }
+    callables.clear();
+
     wrapped.call_void_method(p_env, CLEAN_UP);
     JVM_LOG_VERBOSE("JVM Memory cleaned!");
 }
@@ -175,5 +185,11 @@ void MemoryManager::direct_object_deletion(jni::Env& p_env, Object* p_obj) {
     memdelete(p_obj);
 }
 
+void MemoryManager::add_callable(SelfList<JvmCallableCustom>* callable) {
+    callable_mutex.lock();
+    callables.add(callable);
+    callable_mutex.unlock();
+}
+
 MemoryManager::~MemoryManager() = default;
 

src/jvm_wrapper/memory/memory_manager.h

@@ -3,6 +3,7 @@
 
 #include "jni/types.h"
 #include "jvm_wrapper/jvm_singleton_wrapper.h"
+#include "jvm_wrapper/kotlin_callable_custom.h"
 #include "jvm_wrapper/registration/kt_object.h"
 #include "script/jvm_instance.h"
 
@@ -26,6 +27,9 @@ JVM_SINGLETON_WRAPPER(MemoryManager, "godot.internal.memory.MemoryManager") {
         INIT_NATIVE_METHOD("releaseBinding", "(J)V", MemoryManager::release_binding)
       )
 
+    Mutex callable_mutex;
+    SelfList<JvmCallableCustom>::List callables;
+
     Mutex dead_objects_mutex;
     LocalVector<ObjectID> dead_objects;
 
@@ -39,6 +43,7 @@ JVM_SINGLETON_WRAPPER(MemoryManager, "godot.internal.memory.MemoryManager") {
 
 public:
     void direct_object_deletion(jni::Env& p_env, Object* obj);
+    void add_callable(SelfList<JvmCallableCustom>* callable);
     void queue_dead_object(Object* obj);
     void queue_demotion(JvmInstance* script_instance);
     void cancel_demotion(JvmInstance* script_instance);