From c450260e3ed4357a477eaee733a091e8206c4ea1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pedro=20J=2E=20Est=C3=A9banez?= <pedrojrulez@gmail.com>
Date: Fri, 19 May 2023 10:46:49 +0200
Subject: [PATCH] Fix message queue issues

- Missing flush in resource loading.
- Wrong checks about message queue instance.
---
 core/io/resource_loader.cpp   |  4 +++-
 core/object/message_queue.cpp | 32 ++++++++++++++++++++++++++++++--
 core/object/message_queue.h   |  6 ++++++
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/core/io/resource_loader.cpp b/core/io/resource_loader.cpp
index c68191554c35..e9f812ab1c7e 100644
--- a/core/io/resource_loader.cpp
+++ b/core/io/resource_loader.cpp
@@ -309,6 +309,9 @@ void ResourceLoader::_thread_load_function(void *p_userdata) {
 	// --
 
 	Ref<Resource> res = _load(load_task.remapped_path, load_task.remapped_path != load_task.local_path ? load_task.local_path : String(), load_task.type_hint, load_task.cache_mode, &load_task.error, load_task.use_sub_threads, &load_task.progress);
+	if (mq_override) {
+		mq_override->flush();
+	}
 
 	thread_load_mutex.lock();
 
@@ -354,7 +357,6 @@ void ResourceLoader::_thread_load_function(void *p_userdata) {
 
 	if (load_nesting == 0 && mq_override) {
 		memdelete(mq_override);
-		MessageQueue::set_thread_singleton_override(nullptr);
 	}
 }
 
diff --git a/core/object/message_queue.cpp b/core/object/message_queue.cpp
index 78fdb08f853b..18ba5d5b3006 100644
--- a/core/object/message_queue.cpp
+++ b/core/object/message_queue.cpp
@@ -35,10 +35,22 @@
 #include "core/object/class_db.h"
 #include "core/object/script_language.h"
 
+#ifdef DEV_ENABLED
+// Includes sanity checks to ensure that a queue set as a thread singleton override
+// is only ever called from the thread it was set for.
+#define LOCK_MUTEX                                     \
+	if (this != MessageQueue::thread_singleton) {      \
+		DEV_ASSERT(!this->is_current_thread_override); \
+		mutex.lock();                                  \
+	} else {                                           \
+		DEV_ASSERT(this->is_current_thread_override);  \
+	}
+#else
 #define LOCK_MUTEX                                \
 	if (this != MessageQueue::thread_singleton) { \
 		mutex.lock();                             \
 	}
+#endif
 
 #define UNLOCK_MUTEX                              \
 	if (this != MessageQueue::thread_singleton) { \
@@ -213,8 +225,8 @@ void CallQueue::_call_function(const Callable &p_callable, const Variant *p_args
 Error CallQueue::flush() {
 	LOCK_MUTEX;
 
-	// Non-main threads are not meant to be flushed, but appended to the main one.
-	if (this != MessageQueue::main_singleton) {
+	// Thread overrides are not meant to be flushed, but appended to the main one.
+	if (this == MessageQueue::thread_singleton) {
 		if (pages.size() == 0) {
 			return OK;
 		}
@@ -237,6 +249,7 @@ Error CallQueue::flush() {
 				uint32_t dst_offset = mq->page_bytes[dst_page];
 				if (dst_offset + page_bytes[0] < uint32_t(PAGE_SIZE_BYTES)) {
 					memcpy(mq->pages[dst_page]->data + dst_offset, pages[0]->data, page_bytes[0]);
+					mq->page_bytes[dst_page] += page_bytes[0];
 					src_page++;
 				}
 			}
@@ -520,6 +533,10 @@ CallQueue::~CallQueue() {
 	if (!allocator_is_custom) {
 		memdelete(allocator);
 	}
+	// This is done here to avoid a circular dependency between the sanity checks and the thread singleton pointer.
+	if (this == MessageQueue::thread_singleton) {
+		MessageQueue::thread_singleton = nullptr;
+	}
 }
 
 //////////////////////
@@ -528,7 +545,18 @@ CallQueue *MessageQueue::main_singleton = nullptr;
 thread_local CallQueue *MessageQueue::thread_singleton = nullptr;
 
 void MessageQueue::set_thread_singleton_override(CallQueue *p_thread_singleton) {
+	DEV_ASSERT(p_thread_singleton); // To unset the thread singleton, don't call this with nullptr, but just memfree() it.
+#ifdef DEV_ENABLED
+	if (thread_singleton) {
+		thread_singleton->is_current_thread_override = false;
+	}
+#endif
 	thread_singleton = p_thread_singleton;
+#ifdef DEV_ENABLED
+	if (thread_singleton) {
+		thread_singleton->is_current_thread_override = true;
+	}
+#endif
 }
 
 MessageQueue::MessageQueue() :
diff --git a/core/object/message_queue.h b/core/object/message_queue.h
index c6fcccbd5867..9f567e4dd045 100644
--- a/core/object/message_queue.h
+++ b/core/object/message_queue.h
@@ -40,6 +40,8 @@
 class Object;
 
 class CallQueue {
+	friend class MessageQueue;
+
 public:
 	enum {
 		PAGE_SIZE_BYTES = 4096
@@ -75,6 +77,10 @@ private:
 	uint32_t pages_used = 0;
 	bool flushing = false;
 
+#ifdef DEV_ENABLED
+	bool is_current_thread_override = false;
+#endif
+
 	struct Message {
 		Callable callable;
 		int16_t type;