From 2a67bef38b6d3bbbac01f8f03ca83025f77a6bab Mon Sep 17 00:00:00 2001
From: Nico de Poel <ndepoel@gmail.com>
Date: Thu, 13 Mar 2025 22:09:29 +0100
Subject: [PATCH] Added lock guards to all exported functions that interact
 with the features or the graphics contexts, as well as additional validity
 checks. Protects against most of Unity's weird unsafe behavior when using the
 plugin, fixing the majority of crashes.

---
 FSR2UnityPlugin.cpp | 44 ++++++++++++++++++++++++++------------------
 1 file changed, 26 insertions(+), 18 deletions(-)

diff --git a/FSR2UnityPlugin.cpp b/FSR2UnityPlugin.cpp
index 33fa9b6..5293552 100644
--- a/FSR2UnityPlugin.cpp
+++ b/FSR2UnityPlugin.cpp
@@ -75,6 +75,8 @@ extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API UnityPluginUnload()
 
 static void UNITY_INTERFACE_API OnGraphicsDeviceEvent(UnityGfxDeviceEventType eventType)
 {
+    std::lock_guard<std::mutex> lock(s_FeatureMutex);
+
     switch (eventType)
     {
     case kUnityGfxDeviceEventInitialize:
@@ -135,16 +137,14 @@ static void UNITY_INTERFACE_API OnGraphicsDeviceEvent(UnityGfxDeviceEventType ev
     };
 }
 
-// Thread-safe allocation of feature slot
 static uint32_t AllocateFeatureSlot()
 {
-    std::lock_guard<std::mutex> lock(s_FeatureMutex);
-
     if (s_FeatureSlots.empty())
     {
         // Create a new feature if there are no free slots
         uint32_t featureSlot = (uint32_t)s_Features.size();
         s_Features.push_back(std::move(FSR2Feature()));
+        memset(&s_Features[featureSlot], 0, sizeof(FSR2Feature));
         return featureSlot;
     }
 
@@ -154,11 +154,8 @@ static uint32_t AllocateFeatureSlot()
     return featureSlot;
 }
 
-// Thread-safe freeing and clearing of feature slot
 static void FreeFeatureSlot(uint32_t featureSlot)
 {
-    std::lock_guard<std::mutex> lock(s_FeatureMutex);
-
     s_FeatureSlots.push(featureSlot);
     memset(&s_Features[featureSlot], 0, sizeof(FSR2Feature));
 }
@@ -171,6 +168,8 @@ static PFN_vkVoidFunction GetVulkanDeviceProcAddr(VkDevice device, const char* p
 
 extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_InitApi()
 {
+    std::lock_guard<std::mutex> lock(s_FeatureMutex);
+
     if (s_GraphicsD3D11 != nullptr)
     {
         ID3D11Device* device = s_GraphicsD3D11->GetDevice();
@@ -214,16 +213,23 @@ extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_InitApi()
     return false;
 }
 
+static void DestroyFeature(uint32_t featureSlot)
+{
+    auto& feature = s_Features[featureSlot];
+    if (feature.isValid)
+    {
+        ffxFsr2ContextDestroy(&feature.upscalingContext);
+        FreeFeatureSlot(featureSlot);
+    }
+}
+
 extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_ShutdownApi()
 {
+    std::lock_guard<std::mutex> lock(s_FeatureMutex);
+
     for (uint32_t slot = 0; slot < s_Features.size(); ++slot)
     {
-        auto& feature = s_Features[slot];
-        if (feature.isValid)
-        {
-            ffxFsr2ContextDestroy(&feature.upscalingContext);
-            FreeFeatureSlot(slot);
-        }
+        DestroyFeature(slot);
     }
 
     if (s_Fsr2Interface.scratchBuffer != nullptr)
@@ -252,6 +258,7 @@ extern "C" UnityRenderingEventAndData UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API
 
 extern "C" uint32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_CreateFeatureSlot()
 {
+    std::lock_guard<std::mutex> lock(s_FeatureMutex);
     return AllocateFeatureSlot();
 }
 
@@ -322,6 +329,8 @@ static void DestroyVulkanImageView(UnityVulkanInstance& instance, FSR2TextureDes
 // Plugin function to handle a specific rendering event
 static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
 {
+    std::lock_guard<std::mutex> lock(s_FeatureMutex);
+
     if (s_Device == nullptr)
         return;
 
@@ -334,12 +343,7 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
         if (featureSlot < 0 || featureSlot >= s_Features.size())
             return;
 
-        auto& feature = s_Features[featureSlot];
-        if (feature.isValid)
-        {
-            ffxFsr2ContextDestroy(&feature.upscalingContext);
-            FreeFeatureSlot(featureSlot);
-        }
+        DestroyFeature(featureSlot);
         break;
     }
     case BaseEventId + FSR2PluginEvent::eExecute:
@@ -349,6 +353,8 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
             return;
 
         auto& feature = s_Features[params->featureSlot];
+        if (!feature.isValid)
+            return;
 
         FfxFsr2DispatchDescription dispatchDescription{};
         
@@ -469,6 +475,8 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
 
 static void UNITY_INTERFACE_API OnSetTextureEvent(int eventID, void* data)
 {
+    std::lock_guard<std::mutex> lock(s_FeatureMutex);
+
     auto* params = (UnityRenderingExtTextureUpdateParamsV2*)data;
 
     // userData = (featureId & (int) ushort.MaxValue) << 16 | (textureSlot & (int) short.MaxValue) << 1 | (clearTextureTable ? 1 : 0);