From c303bf9283302b0eca291c9f3800aedc3c894c7e Mon Sep 17 00:00:00 2001
From: SirLynix <lynix680@gmail.com>
Date: Fri, 2 Feb 2024 23:19:29 +0100
Subject: [PATCH] Core/TaskScheduler: Use WorkStealingQueue::pop

steal() can incorrectly return nullptr even if the list is not empty in case of concurrent access, but push and pop are not threadsafe so we use a spinlock to prevent concurrent uses
---
 src/Nazara/Core/TaskScheduler.cpp             | 68 +++++++++++++++++--
 .../Engine/Core/TaskSchedulerTests.cpp        |  3 +-
 2 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/src/Nazara/Core/TaskScheduler.cpp b/src/Nazara/Core/TaskScheduler.cpp
index 2c25aebb2..d7a8900f4 100644
--- a/src/Nazara/Core/TaskScheduler.cpp
+++ b/src/Nazara/Core/TaskScheduler.cpp
@@ -7,6 +7,7 @@
 #include <Nazara/Core/ThreadExt.hpp>
 #include <NazaraUtils/StackArray.hpp>
 #include <wsq.hpp>
+#include <mutex>
 #include <new>
 #include <random>
 #include <semaphore>
@@ -30,6 +31,43 @@ namespace Nz
 #else
 		constexpr std::size_t hardware_destructive_interference_size = 64;
 #endif
+
+		class Spinlock
+		{
+			public:
+				Spinlock() = default;
+				Spinlock(const Spinlock&) = delete;
+				Spinlock(Spinlock&&) = delete;
+				~Spinlock() = default;
+
+				void lock()
+				{
+					while (m_flag.test_and_set());
+				}
+
+				bool try_lock(unsigned int maxLockCount = 1)
+				{
+					unsigned int lockCount = 0;
+					while (m_flag.test_and_set())
+					{
+						if (++lockCount >= maxLockCount)
+							return false;
+					}
+
+					return true;
+				}
+
+				void unlock()
+				{
+					m_flag.clear();
+				}
+
+				Spinlock& operator=(const Spinlock&) = delete;
+				Spinlock& operator=(Spinlock&&) = delete;
+
+			private:
+				std::atomic_flag m_flag;
+		};
 	}
 
 	class alignas(NAZARA_ANONYMOUS_NAMESPACE_PREFIX(hardware_destructive_interference_size)) TaskScheduler::Worker
@@ -67,7 +105,10 @@ namespace Nz
 
 			void AddTask(TaskScheduler::Task* task)
 			{
-				m_tasks.push(task);
+				std::unique_lock lock(m_queueSpinlock);
+				{
+					m_tasks.push(task);
+				}
 				WakeUp();
 			}
 
@@ -92,9 +133,13 @@ namespace Nz
 				bool idle = false;
 				do
 				{
-					// FIXME: We can't use pop() because push() and pop() are not thread-safe (and push is called on another thread), but steal() is
-					// is it an issue?
-					TaskScheduler::Task* task = m_tasks.steal();
+					// Get a task
+					TaskScheduler::Task* task;
+					{
+						std::unique_lock lock(m_queueSpinlock);
+						task = m_tasks.pop();
+					}
+
 					if (!task)
 					{
 						for (unsigned int workerIndex : randomWorkerIndices)
@@ -115,10 +160,15 @@ namespace Nz
 
 #ifdef NAZARA_WITH_TSAN
 						// Workaround for TSan false-positive
-						__tsan_acquire(taskPtr);
+						__tsan_acquire(task);
 #endif
 
 						(*task)();
+
+#ifdef NAZARA_WITH_TSAN
+						// Workaround for TSan false-positive
+						__tsan_release(task);
+#endif
 					}
 					else
 					{
@@ -159,6 +209,7 @@ namespace Nz
 			std::atomic_bool m_running;
 			std::atomic_flag m_notifier;
 			std::thread m_thread; //< std::jthread is not yet widely implemented
+			NAZARA_ANONYMOUS_NAMESPACE_PREFIX(Spinlock) m_queueSpinlock;
 			WorkStealingQueue<TaskScheduler::Task*, TaskScheduler::Task*> m_tasks;
 			TaskScheduler& m_owner;
 			unsigned int m_workerIndex;
@@ -214,6 +265,13 @@ namespace Nz
 	void TaskScheduler::WaitForTasks()
 	{
 		m_idle.wait(false);
+
+#ifdef NAZARA_WITH_TSAN
+		// Workaround for TSan false-positive
+		for (Task& task : m_tasks)
+			__tsan_acquire(&task);
+#endif
+
 		m_tasks.Clear();
 	}
 
diff --git a/tests/UnitTests/Engine/Core/TaskSchedulerTests.cpp b/tests/UnitTests/Engine/Core/TaskSchedulerTests.cpp
index 51b0bb802..8ef643e75 100644
--- a/tests/UnitTests/Engine/Core/TaskSchedulerTests.cpp
+++ b/tests/UnitTests/Engine/Core/TaskSchedulerTests.cpp
@@ -35,7 +35,8 @@ SCENARIO("TaskScheduler", "[CORE][TaskScheduler]")
 				}
 				scheduler.WaitForTasks();
 
-				CHECK(count == taskCount);
+				unsigned int c = count.load(); //< load it once before checking to avoid race condition when testing and printing
+				CHECK(c == taskCount);
 
 				for (std::size_t i = 0; i < taskCount; ++i)
 				{