From 9809265be4cf61c5406a79642254b7cbbcb00ee1 Mon Sep 17 00:00:00 2001
From: Stenzek <stenzek@gmail.com>
Date: Sun, 28 Jan 2024 21:00:06 +1000
Subject: [PATCH] GS/Vulkan: Fix potential race between submit and main thread

Backport of https://github.com/stenzek/duckstation/commit/bcf7f55b93d6c61a66c1aabec5b8568d10a5b26c
---
 pcsx2/GS/Renderers/Vulkan/GSDeviceVK.cpp | 33 ++++++++++++++----------
 pcsx2/GS/Renderers/Vulkan/GSDeviceVK.h   |  2 +-
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/pcsx2/GS/Renderers/Vulkan/GSDeviceVK.cpp b/pcsx2/GS/Renderers/Vulkan/GSDeviceVK.cpp
index 852d6267e790c..0c5e969b18cf7 100644
--- a/pcsx2/GS/Renderers/Vulkan/GSDeviceVK.cpp
+++ b/pcsx2/GS/Renderers/Vulkan/GSDeviceVK.cpp
@@ -1058,6 +1058,13 @@ void GSDeviceVK::ScanForCommandBufferCompletion()
 
 void GSDeviceVK::WaitForCommandBufferCompletion(u32 index)
 {
+	// We might be waiting for the buffer we just submitted to the worker thread.
+	if (m_queued_present.command_buffer_index == index && !m_present_done.load(std::memory_order_acquire))
+	{
+		Console.WarningFmt("Waiting for threaded submission of cmdbuffer {}", index);
+		WaitForPresentComplete();
+	}
+
 	// Wait for this command buffer to be completed.
 	const VkResult res = vkWaitForFences(m_device, 1, &m_frame_resources[index].fence, VK_TRUE, UINT64_MAX);
 	if (res != VK_SUCCESS)
@@ -1174,7 +1181,7 @@ void GSDeviceVK::SubmitCommandBuffer(
 	m_queued_present.command_buffer_index = m_current_frame;
 	m_queued_present.swap_chain = present_swap_chain;
 	m_queued_present.spin_cycles = spin_cycles;
-	m_present_done.store(false);
+	m_present_done.store(false, std::memory_order_release);
 	m_present_queued_cv.notify_one();
 }
 
@@ -1253,7 +1260,7 @@ void GSDeviceVK::DoPresent(VKSwapChain* present_swap_chain)
 
 void GSDeviceVK::WaitForPresentComplete()
 {
-	if (m_present_done.load())
+	if (m_present_done.load(std::memory_order_acquire))
 		return;
 
 	std::unique_lock<std::mutex> lock(m_present_mutex);
@@ -1262,27 +1269,30 @@ void GSDeviceVK::WaitForPresentComplete()
 
 void GSDeviceVK::WaitForPresentComplete(std::unique_lock<std::mutex>& lock)
 {
-	if (m_present_done.load())
+	if (m_present_done.load(std::memory_order_acquire))
 		return;
 
-	m_present_done_cv.wait(lock, [this]() { return m_present_done.load(); });
+	m_present_done_cv.wait(lock, [this]() { return m_present_done.load(std::memory_order_acquire); });
 }
 
 void GSDeviceVK::PresentThread()
 {
 	std::unique_lock<std::mutex> lock(m_present_mutex);
-	while (!m_present_thread_done.load())
+	while (!m_present_thread_done.load(std::memory_order_acquire))
 	{
-		m_present_queued_cv.wait(lock, [this]() { return !m_present_done.load() || m_present_thread_done.load(); });
+		m_present_queued_cv.wait(lock, [this]() {
+			return !m_present_done.load(std::memory_order_acquire) ||
+				   m_present_thread_done.load(std::memory_order_acquire);
+		});
 
-		if (m_present_done.load())
+		if (m_present_done.load(std::memory_order_acquire))
 			continue;
 
 		DoSubmitCommandBuffer(
 			m_queued_present.command_buffer_index, m_queued_present.swap_chain, m_queued_present.spin_cycles);
 		if (m_queued_present.swap_chain)
 			DoPresent(m_queued_present.swap_chain);
-		m_present_done.store(true);
+		m_present_done.store(true, std::memory_order_release);
 		m_present_done_cv.notify_one();
 	}
 }
@@ -1290,7 +1300,7 @@ void GSDeviceVK::PresentThread()
 void GSDeviceVK::StartPresentThread()
 {
 	pxAssert(!m_present_thread.joinable());
-	m_present_thread_done.store(false);
+	m_present_thread_done.store(false, std::memory_order_release);
 	m_present_thread = std::thread(&GSDeviceVK::PresentThread, this);
 }
 
@@ -1302,7 +1312,7 @@ void GSDeviceVK::StopPresentThread()
 	{
 		std::unique_lock<std::mutex> lock(m_present_mutex);
 		WaitForPresentComplete(lock);
-		m_present_thread_done.store(true);
+		m_present_thread_done.store(true, std::memory_order_acquire);
 		m_present_queued_cv.notify_one();
 	}
 
@@ -1367,9 +1377,6 @@ void GSDeviceVK::ActivateCommandBuffer(u32 index)
 {
 	FrameResources& resources = m_frame_resources[index];
 
-	if (!m_present_done.load() && m_queued_present.command_buffer_index == index)
-		WaitForPresentComplete();
-
 	// Wait for the GPU to finish with all resources for this command buffer.
 	if (resources.fence_counter > m_completed_fence_counter)
 		WaitForCommandBufferCompletion(index);
diff --git a/pcsx2/GS/Renderers/Vulkan/GSDeviceVK.h b/pcsx2/GS/Renderers/Vulkan/GSDeviceVK.h
index c64bfe5183129..a36355ad93524 100644
--- a/pcsx2/GS/Renderers/Vulkan/GSDeviceVK.h
+++ b/pcsx2/GS/Renderers/Vulkan/GSDeviceVK.h
@@ -303,7 +303,7 @@ class GSDeviceVK final : public GSDevice
 		u32 spin_cycles;
 	};
 
-	QueuedPresent m_queued_present = {};
+	QueuedPresent m_queued_present = {nullptr, 0xFFFFFFFFu, 0};
 
 	std::map<u32, VkRenderPass> m_render_pass_cache;