1 files changed, 16 insertions, 22 deletions

diff --git a/src/thread_pool.cpp b/src/thread_pool.cpp
index 5f21e0c3d..54d6cd72c 100644
--- a/src/thread_pool.cpp
+++ b/src/thread_pool.cpp
@@ -11,30 +11,27 @@ struct WorkerTask {
 
 struct ThreadPool {
 	std::atomic<isize> outstanding_task_count;
-	WorkerTask *next_task;
+	WorkerTask *volatile next_task;
 	BlockingMutex task_list_mutex;
 };
 
-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);
+			}
 		}
 	}
 }
@@ -77,10 +74,7 @@ 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();
-	}
+	thread_pool_thread_entry(pool);
 }
 
 void thread_pool_add_task(ThreadPool *pool, WorkerTaskProc *proc, void *data) {