From faba366f89c67c98d79349fd66954d19a5d39509 Mon Sep 17 00:00:00 2001
From: Nekotekina <nekotekina@gmail.com>
Date: Sun, 8 Sep 2019 23:27:57 +0300
Subject: [PATCH] named_thread: install atomic wait callback

---
 Utilities/Thread.cpp | 13 +++++++-
 Utilities/Thread.h   | 72 ++++++++++++++++++++++++++++++++++++--------
 2 files changed, 72 insertions(+), 13 deletions(-)

diff --git a/Utilities/Thread.cpp b/Utilities/Thread.cpp
index 2c2e352616..fe6fe0c1e4 100644
--- a/Utilities/Thread.cpp
+++ b/Utilities/Thread.cpp
@@ -1667,11 +1667,14 @@ void thread_base::start(native_entry entry)
 #endif
 }
 
-void thread_base::initialize()
+void thread_base::initialize(bool(*wait_cb)(const void*))
 {
 	// Initialize TLS variable
 	thread_ctrl::g_tls_this_thread = this;
 
+	// Initialize atomic wait callback
+	atomic_storage_futex::set_wait_callback(wait_cb);
+
 	g_tls_log_prefix = []
 	{
 		return thread_ctrl::g_tls_this_thread->m_name.get();
@@ -1718,6 +1721,14 @@ void thread_base::initialize()
 #endif
 }
 
+void thread_base::notify_abort() noexcept
+{
+	// For now
+	notify();
+
+	atomic_storage_futex::raw_notify(+m_state_notifier);
+}
+
 bool thread_base::finalize(int) noexcept
 {
 	// Report pending errors
diff --git a/Utilities/Thread.h b/Utilities/Thread.h
index ea651263fc..d37a13a464 100644
--- a/Utilities/Thread.h
+++ b/Utilities/Thread.h
@@ -34,7 +34,7 @@ enum class thread_class : u32
 	ppu
 };
 
-enum class thread_state
+enum class thread_state : u32
 {
 	created,  // Initial state
 	detached, // The thread has been detached to destroy its own named_thread object (can be dangerously misused)
@@ -95,6 +95,13 @@ struct thread_on_cleanup : std::bool_constant<false> {};
 template <typename T>
 struct thread_on_cleanup<T, decltype(named_thread<T>::on_cleanup(std::declval<named_thread<T>*>()))> : std::bool_constant<true> {};
 
+// Detect on_wait() method (should return bool)
+template <typename T, typename = bool>
+struct thread_on_wait : std::bool_constant<false> {};
+
+template <typename T>
+struct thread_on_wait<T, decltype(std::declval<named_thread<T>&>().on_wait())> : std::bool_constant<true> {};
+
 // Thread base class
 class thread_base
 {
@@ -123,6 +130,9 @@ class thread_base
 	// Thread state
 	atomic_t<thread_state> m_state = thread_state::created;
 
+	// Thread state notification info
+	atomic_t<const void*> m_state_notifier{nullptr};
+
 	// Thread name
 	lf_value<std::string> m_name;
 
@@ -133,7 +143,10 @@ class thread_base
 	void start(native_entry);
 
 	// Called at the thread start
-	void initialize();
+	void initialize(bool(*wait_cb)(const void*));
+
+	// May be called in destructor
+	void notify_abort() noexcept;
 
 	// Called at the thread end, returns true if needs destruction
 	bool finalize(int) noexcept;
@@ -314,7 +327,47 @@ class named_thread final : public Context, result_storage_t<Context>, thread_bas
 
 	bool entry_point()
 	{
-		thread::initialize();
+		thread::initialize([](const void* data)
+		{
+			const auto _this = thread_ctrl::get_current();
+
+			if (_this->m_state >= thread_state::aborting)
+			{
+				return false;
+			}
+
+			if constexpr (thread_on_wait<Context>())
+			{
+				if (!static_cast<named_thread*>(_this)->on_wait())
+				{
+					return false;
+				}
+			}
+
+			_this->m_state_notifier.release(data);
+
+			if (!data)
+			{
+				return true;
+			}
+
+			if (_this->m_state >= thread_state::aborting)
+			{
+				_this->m_state_notifier.release(nullptr);
+				return false;
+			}
+
+			if constexpr (thread_on_wait<Context>())
+			{
+				if (!static_cast<named_thread*>(_this)->on_wait())
+				{
+					_this->m_state_notifier.release(nullptr);
+					return false;
+				}
+			}
+
+			return true;
+		});
 
 		if constexpr (result::empty)
 		{
@@ -394,12 +447,7 @@ public:
 	// Try to abort/detach
 	named_thread& operator=(thread_state s)
 	{
-		if (s != thread_state::aborting && s != thread_state::detached)
-		{
-			ASSUME(0);
-		}
-
-		if (thread::m_state.compare_and_swap_test(thread_state::created, s))
+		if (s < thread_state::finished && thread::m_state.compare_and_swap_test(thread_state::created, s))
 		{
 			if (s == thread_state::aborting)
 			{
@@ -408,9 +456,9 @@ public:
 				{
 					Context::on_abort();
 				}
-			}
 
-			thread::notify();
+				thread::notify_abort();
+			}
 		}
 
 		return *this;
@@ -419,7 +467,7 @@ public:
 	// Context type doesn't need virtual destructor
 	~named_thread()
 	{
-		*this = thread_state::aborting;
+		operator=(thread_state::aborting);
 		thread::join();
 
 		if constexpr (!result::empty)