From ecd68dfc703409496eb44497ce83878995d6e4c5 Mon Sep 17 00:00:00 2001
From: Nekotekina <nekotekina@gmail.com>
Date: Thu, 27 Feb 2020 17:28:42 +0300
Subject: [PATCH] overlays: add "thread bits" to wait on and avoid lockup

Add TLS variable to store its own bit.
---
 .../RSX/Overlays/overlay_message_dialog.cpp   | 14 ++++----
 rpcs3/Emu/RSX/Overlays/overlay_osk.cpp        | 12 +++----
 rpcs3/Emu/RSX/Overlays/overlays.cpp           | 36 ++++++++++++++++---
 rpcs3/Emu/RSX/Overlays/overlays.h             |  6 +++-
 4 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/rpcs3/Emu/RSX/Overlays/overlay_message_dialog.cpp b/rpcs3/Emu/RSX/Overlays/overlay_message_dialog.cpp
index 040c5b8811..0c46e79b87 100644
--- a/rpcs3/Emu/RSX/Overlays/overlay_message_dialog.cpp
+++ b/rpcs3/Emu/RSX/Overlays/overlay_message_dialog.cpp
@@ -242,10 +242,10 @@ namespace rsx
 			{
 				if (!exit)
 				{
-					thread_count++;
-
-					g_fxo->init<named_thread>("MsgDialog Thread", [&]()
+					g_fxo->init<named_thread>("MsgDialog Thread", [&, tbit = alloc_thread_bit()]()
 					{
+						g_thread_bit = tbit;
+
 						if (interactive)
 						{
 							auto ref = g_fxo->get<display_manager>()->get(uid);
@@ -267,15 +267,13 @@ namespace rsx
 								if (!g_fxo->get<display_manager>())
 								{
 									rsx_log.fatal("display_manager was improperly destroyed");
-									return;
+									break;
 								}
 							}
 						}
 
-						if (!--thread_count)
-						{
-							thread_count.notify_all();
-						}
+						thread_bits &= ~tbit;
+						thread_bits.notify_all();
 					});
 				}
 			}
diff --git a/rpcs3/Emu/RSX/Overlays/overlay_osk.cpp b/rpcs3/Emu/RSX/Overlays/overlay_osk.cpp
index b375b01bfb..645e202f4a 100644
--- a/rpcs3/Emu/RSX/Overlays/overlay_osk.cpp
+++ b/rpcs3/Emu/RSX/Overlays/overlay_osk.cpp
@@ -208,19 +208,17 @@ namespace rsx
 			fade_animation.duration = 0.5f;
 			fade_animation.active = true;
 
-			thread_count++;
-
-			g_fxo->init<named_thread>("OSK Thread", [this]
+			g_fxo->init<named_thread>("OSK Thread", [this, tbit = alloc_thread_bit()]
 			{
+				g_thread_bit = tbit;
+
 				if (auto error = run_input_loop())
 				{
 					rsx_log.error("Osk input loop exited with error code=%d", error);
 				}
 
-				if (!--thread_count)
-				{
-					thread_count.notify_all();
-				}
+				thread_bits &= ~tbit;
+				thread_bits.notify_all();
 			});
 		}
 
diff --git a/rpcs3/Emu/RSX/Overlays/overlays.cpp b/rpcs3/Emu/RSX/Overlays/overlays.cpp
index 52d24eb6fd..1aad08bad1 100644
--- a/rpcs3/Emu/RSX/Overlays/overlays.cpp
+++ b/rpcs3/Emu/RSX/Overlays/overlays.cpp
@@ -10,6 +10,33 @@ namespace rsx
 {
 	namespace overlays
 	{
+		thread_local DECLARE(user_interface::g_thread_bit) = 0;
+
+		u64 user_interface::alloc_thread_bit()
+		{
+			auto [_old, ok] = this->thread_bits.fetch_op([](u64& bits)
+			{
+				if (~bits)
+				{
+					// Set lowest clear bit
+					bits |= bits + 1;
+					return true;
+				}
+
+				return false;
+			});
+
+			if (!ok)
+			{
+				::overlays.fatal("Out of thread bits in user interface");
+				return 0;
+			}
+
+			const u64 r = u64{1} << utils::cnttz64(~_old, false);
+			::overlays.trace("Bit allocated (%u)", r);
+			return r;
+		}
+
 		// Singleton instance declaration
 		fontmgr* fontmgr::m_instance = nullptr;
 
@@ -159,14 +186,15 @@ namespace rsx
 			// Force unload
 			exit.release(true);
 
-			while (u32 i = thread_count)
+			while (u64 b = thread_bits)
 			{
-				thread_count.wait(i);
-
-				if (thread_ctrl::state() == thread_state::aborting)
+				if (b == g_thread_bit)
 				{
+					// Don't wait for its own bit
 					break;
 				}
+
+				thread_bits.wait(b);
 			}
 
 			pad::SetIntercepted(false);
diff --git a/rpcs3/Emu/RSX/Overlays/overlays.h b/rpcs3/Emu/RSX/Overlays/overlays.h
index 126b7aa341..e7ddf88a97 100644
--- a/rpcs3/Emu/RSX/Overlays/overlays.h
+++ b/rpcs3/Emu/RSX/Overlays/overlays.h
@@ -72,7 +72,11 @@ namespace rsx
 		protected:
 			Timer input_timer;
 			atomic_t<bool> exit = false;
-			atomic_t<u32> thread_count = 0;
+			atomic_t<u64> thread_bits = 0;
+
+			static thread_local u64 g_thread_bit;
+
+			u64 alloc_thread_bit();
 
 			std::function<void(s32 status)> on_close;