From 21dd779116074f5331887224a631832c5a0c073a Mon Sep 17 00:00:00 2001
From: Eugen <97041180+eugen15@users.noreply.github.com>
Date: Wed, 6 Sep 2023 12:42:23 -0400
Subject: [PATCH] Notify crash handler about GPU reset exception

---
 obs-studio-server/source/nodeobs_api.cpp | 52 +++++++++++++++++++++++-
 1 file changed, 51 insertions(+), 1 deletion(-)

diff --git a/obs-studio-server/source/nodeobs_api.cpp b/obs-studio-server/source/nodeobs_api.cpp
index 327ab028d..7a56f7fb9 100644
--- a/obs-studio-server/source/nodeobs_api.cpp
+++ b/obs-studio-server/source/nodeobs_api.cpp
@@ -73,6 +73,7 @@
 #endif
 
 #ifdef WIN32
+#include <comdef.h>
 #include <audiopolicy.h>
 #include <mmdeviceapi.h>
 
@@ -99,7 +100,14 @@
 #define GBYTE (1024ULL * 1024ULL * 1024ULL)
 #define TBYTE (1024ULL * 1024ULL * 1024ULL * 1024ULL)
 
-enum crashHandlerCommand { REGISTER = 0, UNREGISTER = 1, REGISTERMEMORYDUMP = 2, CRASHWITHCODE = 3 };
+enum crashHandlerCommand
+{
+	REGISTER = 0,
+	UNREGISTER = 1,
+	REGISTERMEMORYDUMP = 2,
+	CRASHWITHCODE = 3,
+	HANDLEOUTOFGPU = 4
+};
 
 struct NodeOBSLogParam final {
 	std::fstream logStream;
@@ -698,6 +706,34 @@ std::vector<char> crashedProcess(uint32_t crash_id)
 	return buffer;
 }
 
+#if defined(_WIN32)
+static std::vector<char> handleOutOfGPU(int category, unsigned long long code)
+{
+	std::vector<char> buffer;
+
+	// Prepare
+	std::uint8_t messageAction = crashHandlerCommand::HANDLEOUTOFGPU;
+	std::uint64_t messageCode(code);
+	std::wstring messageCodeDesc = (code) ? _com_error(static_cast<HRESULT>(code)).ErrorMessage() : L"Undefined error";
+	std::uint32_t messageCodeDescSize = (messageCodeDesc.size() + 1) * sizeof(wchar_t);
+
+	buffer.resize(sizeof(messageAction) + sizeof(messageCode) + sizeof(std::uint32_t) + messageCodeDescSize);
+
+	// Pack	
+	uint32_t offset = 0;
+	memcpy(buffer.data(), &messageAction, sizeof(messageAction));
+	offset++;
+	memcpy(buffer.data() + offset, &messageCode, sizeof(messageCode));
+	offset += sizeof(messageCode);
+	memcpy(buffer.data() + offset, &messageCodeDescSize, sizeof(messageCodeDescSize));
+	offset += sizeof(messageCodeDescSize);
+	memcpy(buffer.data() + offset, messageCodeDesc.data(), messageCode);
+	offset += messageCodeDescSize;
+
+	return buffer;
+}
+#endif
+
 #ifdef WIN32
 std::wstring crash_handler_pipe;
 
@@ -747,6 +783,14 @@ void writeCrashHandler(std::vector<char> buffer)
 }
 #endif
 
+static void handleGSError(void* param, int category, unsigned long long code)
+{
+#if defined(_WIN32)
+	writeCrashHandler(handleOutOfGPU(category, code));
+#endif
+	blog(LOG_ERROR, ">>> PROCESSED A CRITICAL ERROR");
+}
+
 static bool checkIfDebugLogsEnabled(const std::string &appdata)
 {
 #if defined(_DEBUG)
@@ -908,6 +952,12 @@ void OBS_API::OBS_API_initAPI(void *data, const int64_t id, const std::vector<ip
 	SetPrivilegeForGPUPriority();
 #endif
 
+	blog(LOG_INFO, ">>> BEFORE CALLBACK");
+
+	obs_set_gs_error_handler(handleGSError, nullptr);
+
+	blog(LOG_INFO, ">>> AFTER CALLBACK");
+
 	osn::Source::initialize_global_signals();
 
 	cpuUsageInfo = os_cpu_usage_info_start();