From 271511726b02848783904429c8dde857f6266429 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pedro=20J=2E=20Est=C3=A9banez?= <pedrojrulez@gmail.com>
Date: Thu, 9 Nov 2023 12:38:01 +0100
Subject: [PATCH] Fixup thread-owned lambda bookkeeping on thread exit

---
 modules/gdscript/gdscript.cpp                 | 104 +++++++++++++++---
 modules/gdscript/gdscript.h                   |  13 ++-
 modules/gdscript/gdscript_lambda_callable.cpp |   4 +-
 modules/gdscript/gdscript_lambda_callable.h   |   4 +-
 4 files changed, 107 insertions(+), 18 deletions(-)

diff --git a/modules/gdscript/gdscript.cpp b/modules/gdscript/gdscript.cpp
index 6245cc85a0c..a6e5e0eaeaa 100644
--- a/modules/gdscript/gdscript.cpp
+++ b/modules/gdscript/gdscript.cpp
@@ -1387,33 +1387,106 @@ String GDScript::debug_get_script_name(const Ref<Script> &p_script) {
 #endif
 
 thread_local GDScript::UpdatableFuncPtr GDScript::func_ptrs_to_update_thread_local;
+GDScript::UpdatableFuncPtr *GDScript::func_ptrs_to_update_main_thread = &func_ptrs_to_update_thread_local;
 
-GDScript::UpdatableFuncPtrElement GDScript::_add_func_ptr_to_update(GDScriptFunction **p_func_ptr_ptr) {
-	UpdatableFuncPtrElement result = {};
+GDScript::UpdatableFuncPtrElement *GDScript::_add_func_ptr_to_update(GDScriptFunction **p_func_ptr_ptr) {
+	MutexLock lock(func_ptrs_to_update_mutex);
+
+	List<UpdatableFuncPtrElement>::Element *result = func_ptrs_to_update_elems.push_back(UpdatableFuncPtrElement());
 
 	{
-		MutexLock lock(func_ptrs_to_update_thread_local.mutex);
-		result.element = func_ptrs_to_update_thread_local.ptrs.push_back(p_func_ptr_ptr);
-		result.mutex = &func_ptrs_to_update_thread_local.mutex;
+		MutexLock lock2(func_ptrs_to_update_thread_local.mutex);
+		result->get().element = func_ptrs_to_update_thread_local.ptrs.push_back(p_func_ptr_ptr);
+		result->get().mutex = &func_ptrs_to_update_thread_local.mutex;
 
 		if (likely(func_ptrs_to_update_thread_local.initialized)) {
-			return result;
+			return &result->get();
 		}
 
 		func_ptrs_to_update_thread_local.initialized = true;
 	}
 
-	MutexLock lock(func_ptrs_to_update_mutex);
 	func_ptrs_to_update.push_back(&func_ptrs_to_update_thread_local);
 
-	return result;
+	return &result->get();
 }
 
-void GDScript::_remove_func_ptr_to_update(const UpdatableFuncPtrElement p_func_ptr_element) {
-	ERR_FAIL_NULL(p_func_ptr_element.element);
-	ERR_FAIL_NULL(p_func_ptr_element.mutex);
-	MutexLock lock(*p_func_ptr_element.mutex);
-	p_func_ptr_element.element->erase();
+void GDScript::_remove_func_ptr_to_update(const UpdatableFuncPtrElement *p_func_ptr_element) {
+	// None of these checks should ever fail, unless there's a bug.
+	// They can be removed once we are sure they never catch anything.
+	// Left here now due to extra safety needs late in the release cycle.
+	ERR_FAIL_NULL(p_func_ptr_element);
+	MutexLock lock(*p_func_ptr_element->mutex);
+	ERR_FAIL_NULL(p_func_ptr_element->element);
+	ERR_FAIL_NULL(p_func_ptr_element->mutex);
+	p_func_ptr_element->element->erase();
+}
+
+void GDScript::_fixup_thread_function_bookkeeping() {
+	// Transfer the ownership of these update items to the main thread,
+	// because the current one is dying, leaving theirs orphan, dangling.
+
+	HashSet<GDScript *> scripts;
+
+	DEV_ASSERT(!Thread::is_main_thread());
+	MutexLock lock(func_ptrs_to_update_main_thread->mutex);
+
+	{
+		MutexLock lock2(func_ptrs_to_update_thread_local.mutex);
+
+		while (!func_ptrs_to_update_thread_local.ptrs.is_empty()) {
+			// Transfer the thread-to-script records from the dying thread to the main one.
+
+			List<GDScriptFunction **>::Element *E = func_ptrs_to_update_thread_local.ptrs.front();
+			List<GDScriptFunction **>::Element *new_E = func_ptrs_to_update_main_thread->ptrs.push_front(E->get());
+
+			GDScript *script = (*E->get())->get_script();
+			if (!scripts.has(script)) {
+				scripts.insert(script);
+
+				// Replace dying thread by the main thread in the script-to-thread records.
+
+				MutexLock lock3(script->func_ptrs_to_update_mutex);
+				DEV_ASSERT(script->func_ptrs_to_update.find(&func_ptrs_to_update_thread_local));
+				{
+					for (List<UpdatableFuncPtrElement>::Element *F = script->func_ptrs_to_update_elems.front(); F; F = F->next()) {
+						bool is_dying_thread_entry = F->get().mutex == &func_ptrs_to_update_thread_local.mutex;
+						if (is_dying_thread_entry) {
+							// This may lead to multiple main-thread entries, but that's not a problem
+							// and allows to reuse the element, which is needed, since it's tracked by pointer.
+							F->get().element = new_E;
+							F->get().mutex = &func_ptrs_to_update_main_thread->mutex;
+						}
+					}
+				}
+			}
+
+			E->erase();
+		}
+	}
+	func_ptrs_to_update_main_thread->initialized = true;
+
+	{
+		// Remove orphan thread-to-script entries from every script.
+		// FIXME: This involves iterating through every script whenever a thread dies.
+		//        While it's OK that thread creation/destruction are heavy operations,
+		//        additional bookkeeping can be used to outperform this brute-force approach.
+
+		GDScriptLanguage *gd_lang = GDScriptLanguage::get_singleton();
+
+		MutexLock lock2(gd_lang->mutex);
+
+		for (SelfList<GDScript> *s = gd_lang->script_list.first(); s; s = s->next()) {
+			GDScript *script = s->self();
+			for (List<UpdatableFuncPtr *>::Element *E = script->func_ptrs_to_update.front(); E; E = E->next()) {
+				bool is_dying_thread_entry = &E->get()->mutex == &func_ptrs_to_update_thread_local.mutex;
+				if (is_dying_thread_entry) {
+					E->erase();
+					break;
+				}
+			}
+		}
+	}
 }
 
 void GDScript::clear(ClearData *p_clear_data) {
@@ -1441,6 +1514,7 @@ void GDScript::clear(ClearData *p_clear_data) {
 				*func_ptr_ptr = nullptr;
 			}
 		}
+		func_ptrs_to_update_elems.clear();
 	}
 
 	RBSet<GDScript *> must_clear_dependencies = get_must_clear_dependencies();
@@ -2060,6 +2134,10 @@ void GDScriptLanguage::remove_named_global_constant(const StringName &p_name) {
 	named_globals.erase(p_name);
 }
 
+void GDScriptLanguage::thread_exit() {
+	GDScript::_fixup_thread_function_bookkeeping();
+}
+
 void GDScriptLanguage::init() {
 	//populate global constants
 	int gcc = CoreConstants::get_global_constant_count();
diff --git a/modules/gdscript/gdscript.h b/modules/gdscript/gdscript.h
index 04b0a1d7867..342df8e9cb0 100644
--- a/modules/gdscript/gdscript.h
+++ b/modules/gdscript/gdscript.h
@@ -128,11 +128,16 @@ class GDScript : public Script {
 		Mutex *mutex = nullptr;
 	};
 	static thread_local UpdatableFuncPtr func_ptrs_to_update_thread_local;
+	static thread_local LocalVector<List<UpdatableFuncPtr *>::Element> func_ptrs_to_update_entries_thread_local;
+	static UpdatableFuncPtr *func_ptrs_to_update_main_thread;
 	List<UpdatableFuncPtr *> func_ptrs_to_update;
+	List<UpdatableFuncPtrElement> func_ptrs_to_update_elems;
 	Mutex func_ptrs_to_update_mutex;
 
-	UpdatableFuncPtrElement _add_func_ptr_to_update(GDScriptFunction **p_func_ptr_ptr);
-	static void _remove_func_ptr_to_update(const UpdatableFuncPtrElement p_func_ptr_element);
+	UpdatableFuncPtrElement *_add_func_ptr_to_update(GDScriptFunction **p_func_ptr_ptr);
+	static void _remove_func_ptr_to_update(const UpdatableFuncPtrElement *p_func_ptr_element);
+
+	static void _fixup_thread_function_bookkeeping();
 
 #ifdef TOOLS_ENABLED
 	// For static data storage during hot-reloading.
@@ -553,6 +558,10 @@ public:
 	virtual void add_named_global_constant(const StringName &p_name, const Variant &p_value) override;
 	virtual void remove_named_global_constant(const StringName &p_name) override;
 
+	/* MULTITHREAD FUNCTIONS */
+
+	virtual void thread_exit() override;
+
 	/* DEBUGGER FUNCTIONS */
 
 	virtual String debug_get_error() const override;
diff --git a/modules/gdscript/gdscript_lambda_callable.cpp b/modules/gdscript/gdscript_lambda_callable.cpp
index 547f5607d33..339d1ac08ec 100644
--- a/modules/gdscript/gdscript_lambda_callable.cpp
+++ b/modules/gdscript/gdscript_lambda_callable.cpp
@@ -296,5 +296,7 @@ GDScriptLambdaSelfCallable::GDScriptLambdaSelfCallable(Object *p_self, GDScriptF
 }
 
 GDScriptLambdaSelfCallable::~GDScriptLambdaSelfCallable() {
-	GDScript::_remove_func_ptr_to_update(updatable_func_ptr_element);
+	if (updatable_func_ptr_element) {
+		GDScript::_remove_func_ptr_to_update(updatable_func_ptr_element);
+	}
 }
diff --git a/modules/gdscript/gdscript_lambda_callable.h b/modules/gdscript/gdscript_lambda_callable.h
index ee7d547544c..d961f18852a 100644
--- a/modules/gdscript/gdscript_lambda_callable.h
+++ b/modules/gdscript/gdscript_lambda_callable.h
@@ -45,7 +45,7 @@ class GDScriptLambdaCallable : public CallableCustom {
 	GDScriptFunction *function = nullptr;
 	Ref<GDScript> script;
 	uint32_t h;
-	GDScript::UpdatableFuncPtrElement updatable_func_ptr_element;
+	GDScript::UpdatableFuncPtrElement *updatable_func_ptr_element = nullptr;
 
 	Vector<Variant> captures;
 
@@ -72,7 +72,7 @@ class GDScriptLambdaSelfCallable : public CallableCustom {
 	Ref<RefCounted> reference; // For objects that are RefCounted, keep a reference.
 	Object *object = nullptr; // For non RefCounted objects, use a direct pointer.
 	uint32_t h;
-	GDScript::UpdatableFuncPtrElement updatable_func_ptr_element;
+	GDScript::UpdatableFuncPtrElement *updatable_func_ptr_element = nullptr;
 
 	Vector<Variant> captures;