Lua: Fix occasional deadlocks on script reset

praydog · Jul 13, 2024 · a854e5a · a854e5a
1 parent 8e0f077
commit a854e5a
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 40 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -760,6 +760,7 @@ list(APPEND uevr_SOURCES
 	"src/hooks/XInputHook.hpp"
 	"src/mods/FrameworkConfig.hpp"
 	"src/mods/ImGuiThemeHelpers.hpp"
+	"src/mods/LuaLoader.hpp"
 	"src/mods/PluginLoader.hpp"
 	"src/mods/UObjectHook.hpp"
 	"src/mods/VR.hpp"

diff --git a/lua-api/lib/include/ScriptContext.hpp b/lua-api/lib/include/ScriptContext.hpp
@@ -12,7 +12,17 @@ namespace uevr {
 class ScriptContext : public std::enable_shared_from_this<ScriptContext> {
 public:
     static std::shared_ptr<ScriptContext> create(lua_State* l, UEVR_PluginInitializeParam* param = nullptr) {
-        return std::shared_ptr<ScriptContext>(new ScriptContext(l, param));
+        auto ctx = std::shared_ptr<ScriptContext>(new ScriptContext(l, param));
+        ctx->initialize();
+
+        return ctx;
+    }
+
+    static std::shared_ptr<ScriptContext> create(std::shared_ptr<sol::state> l, UEVR_PluginInitializeParam* param = nullptr) {
+        auto ctx = std::shared_ptr<ScriptContext>(new ScriptContext(l->lua_state(), param));
+        ctx->initialize(l);
+
+        return ctx;
     }
 
     ScriptContext() = delete;
@@ -85,10 +95,12 @@ class ScriptContext : public std::enable_shared_from_this<ScriptContext> {
 private:
     // Private constructor to prevent direct instantiation
     ScriptContext(lua_State* l, UEVR_PluginInitializeParam* param = nullptr);
+    void initialize(std::shared_ptr<sol::state> l = nullptr);
 
     std::vector<void*> m_callbacks_to_remove{};
 
     sol::state_view m_lua;
+    std::shared_ptr<sol::state> m_lua_shared{}; // This allows us to keep the state alive (if it was created by ScriptState)
     std::recursive_mutex m_mtx{};
     UEVR_PluginInitializeParam* m_plugin_initialize_param{nullptr};
     std::vector<sol::protected_function> m_on_pre_engine_tick_callbacks{};

diff --git a/lua-api/lib/include/ScriptState.hpp b/lua-api/lib/include/ScriptState.hpp
@@ -60,7 +60,8 @@ class ScriptState {
     auto& lua() { return m_lua; }
 
 private:
-    sol::state m_lua{};
+    std::shared_ptr<sol::state> m_lua_impl{std::make_shared<sol::state>()};
+    sol::state& m_lua{*m_lua_impl.get()};
     std::shared_ptr<ScriptContext> m_context{nullptr};
 
     GarbageCollectionData m_gc_data{};

diff --git a/lua-api/lib/src/ScriptContext.cpp b/lua-api/lib/src/ScriptContext.cpp
@@ -15,39 +15,26 @@
 namespace uevr {
 class ScriptContexts {
 public:
-    void add(ScriptContext* ctx) {
+    void add(std::shared_ptr<ScriptContext> ctx) {
         std::scoped_lock _{mtx};
-
-        // Check if the context is already in the list
-        for (ScriptContext* c : list) {
-            if (c == ctx) {
-                return;
-            }
-        }
-
         list.push_back(ctx);
     }
 
-    void remove(ScriptContext* ctx) {
-        std::scoped_lock _{mtx};
-
-        ScriptContext::log("Removing context from list");
-        std::erase_if(list, [ctx](ScriptContext* c) {
-            return c == ctx;
-        });
-        ScriptContext::log(std::format("New context count: {}", list.size()));
-    }
-
     template<typename T>
     void for_each(T&& fn) {
         std::scoped_lock _{mtx};
-        for (auto& ctx : list) {
-            fn(ctx->shared_from_this());
+        for (auto it = list.begin(); it != list.end();) {
+            if (auto ctx = it->lock()) {
+                fn(ctx);
+                ++it;
+            } else {
+                it = list.erase(it); // Naturally removes the weak_ptr from the list
+            }
         }
     }
 
 private:
-    std::vector<ScriptContext*> list{};
+    std::vector<std::weak_ptr<ScriptContext>> list;
     std::mutex mtx{};
 } g_contexts{};
 
@@ -56,8 +43,6 @@ ScriptContext::ScriptContext(lua_State* l, UEVR_PluginInitializeParam* param)
 {
     std::scoped_lock _{m_mtx};
 
-    g_contexts.add(this);
-
     if (param != nullptr) {
         m_plugin_initialize_param = param;
         uevr::API::initialize(m_plugin_initialize_param);
@@ -74,20 +59,26 @@ ScriptContext::ScriptContext(lua_State* l, UEVR_PluginInitializeParam* param)
     uevr::API::initialize(m_plugin_initialize_param);
 }
 
+void ScriptContext::initialize(std::shared_ptr<sol::state> l) {
+    m_lua_shared = l;
+    g_contexts.add(shared_from_this());
+}
+
 ScriptContext::~ScriptContext() {
     std::scoped_lock _{m_mtx};
     ScriptContext::log("ScriptContext destructor called");
 
     // TODO: this probably does not support multiple states
-    if (m_plugin_initialize_param != nullptr) {
+    // Addendum: I decided this is not necessary, for now...
+    // because all of the functions are static
+    // and this sometimes introduces a deadlock
+    /*if (m_plugin_initialize_param != nullptr) {
         for (auto& cb : m_callbacks_to_remove) {
             m_plugin_initialize_param->functions->remove_callback(cb);
         }
 
         m_callbacks_to_remove.clear();
-    }
-
-    g_contexts.remove(this);
+    }*/
 }
 
 void ScriptContext::log(const std::string& message) {
@@ -100,14 +91,18 @@ void ScriptContext::setup_callback_bindings() {
 
     auto cbs = m_plugin_initialize_param->sdk->callbacks;
 
-    add_callback(cbs->on_pre_engine_tick, on_pre_engine_tick);
-    add_callback(cbs->on_post_engine_tick, on_post_engine_tick);
-    add_callback(cbs->on_pre_slate_draw_window_render_thread, on_pre_slate_draw_window_render_thread);
-    add_callback(cbs->on_post_slate_draw_window_render_thread, on_post_slate_draw_window_render_thread);
-    add_callback(cbs->on_pre_calculate_stereo_view_offset, on_pre_calculate_stereo_view_offset);
-    add_callback(cbs->on_post_calculate_stereo_view_offset, on_post_calculate_stereo_view_offset);
-    add_callback(cbs->on_pre_viewport_client_draw, on_pre_viewport_client_draw);
-    add_callback(cbs->on_post_viewport_client_draw, on_post_viewport_client_draw);
+    // Static callbacks, so we only need to initialize them once
+    static bool init_callbacks_once = [&]() {
+        add_callback(cbs->on_pre_engine_tick, on_pre_engine_tick);
+        add_callback(cbs->on_post_engine_tick, on_post_engine_tick);
+        add_callback(cbs->on_pre_slate_draw_window_render_thread, on_pre_slate_draw_window_render_thread);
+        add_callback(cbs->on_post_slate_draw_window_render_thread, on_post_slate_draw_window_render_thread);
+        add_callback(cbs->on_pre_calculate_stereo_view_offset, on_pre_calculate_stereo_view_offset);
+        add_callback(cbs->on_post_calculate_stereo_view_offset, on_post_calculate_stereo_view_offset);
+        add_callback(cbs->on_pre_viewport_client_draw, on_pre_viewport_client_draw);
+        add_callback(cbs->on_post_viewport_client_draw, on_post_viewport_client_draw);
+        return true;
+    }();
 
     m_lua.new_usertype<UEVR_SDKCallbacks>("UEVR_SDKCallbacks",
         "on_pre_engine_tick", [this](sol::function fn) {

diff --git a/lua-api/lib/src/ScriptState.cpp b/lua-api/lib/src/ScriptState.cpp
@@ -12,7 +12,8 @@ void msg(const char* text) {
 }
 
 namespace uevr {
-ScriptState::ScriptState(const ScriptState::GarbageCollectionData& gc_data, UEVR_PluginInitializeParam* param, bool is_main_state) {
+ScriptState::ScriptState(const ScriptState::GarbageCollectionData& gc_data, UEVR_PluginInitializeParam* param, bool is_main_state)
+{
     if (param != nullptr) {
         uevr::API::initialize(param);
     }
@@ -40,7 +41,8 @@ ScriptState::ScriptState(const ScriptState::GarbageCollectionData& gc_data, UEVR
 
     // TODO: Make this actually support multiple states
     // This stores a global reference to itself, meaning it doesn't support multiple states
-    m_context = ScriptContext::create(m_lua.lua_state(), param);
+    // We pass along the shared_ptr (impl) so the context can keep it alive if for some reason the state is destroyed before the context
+    m_context = ScriptContext::create(m_lua_impl, param);
 
     if (!m_context->valid()) {
         if (param != nullptr && param->functions != nullptr) {