From d2005ec42d93d1dbd7add2f63d5bc02200e4f703 Mon Sep 17 00:00:00 2001
From: Andreas Maerten <24669514+Yimura@users.noreply.github.com>
Date: Tue, 31 Oct 2023 20:21:40 +0100
Subject: [PATCH] fix: thread pool rescaling incorrectly (#2357)

---
 src/thread_pool.cpp                 | 17 +++++++++--------
 src/thread_pool.hpp                 |  6 ++++++
 src/views/debug/view_debug_misc.cpp |  2 --
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/src/thread_pool.cpp b/src/thread_pool.cpp
index 347c7484..39d67b3c 100644
--- a/src/thread_pool.cpp
+++ b/src/thread_pool.cpp
@@ -4,7 +4,8 @@ namespace big
 {
 	thread_pool::thread_pool(const std::size_t preallocated_thread_count) :
 	    m_accept_jobs(true),
-		m_allocated_thread_count(preallocated_thread_count)
+		m_allocated_thread_count(preallocated_thread_count),
+		m_busy_threads(0)
 	{
 		rescale_thread_pool();
 
@@ -23,7 +24,7 @@ namespace big
 
 		if (m_thread_pool.size() < m_allocated_thread_count)
 		{
-			for (uint32_t i = 0; i < m_allocated_thread_count; i++)
+			for (auto i = m_thread_pool.size(); i < m_allocated_thread_count; i++)
 				m_thread_pool.emplace_back(std::thread(&thread_pool::run, this));
 		}
 	}
@@ -50,16 +51,17 @@ namespace big
 				std::unique_lock lock(m_lock);
 				m_job_stack.push({func, location});
 
-				if (m_allocated_thread_count < m_job_stack.size())
+				if (m_allocated_thread_count - m_busy_threads < m_job_stack.size())
 				{
 					LOG(WARNING) << "Thread pool potentially starved, resizing to accommodate for load.";
 
-					if (m_allocated_thread_count++ >= MAX_POOL_SIZE)
+					if (m_allocated_thread_count >= MAX_POOL_SIZE)
 					{
 						LOG(FATAL) << "The thread pool limit has been reached, whatever you did this should not occur in production.";
 					}
-					if (m_accept_jobs && m_allocated_thread_count <= MAX_POOL_SIZE)
+					if (m_accept_jobs && m_allocated_thread_count + 1 <= MAX_POOL_SIZE)
 					{
+						++m_allocated_thread_count;
 						rescale_thread_pool();
 					}
 				}
@@ -73,7 +75,6 @@ namespace big
 		for (;;)
 		{
 			std::unique_lock lock(m_lock);
-
 			m_data_condition.wait(lock, [this]() {
 				return !m_job_stack.empty() || !m_accept_jobs;
 			});
@@ -87,7 +88,7 @@ namespace big
 			m_job_stack.pop();
 			lock.unlock();
 
-			m_allocated_thread_count--;
+			++m_busy_threads;
 
 			try
 			{
@@ -102,7 +103,7 @@ namespace big
 				LOG(WARNING) << "Exception thrown while executing job in thread:" << std::endl << e.what();
 			}
 
-			m_allocated_thread_count++;
+			--m_busy_threads;
 		}
 
 		LOG(VERBOSE) << "Thread " << std::this_thread::get_id() << " exiting...";
diff --git a/src/thread_pool.hpp b/src/thread_pool.hpp
index 2246a7d7..d2a78e3b 100644
--- a/src/thread_pool.hpp
+++ b/src/thread_pool.hpp
@@ -20,7 +20,10 @@ namespace big
 		std::mutex m_lock;
 		std::vector<std::thread> m_thread_pool;
 
+		// the amount of threads active in the pool
 		std::atomic<size_t> m_allocated_thread_count;
+		// the amount of threads currently on a job
+		std::atomic<size_t> m_busy_threads;
 
 	public:
 		// YimMenu only has 2 blocking threads, 4 should be sufficient but the pool should automatically allocate more if needed
@@ -30,6 +33,9 @@ namespace big
 		void destroy();
 		void push(std::function<void()> func, std::source_location location = std::source_location::current());
 
+		std::pair<size_t, size_t> usage() const
+		{ return { m_busy_threads, m_allocated_thread_count }; }
+
 	private:
 		void run();
 		void rescale_thread_pool();
diff --git a/src/views/debug/view_debug_misc.cpp b/src/views/debug/view_debug_misc.cpp
index 26d2320a..08b73be6 100644
--- a/src/views/debug/view_debug_misc.cpp
+++ b/src/views/debug/view_debug_misc.cpp
@@ -21,9 +21,7 @@ namespace big
 			components::command_checkbox<"windowhook">("VIEW_DEBUG_MISC_DISABLE_GTA_WINDOW_HOOK"_T);
 
 			ImGui::Text(std::format("{}: {}/{}", "VIEW_DEBUG_MISC_FIBER_POOL_USAGE"_T, g_fiber_pool->get_used_fibers(), g_fiber_pool->get_total_fibers()).c_str());
-
 			ImGui::SameLine();
-
 			if (components::button("RESET"_T))
 			{
 				g_fiber_pool->reset();