From c1a7235767ba4e3cb05c62ae7d08c54181b074e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrik=20Rydg=C3=A5rd?= <hrydgard@gmail.com>
Date: Mon, 15 Jan 2024 17:06:00 +0100
Subject: [PATCH] Fix another shutdown race condition in the Vulkan backend

---
 Common/GPU/Vulkan/VulkanRenderManager.cpp | 22 ++++++++++++++++++++--
 Common/GPU/Vulkan/thin3d_vulkan.cpp       |  2 ++
 Common/Thread/ThreadManager.h             |  1 +
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/Common/GPU/Vulkan/VulkanRenderManager.cpp b/Common/GPU/Vulkan/VulkanRenderManager.cpp
index 0af2eefb5d54..a26a58369a3c 100644
--- a/Common/GPU/Vulkan/VulkanRenderManager.cpp
+++ b/Common/GPU/Vulkan/VulkanRenderManager.cpp
@@ -411,7 +411,9 @@ struct SinglePipelineTask {
 
 class CreateMultiPipelinesTask : public Task {
 public:
-	CreateMultiPipelinesTask(VulkanContext *vulkan, std::vector<SinglePipelineTask> tasks) : vulkan_(vulkan), tasks_(tasks) {}
+	CreateMultiPipelinesTask(VulkanContext *vulkan, std::vector<SinglePipelineTask> tasks) : vulkan_(vulkan), tasks_(tasks) {
+		tasksInFlight_.fetch_add(1);
+	}
 	~CreateMultiPipelinesTask() {}
 
 	TaskType Type() const override {
@@ -426,12 +428,25 @@ class CreateMultiPipelinesTask : public Task {
 		for (auto &task : tasks_) {
 			task.pipeline->Create(vulkan_, task.compatibleRenderPass, task.rpType, task.sampleCount, task.scheduleTime, task.countToCompile);
 		}
+		tasksInFlight_.fetch_sub(1);
 	}
 
 	VulkanContext *vulkan_;
 	std::vector<SinglePipelineTask> tasks_;
+
+	// Use during shutdown to make sure there aren't any leftover tasks sitting queued.
+	// Could probably be done more elegantly. Like waiting for all tasks of a type, or saving pointers to them, or something...
+	static void WaitForAll() {
+		while (tasksInFlight_.load() > 0) {
+			sleep_ms(2);
+		}
+	}
+
+	static std::atomic<int> tasksInFlight_;
 };
 
+std::atomic<int> CreateMultiPipelinesTask::tasksInFlight_;
+
 void VulkanRenderManager::CompileThreadFunc() {
 	SetCurrentThreadName("ShaderCompile");
 	while (true) {
@@ -465,7 +480,8 @@ void VulkanRenderManager::CompileThreadFunc() {
 		for (auto &entry : toCompile) {
 			switch (entry.type) {
 			case CompileQueueEntry::Type::GRAPHICS:
-				map[std::pair< Promise<VkShaderModule> *, Promise<VkShaderModule> *>(entry.graphics->desc->vertexShader, entry.graphics->desc->fragmentShader)].push_back(
+			{
+				map[std::make_pair(entry.graphics->desc->vertexShader, entry.graphics->desc->fragmentShader)].push_back(
 					SinglePipelineTask{
 						entry.graphics,
 						entry.compatibleRenderPass,
@@ -477,6 +493,7 @@ void VulkanRenderManager::CompileThreadFunc() {
 				);
 				break;
 			}
+			}
 		}
 
 		for (auto iter : map) {
@@ -500,6 +517,7 @@ void VulkanRenderManager::DrainAndBlockCompileQueue() {
 	while (!compileQueue_.empty()) {
 		queueRunner_.WaitForCompileNotification();
 	}
+	CreateMultiPipelinesTask::WaitForAll();
 }
 
 void VulkanRenderManager::ReleaseCompileQueue() {
diff --git a/Common/GPU/Vulkan/thin3d_vulkan.cpp b/Common/GPU/Vulkan/thin3d_vulkan.cpp
index 572cb8fbc848..91cdc461e82f 100644
--- a/Common/GPU/Vulkan/thin3d_vulkan.cpp
+++ b/Common/GPU/Vulkan/thin3d_vulkan.cpp
@@ -191,9 +191,11 @@ class VKShaderModule : public ShaderModule {
 	const std::string &GetSource() const { return source_; }
 	~VKShaderModule() {
 		if (module_) {
+			INFO_LOG(G3D, "~VKShaderModule");
 			VkShaderModule shaderModule = module_->BlockUntilReady();
 			vulkan_->Delete().QueueDeleteShaderModule(shaderModule);
 			vulkan_->Delete().QueueCallback([](VulkanContext *context, void *m) {
+				INFO_LOG(G3D, "destroying shader module promise");
 				auto module = (Promise<VkShaderModule> *)m;
 				delete module;
 			}, module_);
diff --git a/Common/Thread/ThreadManager.h b/Common/Thread/ThreadManager.h
index 66d67fcee527..0b29af773043 100644
--- a/Common/Thread/ThreadManager.h
+++ b/Common/Thread/ThreadManager.h
@@ -30,6 +30,7 @@ class Task {
 	virtual void Cancel() {}
 	virtual uint64_t id() { return 0; }
 	virtual void Release() { delete this; }
+	virtual const char *Kind() const { return nullptr; }  // Useful for selecting task by some qualifier, like, waiting for them all.
 };
 
 class Waitable {