From ed6a8c7daf5b510f046028f787e26018a220cfc8 Mon Sep 17 00:00:00 2001
From: Quahu <quahuu@gmail.com>
Date: Wed, 6 Mar 2024 00:06:35 +0100
Subject: [PATCH] Reverted reference disposal to the original mechanism.

---
 .../Tests/Entities/LuaReferenceTests.cs       |  3 +--
 .../Entities/Reference/LuaReference.cs        |  2 +-
 .../Marshaler/LuaMarshaler.Internal.cs        | 11 ++++++--
 src/Laylua/Library/Marshaler/LuaMarshaler.cs  | 27 ++++++++++++++++---
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/src/Laylua.Tests/Tests/Entities/LuaReferenceTests.cs b/src/Laylua.Tests/Tests/Entities/LuaReferenceTests.cs
index 4bcc86d..93f0f92 100644
--- a/src/Laylua.Tests/Tests/Entities/LuaReferenceTests.cs
+++ b/src/Laylua.Tests/Tests/Entities/LuaReferenceTests.cs
@@ -1,4 +1,4 @@
-using System;
+using System;
 using Laylua.Moon;
 using NUnit.Framework;
 
@@ -75,7 +75,6 @@ public unsafe void NoReferencesToAliveObject_MarshalerFiresLeakedReferenceEvent(
 
         // Assert
         Assert.That(leakedReference, Is.EqualTo(reference));
-        return;
 
         int CreateTable()
         {
diff --git a/src/Laylua/Library/Entities/Reference/LuaReference.cs b/src/Laylua/Library/Entities/Reference/LuaReference.cs
index 56a07b3..c59cc12 100644
--- a/src/Laylua/Library/Entities/Reference/LuaReference.cs
+++ b/src/Laylua/Library/Entities/Reference/LuaReference.cs
@@ -70,7 +70,7 @@ private protected LuaReference()
         if (LuaRegistry.IsPersistentReference(_reference) || _lua == null)
             return;
 
-        _lua.Marshaler.ReturnReference(this);
+        _lua.Marshaler.OnReferenceCollected(this);
     }
 
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
diff --git a/src/Laylua/Library/Marshaler/LuaMarshaler.Internal.cs b/src/Laylua/Library/Marshaler/LuaMarshaler.Internal.cs
index f8c1794..6dd7d19 100644
--- a/src/Laylua/Library/Marshaler/LuaMarshaler.Internal.cs
+++ b/src/Laylua/Library/Marshaler/LuaMarshaler.Internal.cs
@@ -6,15 +6,22 @@ public abstract partial class LuaMarshaler
 {
     protected internal abstract void RemoveUserDataHandle(UserDataHandle handle);
 
-    internal void ReturnReference(LuaReference reference)
+    internal void OnReferenceCollected(LuaReference reference)
     {
         if (LuaReference.IsAlive(reference))
         {
             ReferenceLeaked?.Invoke(this, new LuaReferenceLeakedEventArgs(reference));
 
-            reference.Dispose();
+            _leakedReferences.Push(reference);
         }
+        else
+        {
+            ReturnReference(reference);
+        }
+    }
 
+    private void ReturnReference(LuaReference reference)
+    {
         if (_entityPool.Return(reference))
         {
             GC.ReRegisterForFinalize(reference);
diff --git a/src/Laylua/Library/Marshaler/LuaMarshaler.cs b/src/Laylua/Library/Marshaler/LuaMarshaler.cs
index 6c24a4d..52d731d 100644
--- a/src/Laylua/Library/Marshaler/LuaMarshaler.cs
+++ b/src/Laylua/Library/Marshaler/LuaMarshaler.cs
@@ -1,4 +1,5 @@
-using System;
+using System;
+using System.Collections.Concurrent;
 using System.Runtime.CompilerServices;
 using System.Threading;
 using Laylua.Moon;
@@ -34,8 +35,10 @@ public UserDataDescriptorProvider UserDataDescriptorProvider
     ///     You can utilize this event to find out if your application
     ///     is failing to correctly dispose all <see cref="LuaReference"/> instances.
     ///     <para/>
-    ///     Subscribed event handlers should be as lightweight as possible
-    ///     and must not throw any exceptions.
+    ///     Subscribed event handlers should be lightweight and must not throw exceptions.
+    ///     <para/>
+    ///     Subscribed handlers must not perform any Lua interactions,
+    ///     as they might corrupt the Lua state.
     /// </remarks>
     public event EventHandler<LuaReferenceLeakedEventArgs>? ReferenceLeaked;
 
@@ -53,6 +56,7 @@ public LuaMarshalerEntityPoolConfiguration EntityPoolConfiguration
 
     private UserDataDescriptorProvider _userDataDescriptorProvider = UserDataDescriptorProvider.Default;
     private LuaReferencePool _entityPool = null!;
+    private readonly ConcurrentStack<LuaReference> _leakedReferences = new();
 
     /// <summary>
     ///     Instantiates a new marshaler with the specified Lua instance.
@@ -77,6 +81,7 @@ protected LuaMarshaler(Lua lua)
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     protected LuaTable CreateTable(int reference)
     {
+        DisposeLeakedReferences();
         return _entityPool.RentTable(reference);
     }
 
@@ -90,6 +95,7 @@ protected LuaTable CreateTable(int reference)
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     protected LuaFunction CreateFunction(int reference)
     {
+        DisposeLeakedReferences();
         return _entityPool.RentFunction(reference);
     }
 
@@ -104,6 +110,7 @@ protected LuaFunction CreateFunction(int reference)
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     protected LuaUserData CreateUserData(int reference, IntPtr ptr)
     {
+        DisposeLeakedReferences();
         return _entityPool.RentUserData(reference, ptr);
     }
 
@@ -118,9 +125,21 @@ protected LuaUserData CreateUserData(int reference, IntPtr ptr)
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     protected unsafe LuaThread CreateThread(int reference, lua_State* L)
     {
+        DisposeLeakedReferences();
         return _entityPool.RentThread(reference, L);
     }
 
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    private void DisposeLeakedReferences()
+    {
+        while (_leakedReferences.TryPop(out var reference))
+        {
+            reference.Dispose();
+
+            ReturnReference(reference);
+        }
+    }
+
     /// <summary>
     ///     Tries to convert the Lua value at the specified stack index
     ///     to a .NET value of type <typeparamref name="T"/>.
@@ -158,6 +177,8 @@ protected virtual void Dispose(bool disposing)
     /// </summary>
     public void Dispose()
     {
+        DisposeLeakedReferences();
+
         Dispose(true);
         GC.SuppressFinalize(this);
     }