fix: fix Skyrim hanging on exit

Adds more explicit checks for cache stop_token message and also includes
a thread killing function if it has not stopped cleanly.

closes doodlum#296

src/ShaderCache.cpp

@@ -1476,8 +1476,20 @@ namespace SIE
 
 	ShaderCache::~ShaderCache()
 	{
+		logger::debug("Trying to stop pool with {}/{}/{} running/queued/total", compilationPool.get_tasks_running(), compilationPool.get_tasks_queued(), compilationPool.get_tasks_total());
+		if (!ssource.request_stop())
+			logger::warn("Unable to request a stop");
+		logger::debug("Stop requested: {}", stoken.stop_requested());
+		compilationSet.conditionVariable.notify_all();
+		compilationPool.purge();
 		Clear();
 		StopFileWatcher();
+		if (!compilationPool.wait_for(std::chrono::milliseconds(1000))) {
+			logger::debug("Tasks still running; killing thread!");
+			WaitForSingleObject(managementThread, 1000);
+			TerminateThread(managementThread, 0);
+			CloseHandle(managementThread);
+		}
 	}
 
 	void ShaderCache::Clear()
@@ -1687,8 +1699,6 @@ namespace SIE
 
 	ShaderCache::ShaderCache()
 	{
-		logger::debug("ShaderCache initialized with {} compiler threads", (int)compilationThreadCount);
-		compilationPool.push_task(&ShaderCache::ManageCompilationSet, this, ssource.get_token());
 	}
 
 	bool ShaderCache::UseFileWatcher() const
@@ -1700,12 +1710,23 @@ namespace SIE
 	{
 		auto oldValue = useFileWatcher;
 		useFileWatcher = value;
+		if (!isReady)
+			return;
 		if (useFileWatcher && !oldValue)
 			StartFileWatcher();
 		else if (!useFileWatcher && oldValue)
 			StopFileWatcher();
 	}
 
+	// Whether the threadpool is ready to process tasks
+	void ShaderCache::Init()
+	{
+		logger::debug("ShaderCache initialized with {} compiler threads", (int)compilationThreadCount);
+		stoken = ssource.get_token();
+		compilationPool.detach_task([this] { this->ManageCompilationSet(); });
+		isReady = true;
+	}
+
 	void ShaderCache::StartFileWatcher()
 	{
 		logger::info("Starting FileWatcher");
@@ -1722,7 +1743,9 @@ namespace SIE
 				pathStr += std::format("{}; ", path);
 			}
 			logger::debug("ShaderCache watching for changes in {}", pathStr);
-			compilationPool.push_task(&SIE::UpdateListener::processQueue, listener);
+			compilationPool.detach_task([this] {
+				listener->processQueue();
+			});
 		} else {
 			logger::debug("ShaderCache already enabled");
 		}
@@ -1857,7 +1880,7 @@ namespace SIE
 		compilationSet.cacheHitTasks++;
 	}
 
-	bool ShaderCache::IsHideErrors()
+	bool ShaderCache::IsHideErrors() const
 	{
 		return hideError;
 	}
@@ -1906,22 +1929,31 @@ namespace SIE
 		logger::debug("Stopped blocking shaders");
 	}
 
-	void ShaderCache::ManageCompilationSet(std::stop_token stoken)
+	void ShaderCache::ManageCompilationSet()
 	{
-		SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_BELOW_NORMAL);
+		managementThread = GetCurrentThread();
+		SetThreadPriority(managementThread, THREAD_PRIORITY_BELOW_NORMAL);
 		while (!stoken.stop_requested()) {
-			const auto& task = compilationSet.WaitTake(stoken);
-			if (!task.has_value())
-				break;  // exit because thread told to end
-			compilationPool.push_task(&ShaderCache::ProcessCompilationSet, this, stoken, task.value());
+			const auto& task = compilationSet.WaitTake();
+			if (task.has_value()) {
+				compilationPool.detach_task([this, &task] {
+					this->ProcessCompilationSet(task.value());
+				});
+			}
+			logger::debug("Pool running with {} tasks, stop {}", compilationPool.get_tasks_running(), stoken.stop_requested());
 		}
