From 34fad0bf98af4a4fbbc8331dd33634cccfb3ef31 Mon Sep 17 00:00:00 2001
From: Connor McLaughlin <stenzek@gmail.com>
Date: Mon, 12 Dec 2022 22:17:12 +1000
Subject: [PATCH] GS: Ensure screenshots are saved before shutting down

---
 pcsx2/GS/GS.cpp                          | 18 +++---------
 pcsx2/GS/GS.h                            |  1 +
 pcsx2/GS/Renderers/Common/GSRenderer.cpp | 36 ++++++++++++++++++++++--
 3 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/pcsx2/GS/GS.cpp b/pcsx2/GS/GS.cpp
index d5d1c27281..d57f6a2063 100644
--- a/pcsx2/GS/GS.cpp
+++ b/pcsx2/GS/GS.cpp
@@ -121,20 +121,7 @@ void GSinitConfig()
 
 void GSshutdown()
 {
-#ifndef PCSX2_CORE
-	if (g_gs_renderer)
-	{
-		g_gs_renderer->Destroy();
-		g_gs_renderer.reset();
-	}
-	if (g_gs_device)
-	{
-		g_gs_device->Destroy();
-		g_gs_device.reset();
-	}
-
-	Host::ReleaseHostDisplay(true);
-#endif
+	GSclose();
 
 #ifdef _WIN32
 	if (SUCCEEDED(s_hr))
@@ -163,6 +150,9 @@ void GSclose()
 		g_host_display->SetGPUTimingEnabled(false);
 
 	Host::ReleaseHostDisplay(true);
+
+	// ensure all screenshots have been saved
+	GSJoinSnapshotThreads();
 }
 
 static RenderAPI GetAPIForRenderer(GSRendererType renderer)
diff --git a/pcsx2/GS/GS.h b/pcsx2/GS/GS.h
index 0b6dc16bcc..1363b759a6 100644
--- a/pcsx2/GS/GS.h
+++ b/pcsx2/GS/GS.h
@@ -93,6 +93,7 @@ void GSResetAPIState();
 void GSRestoreAPIState();
 bool GSSaveSnapshotToMemory(u32 window_width, u32 window_height, bool apply_aspect, bool crop_borders,
 	u32* width, u32* height, std::vector<u32>* pixels);
+void GSJoinSnapshotThreads();
 
 class GSApp
 {
diff --git a/pcsx2/GS/Renderers/Common/GSRenderer.cpp b/pcsx2/GS/Renderers/Common/GSRenderer.cpp
index 55173df792..61e42ba2db 100644
--- a/pcsx2/GS/Renderers/Common/GSRenderer.cpp
+++ b/pcsx2/GS/Renderers/Common/GSRenderer.cpp
@@ -28,6 +28,9 @@
 #include "common/Timer.h"
 #include "fmt/core.h"
 #include <array>
+#include <deque>
+#include <thread>
+#include <mutex>
 
 #ifndef PCSX2_CORE
 #include "gui/AppCoreThread.h"
@@ -63,6 +66,9 @@ static constexpr std::array<PresentShader, 6> s_tv_shader_indices = {
 	PresentShader::DIAGONAL_FILTER, PresentShader::TRIANGULAR_FILTER,
 	PresentShader::COMPLEX_FILTER, PresentShader::LOTTES_FILTER};
 
+static std::deque<std::thread> s_screenshot_threads;
+static std::mutex s_screenshot_threads_mutex;
+
 std::unique_ptr<GSRenderer> g_gs_renderer;
 
 GSRenderer::GSRenderer()
@@ -570,7 +576,8 @@ static void CompressAndWriteScreenshot(std::string filename, u32 width, u32 heig
 	Host::AddIconOSDMessage(key, ICON_FA_CAMERA, fmt::format("Saving screenshot to '{}'.", Path::GetFileName(filename)), 60.0f);
 
 	// maybe std::async would be better here.. but it's definitely worth threading, large screenshots take a while to compress.
-	std::thread compress_thread([key = std::move(key), filename = std::move(filename), image = std::move(image), quality = GSConfig.ScreenshotQuality]() {
+	std::unique_lock lock(s_screenshot_threads_mutex);
+	s_screenshot_threads.emplace_back([key = std::move(key), filename = std::move(filename), image = std::move(image), quality = GSConfig.ScreenshotQuality]() {
 		if (image.SaveToFile(filename.c_str(), quality))
 		{
 			Host::AddIconOSDMessage(std::move(key), ICON_FA_CAMERA,
@@ -581,8 +588,33 @@ static void CompressAndWriteScreenshot(std::string filename, u32 width, u32 heig
 			Host::AddIconOSDMessage(std::move(key), ICON_FA_CAMERA,
 				fmt::format("Failed to save screenshot to '{}'.", Path::GetFileName(filename), Host::OSD_ERROR_DURATION));
 		}
+
+		// remove ourselves from the list, if the GS thread is waiting for us, we won't be in there
+		const auto this_id = std::this_thread::get_id();
+		std::unique_lock lock(s_screenshot_threads_mutex);
+		for (auto it = s_screenshot_threads.begin(); it != s_screenshot_threads.end(); ++it)
+		{
+			if (it->get_id() == this_id)
+			{
+				it->detach();
+				s_screenshot_threads.erase(it);
+				break;
+			}
+		}
 	});
-	compress_thread.detach();
+}
+
+void GSJoinSnapshotThreads()
+{
+	std::unique_lock lock(s_screenshot_threads_mutex);
+	while (!s_screenshot_threads.empty())
+	{
+		std::thread save_thread(std::move(s_screenshot_threads.front()));
+		s_screenshot_threads.pop_front();
+		lock.unlock();
+		save_thread.join();
+		lock.lock();
+	}
 }
 
 void GSRenderer::VSync(u32 field, bool registers_written)