From aa40402c28f72424d3394df7f4b15a2359e3c91f Mon Sep 17 00:00:00 2001
From: Nico de Poel <ndepoel@gmail.com>
Date: Thu, 13 Mar 2025 21:44:45 +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.

---
 FSR3UnityPlugin.cpp | 63 ++++++++++++++++++++++++---------------------
 1 file changed, 34 insertions(+), 29 deletions(-)

diff --git a/FSR3UnityPlugin.cpp b/FSR3UnityPlugin.cpp
index d04aa1d..0188ecc 100644
--- a/FSR3UnityPlugin.cpp
+++ b/FSR3UnityPlugin.cpp
@@ -87,6 +87,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:
@@ -147,11 +149,8 @@ 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
@@ -167,11 +166,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(FSR3Feature));
 }
@@ -203,6 +199,8 @@ static bool LoadFidelityFXLibrary(_In_ LPCWSTR lpLibFileName)
 
 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();
@@ -242,21 +240,28 @@ 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.upscalingContext != nullptr)
+    {
+        s_ffxFunctions.DestroyContext(&feature.upscalingContext, nullptr);
+        FreeFeatureSlot(featureSlot);
+    }
+    else if (feature.isValid)
+    {
+        ffxFsr2ContextDestroy(&feature.fsr2Context);
+        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.upscalingContext != nullptr)
-        {
-            s_ffxFunctions.DestroyContext(&feature.upscalingContext, nullptr);
-            FreeFeatureSlot(slot);
-        }
-        else if (feature.isValid)
-        {
-            ffxFsr2ContextDestroy(&feature.fsr2Context);
-            FreeFeatureSlot(slot);
-        }
+        DestroyFeature(slot);
     }
 
     if (s_ffxModule != nullptr)
@@ -297,6 +302,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();
 }
 
@@ -355,6 +361,11 @@ static FfxApiResource GetVulkanTextureResource(UnityVulkanInstance& instance, FS
 // 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_DX11Device == nullptr && s_DX12BackendDesc.device == nullptr && s_VulkanBackendDesc.vkDevice == nullptr)
+        return;
+
     // User rendering code
     switch (eventID)
     {
@@ -364,17 +375,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.upscalingContext != nullptr)
-        {
-            s_ffxFunctions.DestroyContext(&feature.upscalingContext, nullptr);
-            FreeFeatureSlot(featureSlot);
-        }
-        else if (feature.isValid)
-        {
-            ffxFsr2ContextDestroy(&feature.fsr2Context);
-            FreeFeatureSlot(featureSlot);
-        }
+        DestroyFeature(featureSlot);
         break;
     }
     case BaseEventId + FSR3PluginEvent::eExecute:
@@ -384,6 +385,8 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
             return;
 
         auto& feature = s_Features[params->featureSlot];
+        if (feature.upscalingContext == nullptr && !feature.isValid)
+            return;
 
         if (s_GraphicsD3D11 != nullptr)
         {
@@ -429,7 +432,7 @@ static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
             ffxFsr2ContextDispatch(&feature.fsr2Context, &dispatchDescription);
             return;
         }
-        
+
         // Execute FSR 3.1 or higher on DX12 and Vulkan
         ffx::DispatchDescUpscale dispatchUpscale{};
 
@@ -543,6 +546,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);