+		logger::debug("Pool trying to stop with {} tasks, stop {}", compilationPool.get_tasks_running(), stoken.stop_requested());
 	}
 
-	void ShaderCache::ProcessCompilationSet(std::stop_token stoken, SIE::ShaderCompilationTask task)
+	void ShaderCache::ProcessCompilationSet(SIE::ShaderCompilationTask task)
 	{
 		SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_BELOW_NORMAL);
-		task.Perform();
-		compilationSet.Complete(task);
+		if (!stoken.stop_requested()) {
+			task.Perform();
+			compilationSet.Complete(task);
+			logger::debug("Compilation completed: Pool running with {} tasks, stop {}", compilationPool.get_tasks_running(), stoken.stop_requested());
+		}
+		compilationSet.conditionVariable.notify_one();
 	}
 
 	ShaderCompilationTask::ShaderCompilationTask(ShaderClass aShaderClass,
@@ -1956,25 +1988,29 @@ namespace SIE
 		return GetId() == other.GetId();
 	}
 
-	std::optional<ShaderCompilationTask> CompilationSet::WaitTake(std::stop_token stoken)
+	std::optional<ShaderCompilationTask> CompilationSet::WaitTake()
 	{
-		std::unique_lock lock(compilationMutex);
 		auto& shaderCache = ShaderCache::Instance();
-		if (!conditionVariable.wait(
-				lock, stoken,
-				[this, &shaderCache]() { return !availableTasks.empty() &&
-			                                    // check against all tasks in queue to trickle the work. It cannot be the active tasks count because the thread pool itself is maximum.
-			                                    (int)shaderCache.compilationPool.get_tasks_total() <=
-			                                        (!shaderCache.backgroundCompilation ? shaderCache.compilationThreadCount : shaderCache.backgroundCompilationThreadCount); })) {
-			/*Woke up because of a stop request. */
+		std::unique_lock lock(compilationMutex);
+		conditionVariable.wait(
+			lock,
+			[this, &shaderCache]() { return shaderCache.stoken.stop_requested() || (!availableTasks.empty() &&
+																					   // check against all tasks in queue to trickle the work. It cannot be the active tasks count because the thread pool itself is maximum.
+																					   (int)shaderCache.compilationPool.get_tasks_total() <=
+																						   (!shaderCache.backgroundCompilation ? shaderCache.compilationThreadCount : shaderCache.backgroundCompilationThreadCount)); });
+		if (shaderCache.stoken.stop_requested()) {  // this appears to never fire so we have to manually kill it
+			logger::debug("Returning null since stop requested");
 			return std::nullopt;
 		}
-		if (!ShaderCache::Instance().IsCompiling()) {  // we just got woken up because there's a task, start clock
+		if (!shaderCache.IsCompiling()) {  // we just got woken up because there's a task, start clock
 			lastCalculation = lastReset = high_resolution_clock::now();
 		}
 		auto node = availableTasks.extract(availableTasks.begin());
+		if (node.empty())
+			return std::nullopt;
 		auto& task = node.value();
 		tasksInProgress.insert(std::move(node));
+		logger::debug("Pool extracted task: {}", task.GetString());
 		return task;
 	}
 
@@ -2011,7 +2047,6 @@ namespace SIE
 		std::scoped_lock lock(compilationMutex);
 		processedTasks.insert(task);
 		tasksInProgress.erase(task);
-		conditionVariable.notify_one();
 	}
 
 	void CompilationSet::Clear()

src/ShaderCache.h

@@ -62,7 +62,7 @@ namespace SIE
 	class CompilationSet
 	{
 	public:
-		std::optional<ShaderCompilationTask> WaitTake(std::stop_token stoken);
+		std::optional<ShaderCompilationTask> WaitTake();
 		void Add(const ShaderCompilationTask& task);
 		void Complete(const ShaderCompilationTask& task);
 		void Clear();
@@ -74,12 +74,12 @@ namespace SIE
 		std::atomic<uint64_t> failedTasks = 0;
 		std::atomic<uint64_t> cacheHitTasks = 0;  // number of compiles of a previously seen shader combo
 		std::mutex compilationMutex;
+		std::condition_variable conditionVariable;
 
 	private:
 		std::unordered_set<ShaderCompilationTask> availableTasks;
 		std::unordered_set<ShaderCompilationTask> tasksInProgress;
 		std::unordered_set<ShaderCompilationTask> processedTasks;  // completed or failed
-		std::condition_variable_any conditionVariable;
 		std::chrono::steady_clock::time_point lastReset = high_resolution_clock::now();
 		std::chrono::steady_clock::time_point lastCalculation = high_resolution_clock::now();
 		double totalMs = (double)duration_cast<std::chrono::milliseconds>(lastReset - lastReset).count();
@@ -145,6 +145,7 @@ namespace SIE
 		bool UseFileWatcher() const;
 		void SetFileWatcher(bool value);
 
+		void Init();
 		void StartFileWatcher();
 		void StopFileWatcher();
 
@@ -189,7 +190,7 @@ namespace SIE
 		void ToggleErrorMessages();
 		void DisableShaderBlocking();
 		void IterateShaderBlock(bool a_forward = true);
-		bool IsHideErrors();
+		bool IsHideErrors() const;
 
 		void InsertModifiedShaderMap(std::string a_shader, std::chrono::time_point<std::chrono::system_clock> a_time);
 		std::chrono::time_point<std::chrono::system_clock> GetModifiedShaderMapTime(std::string a_shader);
@@ -393,11 +394,13 @@ namespace SIE
 		uint blockedKeyIndex = (uint)-1;  // index in shaderMap; negative value indicates disabled
 		std::string blockedKey = "";
 		std::vector<uint32_t> blockedIDs;  // more than one descriptor could be blocked based on shader hash
+		void ManageCompilationSet();
+		void ProcessCompilationSet(SIE::ShaderCompilationTask task);
+		std::stop_token stoken;
+		HANDLE managementThread = nullptr;
 
 	private:
 		ShaderCache();
-		void ManageCompilationSet(std::stop_token stoken);
-		void ProcessCompilationSet(std::stop_token stoken, SIE::ShaderCompilationTask task);
 
 		~ShaderCache();
 
@@ -414,8 +417,8 @@ namespace SIE
 		bool isDump = false;
 		bool hideError = false;
 		bool useFileWatcher = false;
-
-		std::stop_source ssource;
+		bool isReady = false;
+		std::stop_source ssource{};
 		std::mutex vertexShadersMutex;
 		std::mutex pixelShadersMutex;
 		CompilationSet compilationSet;

src/State.cpp

@@ -207,8 +207,6 @@ static const std::string& GetConfigPath(State::ConfigMode a_configMode)
 void State::Load(ConfigMode a_configMode)
 {
 	ConfigMode configMode = a_configMode;
-	auto& shaderCache = SIE::ShaderCache::Instance();
-
 	std::string configPath = GetConfigPath(configMode);
 	std::ifstream i(configPath);
 	if (!i.is_open()) {
@@ -239,6 +237,8 @@ void State::Load(ConfigMode a_configMode)
 		return;
 	}
 
+	auto& shaderCache = SIE::ShaderCache::Instance();
+
 	if (settings["Menu"].is_object()) {
 		Menu::GetSingleton()->Load(settings["Menu"]);
 	}

src/XSEPlugin.cpp

@@ -96,7 +96,7 @@ void MessageHandler(SKSE::MessagingInterface::Message* message)
 				FrameAnnotations::OnPostPostLoad();
 
 				auto& shaderCache = SIE::ShaderCache::Instance();
-
+				shaderCache.Init();
 				shaderCache.ValidateDiskCache();
 				if (shaderCache.UseFileWatcher())
 					shaderCache.StartFileWatcher();

vcpkg.json

@@ -30,7 +30,7 @@
   "overrides": [
     {
       "name": "bshoshany-thread-pool",
-      "version-string": "3.5.0"
+      "version-string": "4.1.0"
     }
   ],
   "builtin-baseline": "83972272512ce4ede5fc3b2ba98f6468b179f192"