From aeec7f78fd535e07606d74d555c8d0fa59007e5a Mon Sep 17 00:00:00 2001
From: Kamron Batman <3953314+kamronbatman@users.noreply.github.com>
Date: Mon, 30 May 2022 14:04:54 -0700
Subject: [PATCH] fix: Fixes issue with wepoll losing GCHandle. (#1039)

- [X] Fixes issue with wepoll losing GCHandle.
- [X] `NetState.Disconnect()` is no longer thread safe.
  - Use `Core.LoopContext.Post()` to post disconnects
- [X] Optimizes PollGroup by not processing IntPtr -> GCHandle for discard polls.
---
 .../Tests/Network/PollGroupTests.cs           | 37 ++++++++++++
 Projects/Server/Network/NetState/NetState.cs  | 58 +++++++++++--------
 Projects/Server/Server.csproj                 |  2 +-
 3 files changed, 72 insertions(+), 25 deletions(-)
 create mode 100644 Projects/Server.Tests/Tests/Network/PollGroupTests.cs

diff --git a/Projects/Server.Tests/Tests/Network/PollGroupTests.cs b/Projects/Server.Tests/Tests/Network/PollGroupTests.cs
new file mode 100644
index 000000000..34d97e9a2
--- /dev/null
+++ b/Projects/Server.Tests/Tests/Network/PollGroupTests.cs
@@ -0,0 +1,37 @@
+using System;
+using System.Runtime.InteropServices;
+using System.Threading;
+using Server.Network;
+using Xunit;
+
+namespace Server.Tests.Network;
+
+public class PollGroupTests
+{
+    [Fact]
+    public void TestPollGroup()
+    {
+        // var group = new KQueuePollGroup();
+        var nss = new NetState[2048];
+        var handles = new IntPtr[2048];
+        for (var i = 0; i < nss.Length; i++)
+        {
+            nss[i] = PacketTestUtilities.CreateTestNetState();
+            handles[i] = (IntPtr)nss[i].Handle;
+        }
+
+        GC.AddMemoryPressure(10000000000);
+        GC.Collect();
+        GC.RemoveMemoryPressure(10000000000);
+        GC.Collect();
+
+        Thread.Sleep(1000);
+
+        for (var i = 0; i < nss.Length; i++)
+        {
+            Assert.Equal(nss[i].Handle, (GCHandle)handles[i]);
+        }
+
+        // group.Dispose();
+    }
+}
diff --git a/Projects/Server/Network/NetState/NetState.cs b/Projects/Server/Network/NetState/NetState.cs
index 267a0c0da..1bd7cf40f 100755
--- a/Projects/Server/Network/NetState/NetState.cs
+++ b/Projects/Server/Network/NetState/NetState.cs
@@ -15,7 +15,6 @@
 
 using System;
 using System.Buffers;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.IO;
 using System.Net;
@@ -51,9 +50,9 @@ public partial class NetState : IComparable<NetState>
 
     private static GCHandle[] _polledStates = new GCHandle[2048];
     private static readonly IPollGroup _pollGroup = PollGroup.Create();
-    private static readonly Queue<NetState> FlushPending = new(2048);
-    private static readonly Queue<NetState> FlushedPartials = new(2048);
-    private static readonly ConcurrentQueue<NetState> Disposed = new();
+    private static readonly Queue<NetState> _flushPending = new(2048);
+    private static readonly Queue<NetState> _flushedPartials = new(256);
+    private static readonly Queue<NetState> _disposed = new(256);
 
     public static NetStateCreatedCallback CreatedCallback { get; set; }
 
@@ -75,6 +74,8 @@ public partial class NetState : IComparable<NetState>
     internal GCHandle _handle;
     private bool _packetLogging;
 
+    public GCHandle Handle => _handle;
+
     internal enum ParserState
     {
         AwaitingNextPacket,
@@ -508,7 +509,7 @@ public partial class NetState : IComparable<NetState>
 
             if (!_flushQueued)
             {
-                FlushPending.Enqueue(this);
+                _flushPending.Enqueue(this);
                 _flushQueued = true;
             }
 
@@ -930,15 +931,15 @@ public partial class NetState : IComparable<NetState>
 
     public static void FlushAll()
     {
-        while (FlushPending.Count != 0)
+        while (_flushPending.Count != 0)
         {
-            FlushPending.Dequeue()?.Flush();
+            _flushPending.Dequeue()?.Flush();
         }
     }
 
     public static void Slice()
     {
-        int count = _pollGroup.Poll(ref _polledStates);
+        int count = _pollGroup.Poll(_polledStates);
 
         if (count > 0)
         {
@@ -949,24 +950,34 @@ public partial class NetState : IComparable<NetState>
             }
         }
 
-        while (FlushPending.TryDequeue(out var ns))
+        while (_flushPending.TryDequeue(out var ns))
         {
             if (!ns.Flush())
             {
                 // Incomplete data, so we need to requeue
-                FlushedPartials.Enqueue(ns);
+                _flushedPartials.Enqueue(ns);
             }
         }
 
-        var hasDisposes = !Disposed.IsEmpty;
-        while (Disposed.TryDequeue(out var ns))
+        var hasDisposes = false;
+        while (_disposed.TryDequeue(out var ns))
         {
+            hasDisposes = true;
             ns.Dispose();
         }
 
+        // If they weren't disconnected, requeue them
+        while (_flushedPartials.TryDequeue(out var ns))
+        {
+            if (ns.Running)
+            {
+                _flushPending.Enqueue(ns);
+            }
+        }
+
         if (hasDisposes)
         {
-            _pollGroup.Poll(ref _polledStates);
+            _pollGroup.Poll(_polledStates.Length);
         }
     }
 
@@ -1025,20 +1036,19 @@ public partial class NetState : IComparable<NetState>
 
         _running = false;
 
-        try
-        {
-            if (_disconnectReason != string.Empty)
+#if THREADGUARD
+            if (Thread.CurrentThread != Core.Thread)
             {
-                throw new Exception("Attempted to disconnect a netstate twice.");
+                Utility.PushColor(ConsoleColor.Red);
+                Console.WriteLine("Attempting to disconnect a netstate from an invalid thread!");
+                Console.WriteLine(new StackTrace());
+                Utility.PopColor();
+                return;
             }
-        }
-        catch (Exception ex)
-        {
-            TraceException(ex);
-        }
+#endif
 
         _disconnectReason = reason;
-        Disposed.Enqueue(this);
+        _disposed.Enqueue(this);
     }
 
     public static void TraceDisconnect(string reason, string ip)
@@ -1088,7 +1098,7 @@ public partial class NetState : IComparable<NetState>
         TcpServer.Instances.Remove(this);
         try
         {
-            _pollGroup.Remove(Connection);
+            _pollGroup.Remove(Connection, _handle);
         }
         catch (Exception ex)
         {
diff --git a/Projects/Server/Server.csproj b/Projects/Server/Server.csproj
index 50866b6d6..1c7cadb61 100755
--- a/Projects/Server/Server.csproj
+++ b/Projects/Server/Server.csproj
@@ -37,7 +37,7 @@
     </Target>
     <ItemGroup>
         <PackageReference Include="Microsoft.Toolkit.HighPerformance" Version="7.1.2" />
-        <PackageReference Include="PollGroup" Version="1.1.0" />
+        <PackageReference Include="PollGroup" Version="1.2.1" />
         <PackageReference Include="Zlib.Bindings" Version="1.9.2" />
 
         <PackageReference Include="ModernUO.Serialization.Annotations" Version="2.0.2" />