Fixed race conditions and forgotten clear in Lua ref tracking. (#3530)

This fixes occasional crashes on plugin reload.

src/Bindings/LuaState.cpp

@@ -218,7 +218,7 @@ void cLuaState::cTrackedRef::Clear(void)
 	// Free the reference:
 	lua_State * luaState = nullptr;
 	{
-		auto cs = m_CS;
+		auto cs = m_CS.load();
 		if (cs != nullptr)
 		{
 			cCSLock Lock(*cs);
@@ -246,7 +246,7 @@ void cLuaState::cTrackedRef::Clear(void)
 
 bool cLuaState::cTrackedRef::IsValid(void)
 {
-	auto cs = m_CS;
+	auto cs = m_CS.load();
 	if (cs == nullptr)
 	{
 		return false;
@@ -261,7 +261,7 @@ bool cLuaState::cTrackedRef::IsValid(void)
 
 bool cLuaState::cTrackedRef::IsSameLuaState(cLuaState & a_LuaState)
 {
-	auto cs = m_CS;
+	auto cs = m_CS.load();
 	if (cs == nullptr)
 	{
 		return false;
@@ -285,7 +285,7 @@ bool cLuaState::cTrackedRef::IsSameLuaState(cLuaState & a_LuaState)
 
 void cLuaState::cTrackedRef::Invalidate(void)
 {
-	auto cs = m_CS;
+	auto cs = m_CS.load();
 	if (cs == nullptr)
 	{
 		// Already invalid
@@ -548,6 +548,7 @@ void cLuaState::Close(void)
 		{
 			r->Invalidate();
 		}
+		m_TrackedRefs.clear();
 	}
 
 	cCanonLuaStates::Remove(*this);
@@ -2357,6 +2358,7 @@ void cLuaState::cRef::RefStack(cLuaState & a_LuaState, int a_StackPos)
 	{
 		UnRef();
 	}
+	ASSERT(cCanonLuaStates::GetCanonState(a_LuaState)->m_CS.IsLockedByCurrentThread());
 	m_LuaState = a_LuaState;
 	lua_pushvalue(a_LuaState, a_StackPos);  // Push a copy of the value at a_StackPos onto the stack
 	m_Ref = luaL_ref(a_LuaState, LUA_REGISTRYINDEX);
@@ -2370,6 +2372,7 @@ void cLuaState::cRef::UnRef(void)
 {
 	if (IsValid())
 	{
+		ASSERT(cCanonLuaStates::GetCanonState(m_LuaState)->m_CS.IsLockedByCurrentThread());
 		luaL_unref(m_LuaState, LUA_REGISTRYINDEX, m_Ref);
 	}
 	m_LuaState = nullptr;

src/Bindings/LuaState.h

@@ -35,6 +35,7 @@ extern "C"
 	#include "lua/src/lauxlib.h"
 }
 
+#include <atomic>
 #include "../Vector3.h"
 #include "../Defines.h"
 #include "PluginManager.h"
@@ -189,24 +190,29 @@ public:
 		}
 
 		/** Set the contained reference to the object at the specified Lua state's stack position.
-		If another reference has been previously contained, it is freed first. */
+		If another reference has been previously contained, it is Clear()-ed first. */
 		bool RefStack(cLuaState & a_LuaState, int a_StackPos);
 
-		/** Frees the contained reference, if any. */
+		/** Frees the contained reference, if any.
+		Untracks the reference from its canon Lua state. */
 		void Clear(void);
 
-		/** Returns true if the contained reference is valid. */
+		/** Returns true if the contained reference is valid.
+		(Note that depending on this value is not thread-safe, another thread may invalidate the ref in the meantime. It is meant for quick ASSERTs only). */
 		bool IsValid(void);
 
 		/** Returns true if the reference resides in the specified Lua state.
-		Internally, compares the reference's canon Lua state. */
+		Internally, compares the reference's canon Lua state.
+		(Note that depending on this value is not thread-safe, another thread may modify the ref in the meantime. It is meant for quick ASSERTs only). */
 		bool IsSameLuaState(cLuaState & a_LuaState);
 
 	protected:
 		friend class cLuaState;
 
-		/** The mutex protecting m_Ref against multithreaded access */
+		/** The mutex protecting m_Ref against multithreaded access.
-		cCriticalSection * m_CS;
+		Actually points to the canon Lua state's m_CriticalSection.
+		Is nullptr when ref is empty (not bound). */
+		std::atomic<cCriticalSection *> m_CS;
 
 		/** Reference to the Lua callback */
 		cRef m_Ref;
@@ -254,7 +260,7 @@ public:
 		template <typename... Args>
 		bool Call(Args &&... args)
 		{
-			auto cs = m_CS;
+			auto cs = m_CS.load();
 			if (cs == nullptr)
 			{
 				return false;
@@ -336,7 +342,7 @@ public:
 		template <typename... Args>
 		bool CallTableFn(const char * a_FnName, Args &&... args)
 		{
-			auto cs = m_CS;
+			auto cs = m_CS.load();
 			if (cs == nullptr)
 			{
 				return false;
@@ -356,7 +362,7 @@ public:
 		template <typename... Args>
 		bool CallTableFnWithSelf(const char * a_FnName, Args &&... args)
 		{
-			auto cs = m_CS;
+			auto cs = m_CS.load();
 			if (cs == nullptr)
 			{
 				return false;

src/Generating/PrefabPiecePool.cpp

@@ -204,6 +204,7 @@ bool cPrefabPiecePool::LoadFromCubeset(const AString & a_Contents, const AString
 	// Load the file in the Lua interpreter:
 	cLuaState Lua(Printf("LoadablePiecePool %s", a_FileName.c_str()));
 	Lua.Create();
+	cLuaState::cLock lock(Lua);
 	if (!Lua.LoadString(a_Contents, a_FileName, a_LogWarnings))
 	{
 		// Reason for failure has already been logged in LoadFile()