From e9f986f55b34ea76fb0c2ff1084e043131fcb53c Mon Sep 17 00:00:00 2001
From: Kamron Batman <3953314+kamronbatman@users.noreply.github.com>
Date: Sat, 26 Mar 2022 13:57:42 -0700
Subject: [PATCH] fix: Adds Gen2 callback for each STArrayPool (#971)

Adds. Gen2Callback for STArrayPool so it can purge the internal array stacks properly. This only happens if significant timed has passed, 2 Gen 2's have been run, and memory pressure exceeds a threshold.
---
 Projects/Server/Buffers/STArrayPool.cs        |  8 ++
 .../GarbageCollection/Gen2GcCallback.cs       | 75 +++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 Projects/Server/GarbageCollection/Gen2GcCallback.cs

diff --git a/Projects/Server/Buffers/STArrayPool.cs b/Projects/Server/Buffers/STArrayPool.cs
index 72ddca91b..8210aa3de 100644
--- a/Projects/Server/Buffers/STArrayPool.cs
+++ b/Projects/Server/Buffers/STArrayPool.cs
@@ -6,6 +6,7 @@ using System.Buffers;
 using System.Diagnostics;
 using System.Numerics;
 using System.Runtime.CompilerServices;
+using System.Threading;
 
 namespace Server.Buffers;
 
@@ -20,6 +21,7 @@ public class STArrayPool<T> : ArrayPool<T>
 
     public static STArrayPool<T> Shared => _shared;
 
+    private int _trimCallbackCreated;
     private static STArray[] _cacheBuckets;
     private STArrayStack[] _buckets = new STArrayStack[BucketCount];
 
@@ -172,6 +174,12 @@ public class STArrayPool<T> : ArrayPool<T>
     {
         Debug.Assert(_cacheBuckets is null, $"Non-null {nameof(_cacheBuckets)}");
         var buckets = new STArray[BucketCount];
+
+        if (Interlocked.Exchange(ref _trimCallbackCreated, 1) == 0)
+        {
+            Gen2GcCallback.Register(o => ((STArrayPool<T>)o).Trim(), this);
+        }
+
         return _cacheBuckets = buckets;
     }
 
diff --git a/Projects/Server/GarbageCollection/Gen2GcCallback.cs b/Projects/Server/GarbageCollection/Gen2GcCallback.cs
new file mode 100644
index 000000000..fc23db8ce
--- /dev/null
+++ b/Projects/Server/GarbageCollection/Gen2GcCallback.cs
@@ -0,0 +1,75 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Diagnostics;
+using System.Runtime.ConstrainedExecution;
+using System.Runtime.InteropServices;
+
+namespace System;
+
+/// <summary>
+/// Schedules a callback roughly every gen 2 GC (you may see a Gen 0 an Gen 1 but only once)
+/// (We can fix this by capturing the Gen 2 count at startup and testing, but I mostly don't care)
+/// </summary>
+internal sealed class Gen2GcCallback : CriticalFinalizerObject
+{
+    private readonly Func<object, bool> _callback;
+    private GCHandle _weakTargetObj;
+
+    private Gen2GcCallback(Func<object, bool> callback, object targetObj)
+    {
+        _callback = callback;
+        _weakTargetObj = GCHandle.Alloc(targetObj, GCHandleType.Weak);
+    }
+
+    /// <summary>
+    /// Schedule 'callback' to be called in the next GC.  If the callback returns true it is
+    /// rescheduled for the next Gen 2 GC.  Otherwise the callbacks stop.
+    ///
+    /// NOTE: This callback will be kept alive until either the callback function returns false,
+    /// or the target object dies.
+    /// </summary>
+    public static void Register(Func<object, bool> callback, object targetObj)
+    {
+        // Create a unreachable object that remembers the callback function and target object.
+        new Gen2GcCallback(callback, targetObj);
+    }
+
+    ~Gen2GcCallback()
+    {
+        if (_weakTargetObj.IsAllocated)
+        {
+            // Check to see if the target object is still alive.
+            object? targetObj = _weakTargetObj.Target;
+            if (targetObj == null)
+            {
+                // The target object is dead, so this callback object is no longer needed.
+                _weakTargetObj.Free();
+                return;
+            }
+
+            // Execute the callback method.
+            try
+            {
+                Debug.Assert(_callback != null);
+                if (_callback?.Invoke(targetObj) != true)
+                {
+                    // If the callback returns false, this callback object is no longer needed.
+                    _weakTargetObj.Free();
+                    return;
+                }
+            }
+            catch
+            {
+                // Ensure that we still get a chance to resurrect this object, even if the callback throws an exception.
+#if DEBUG
+                    // Except in DEBUG, as we really shouldn't be hitting any exceptions here.
+                    throw;
+#endif
+            }
+        }
+
+        // Resurrect ourselves by re-registering for finalization.
+        GC.ReRegisterForFinalize(this);
+    }
+}