From a7317748135e064ff150072f2f46ccc1d2c4b358 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pedro=20J=2E=20Est=C3=A9banez?= <pedrojrulez@gmail.com>
Date: Thu, 28 Dec 2023 19:31:28 +0100
Subject: [PATCH] WorkerThreadPool: Avoid most runtime allocations

Just a little optimization.

**NOTE:**
With `RID_Owner` we could replace each pair of `PagedAllocator` and
`HashMap`-of-ids-to-pointers. However, that would force us to expose
`RID` as the task/group id, instead of `int`, which would break the
API. Too bad. Let's wait until Godot 5.0.
---
 core/object/worker_thread_pool.cpp | 20 ++++++++++++++------
 core/object/worker_thread_pool.h   | 23 +++++++++++++++++++----
 core/templates/paged_allocator.h   |  6 +-----
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/core/object/worker_thread_pool.cpp b/core/object/worker_thread_pool.cpp
index 881c825cf1..8e8a2ef06b 100644
--- a/core/object/worker_thread_pool.cpp
+++ b/core/object/worker_thread_pool.cpp
@@ -613,13 +613,14 @@ void WorkerThreadPool::finish() {
 		return;
 	}
 
-	task_mutex.lock();
-	SelfList<Task> *E = low_priority_task_queue.first();
-	while (E) {
-		print_error("Task waiting was never re-claimed: " + E->self()->description);
-		E = E->next();
+	{
+		MutexLock lock(task_mutex);
+		SelfList<Task> *E = low_priority_task_queue.first();
+		while (E) {
+			print_error("Task waiting was never re-claimed: " + E->self()->description);
+			E = E->next();
+		}
 	}
-	task_mutex.unlock();
 
 	{
 		MutexLock lock(task_mutex);
@@ -632,6 +633,13 @@ void WorkerThreadPool::finish() {
 		data.thread.wait_to_finish();
 	}
 
+	{
+		MutexLock lock(task_mutex);
+		for (KeyValue<TaskID, Task *> &E : tasks) {
+			task_allocator.free(E.value);
+		}
+	}
+
 	threads.clear();
 }
 
diff --git a/core/object/worker_thread_pool.h b/core/object/worker_thread_pool.h
index 3ec4fd732f..c9921c808d 100644
--- a/core/object/worker_thread_pool.h
+++ b/core/object/worker_thread_pool.h
@@ -95,8 +95,11 @@ private:
 				task_elem(this) {}
 	};
 
-	PagedAllocator<Task> task_allocator;
-	PagedAllocator<Group> group_allocator;
+	static const uint32_t TASKS_PAGE_SIZE = 1024;
+	static const uint32_t GROUPS_PAGE_SIZE = 256;
+
+	PagedAllocator<Task, false, TASKS_PAGE_SIZE> task_allocator;
+	PagedAllocator<Group, false, GROUPS_PAGE_SIZE> group_allocator;
 
 	SelfList<Task>::List low_priority_task_queue;
 	SelfList<Task>::List task_queue;
@@ -117,8 +120,20 @@ private:
 	bool exit_threads = false;
 
 	HashMap<Thread::ID, int> thread_ids;
-	HashMap<TaskID, Task *> tasks;
-	HashMap<GroupID, Group *> groups;
+	HashMap<
+			TaskID,
+			Task *,
+			HashMapHasherDefault,
+			HashMapComparatorDefault<TaskID>,
+			PagedAllocator<HashMapElement<TaskID, Task *>, false, TASKS_PAGE_SIZE>>
+			tasks;
+	HashMap<
+			GroupID,
+			Group *,
+			HashMapHasherDefault,
+			HashMapComparatorDefault<GroupID>,
+			PagedAllocator<HashMapElement<GroupID, Group *>, false, GROUPS_PAGE_SIZE>>
+			groups;
 
 	uint32_t max_low_priority_threads = 0;
 	uint32_t low_priority_threads_used = 0;
diff --git a/core/templates/paged_allocator.h b/core/templates/paged_allocator.h
index 6f3f78d4d2..d880eae0c3 100644
--- a/core/templates/paged_allocator.h
+++ b/core/templates/paged_allocator.h
@@ -40,7 +40,7 @@
 #include <type_traits>
 #include <typeinfo>
 
-template <class T, bool thread_safe = false>
+template <class T, bool thread_safe = false, uint32_t DEFAULT_PAGE_SIZE = 4096>
 class PagedAllocator {
 	T **page_pool = nullptr;
 	T ***available_pool = nullptr;
@@ -53,10 +53,6 @@ class PagedAllocator {
 	SpinLock spin_lock;
 
 public:
-	enum {
-		DEFAULT_PAGE_SIZE = 4096
-	};
-
 	template <class... Args>
 	T *alloc(Args &&...p_args) {
 		if (thread_safe) {