From 2a483fa9ba46035270c95a9cebf3877f307b071c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pedro=20J=2E=20Est=C3=A9banez?= <pedrojrulez@gmail.com>
Date: Fri, 13 Sep 2024 14:20:11 +0200
Subject: [PATCH] WorkerThreadPool: Refactor running and exit-requested as
 runlevels

---
 core/object/worker_thread_pool.cpp | 32 +++++++++++++++++++++---------
 core/object/worker_thread_pool.h   |  8 +++++++-
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/core/object/worker_thread_pool.cpp b/core/object/worker_thread_pool.cpp
index 91e010b2880..0c19fe06a43 100644
--- a/core/object/worker_thread_pool.cpp
+++ b/core/object/worker_thread_pool.cpp
@@ -186,7 +186,8 @@ void WorkerThreadPool::_thread_function(void *p_user) {
 		{
 			MutexLock lock(singleton->task_mutex);
 
-			if (unlikely(singleton->exit_threads)) {
+			bool exit = singleton->_handle_runlevel();
+			if (unlikely(exit)) {
 				break;
 			}
 
@@ -453,7 +454,8 @@ void WorkerThreadPool::_wait_collaboratively(ThreadData *p_caller_pool_thread, T
 			bool was_signaled = p_caller_pool_thread->signaled;
 			p_caller_pool_thread->signaled = false;
 
-			if (unlikely(exit_threads)) {
+			bool exit = _handle_runlevel();
+			if (unlikely(exit)) {
 				break;
 			}
 
@@ -518,6 +520,20 @@ void WorkerThreadPool::_wait_collaboratively(ThreadData *p_caller_pool_thread, T
 	}
 }
 
+void WorkerThreadPool::_switch_runlevel(Runlevel p_runlevel) {
+	DEV_ASSERT(p_runlevel > runlevel);
+	runlevel = p_runlevel;
+	for (uint32_t i = 0; i < threads.size(); i++) {
+		threads[i].cond_var.notify_one();
+		threads[i].signaled = true;
+	}
+}
+
+// Returns whether threads have to exit. This may perform the check about handling needed.
+bool WorkerThreadPool::_handle_runlevel() {
+	return runlevel == RUNLEVEL_EXIT;
+}
+
 void WorkerThreadPool::yield() {
 	int th_index = get_thread_index();
 	ERR_FAIL_COND_MSG(th_index == -1, "This function can only be called from a worker thread.");
@@ -695,6 +711,9 @@ void WorkerThreadPool::thread_exit_unlock_allowance_zone(uint32_t p_zone_id) {
 
 void WorkerThreadPool::init(int p_thread_count, float p_low_priority_task_ratio) {
 	ERR_FAIL_COND(threads.size() > 0);
+
+	runlevel = RUNLEVEL_NORMAL;
+
 	if (p_thread_count < 0) {
 		p_thread_count = OS::get_singleton()->get_default_thread_pool_size();
 	}
@@ -724,15 +743,10 @@ void WorkerThreadPool::finish() {
 			print_error("Task waiting was never re-claimed: " + E->self()->description);
 			E = E->next();
 		}
+
+		_switch_runlevel(RUNLEVEL_EXIT);
 	}
 
-	{
-		MutexLock lock(task_mutex);
-		exit_threads = true;
-	}
-	for (ThreadData &data : threads) {
-		data.cond_var.notify_one();
-	}
 	for (ThreadData &data : threads) {
 		data.thread.wait_to_finish();
 	}
diff --git a/core/object/worker_thread_pool.h b/core/object/worker_thread_pool.h
index 6374dbe8c7e..ba6efbb0650 100644
--- a/core/object/worker_thread_pool.h
+++ b/core/object/worker_thread_pool.h
@@ -124,7 +124,10 @@ private:
 	};
 
 	TightLocalVector<ThreadData> threads;
-	bool exit_threads = false;
+	enum Runlevel {
+		RUNLEVEL_NORMAL,
+		RUNLEVEL_EXIT,
+	} runlevel = RUNLEVEL_NORMAL;
 
 	HashMap<Thread::ID, int> thread_ids;
 	HashMap<
@@ -193,6 +196,9 @@ private:
 
 	void _wait_collaboratively(ThreadData *p_caller_pool_thread, Task *p_task);
 
+	void _switch_runlevel(Runlevel p_runlevel);
+	bool _handle_runlevel();
+
 #ifdef THREADS_ENABLED
 	static uint32_t _thread_enter_unlock_allowance_zone(THREADING_NAMESPACE::unique_lock<THREADING_NAMESPACE::mutex> &p_ulock);
 #endif