1 files changed, 22 insertions, 26 deletions

diff --git a/src/thread_pool.cpp b/src/thread_pool.cpp
index 5f21e0c3d..e904a2e29 100644
--- a/src/thread_pool.cpp
+++ b/src/thread_pool.cpp
@@ -11,30 +11,28 @@ struct WorkerTask {
 
 struct ThreadPool {
 	std::atomic<isize> outstanding_task_count;
-	WorkerTask *next_task;
+	WorkerTask *volatile next_task;
 	BlockingMutex task_list_mutex;
+	isize thread_count;
 };
 
-void thread_pool_init(ThreadPool *pool, gbAllocator const &a, isize thread_count, char const *worker_prefix = nullptr);
-void thread_pool_destroy(ThreadPool *pool);
-void thread_pool_wait(ThreadPool *pool);
-void thread_pool_add_task(ThreadPool *pool, WorkerTaskProc *proc, void *data);
-void worker_thread_internal();
-
 void thread_pool_thread_entry(ThreadPool *pool) {
 	while (pool->outstanding_task_count) {
-		mutex_lock(&pool->task_list_mutex);
-
-		if (pool->next_task) {
-			WorkerTask *task = pool->next_task;
-			pool->next_task = task->next_task;
-			mutex_unlock(&pool->task_list_mutex);
-			task->do_work(task->data);
-			pool->outstanding_task_count.fetch_sub(1);
-			gb_free(heap_allocator(), task);
+		if (!pool->next_task) {
+			yield(); // No need to grab the mutex.
 		} else {
-			mutex_unlock(&pool->task_list_mutex);
-			yield();
+			mutex_lock(&pool->task_list_mutex);
+
+			if (pool->next_task) {
+				WorkerTask *task = pool->next_task;
+				pool->next_task = task->next_task;
+				mutex_unlock(&pool->task_list_mutex);
+				task->do_work(task->data);
+				pool->outstanding_task_count.fetch_sub(1);
+				gb_free(heap_allocator(), task);
+			} else {
+				mutex_unlock(&pool->task_list_mutex);
+			}
 		}
 	}
 }
@@ -65,10 +63,7 @@ void thread_pool_init(ThreadPool *pool, gbAllocator const &a, isize thread_count
 	memset(pool, 0, sizeof(ThreadPool));
 	mutex_init(&pool->task_list_mutex);
 	pool->outstanding_task_count.store(1);
-
-	for (int i = 0; i < thread_count; i++) {
-		thread_pool_start_thread(pool);
-	}
+	pool->thread_count = thread_count;
 }
 
 void thread_pool_destroy(ThreadPool *pool) {
@@ -76,11 +71,12 @@ void thread_pool_destroy(ThreadPool *pool) {
 }
 
 void thread_pool_wait(ThreadPool *pool) {
-	pool->outstanding_task_count.fetch_sub(1);
-
-	while (pool->outstanding_task_count.load() != 0) {
-		yield();
+	for (int i = 0; i < pool->thread_count; i++) {
+		thread_pool_start_thread(pool);
 	}
+
+	pool->outstanding_task_count.fetch_sub(1);
+	thread_pool_thread_entry(pool);
 }
 
 void thread_pool_add_task(ThreadPool *pool, WorkerTaskProc *proc, void *data) {