From 3df1e86ecb206485a7a2ac3284d28d7525e48dce Mon Sep 17 00:00:00 2001
From: Ari Vuollet <ari.vuollet@kapsi.fi>
Date: Sun, 27 Aug 2023 13:18:32 +0300
Subject: [PATCH] _gchandle weak v2

---
 Source/Engine/Engine/NativeInterop.Managed.cs | 78 ++++++++++++-------
 1 file changed, 49 insertions(+), 29 deletions(-)

diff --git a/Source/Engine/Engine/NativeInterop.Managed.cs b/Source/Engine/Engine/NativeInterop.Managed.cs
index 8f63b788d..e2a46b348 100644
--- a/Source/Engine/Engine/NativeInterop.Managed.cs
+++ b/Source/Engine/Engine/NativeInterop.Managed.cs
@@ -419,7 +419,7 @@ namespace FlaxEngine.Interop
             // Rolling numbers for handles, two bits reserved for the type
             //private static ulong normalHandleAccumulator = ((ulong)GCHandleType.Normal << 62) & 0xC000000000000000;
             //private static ulong pinnedHandleAccumulator = ((ulong)GCHandleType.Pinned << 62) & 0xC000000000000000;
-            //private static ulong weakHandleAccumulator = ((ulong)GCHandleType.Weak << 62) & 0xC000000000000000;
+            private static ulong weakHandleAccumulator = ((ulong)GCHandleType.Normal << 62) & 0xC000000000000000;
 
             // Dictionaries for storing the valid handles.
             // Note: Using locks seems to be generally the fastest when adding or fetching from the dictionary.
@@ -432,8 +432,8 @@ namespace FlaxEngine.Interop
 
             // Manage double-buffered pool for weak handles in order to avoid collecting in-flight handles.
             // Periodically when the pools are being accessed and conditions are met, the other pool is cleared and swapped.
-            private static HashSet<IntPtr> weakPool = new();
-            private static HashSet<IntPtr> weakPoolOther = new();
+            private static Dictionary<IntPtr, object> weakPool = new();
+            private static Dictionary<IntPtr, object> weakPoolOther = new();
             private static object weakPoolLock = new object();
             private static ulong nextWeakPoolCollection;
             private static ulong weakPoolCycle;
@@ -451,12 +451,12 @@ namespace FlaxEngine.Interop
 
                 nextWeakPoolCollection = 0;
 
-                //ulong currentCycle = weakPoolCycle;
+                ulong currentCycle = weakPoolCycle;
 
-                //lock (weakPoolLock)
+                lock (weakPoolLock)
                 {
-                    //if (currentCycle != weakPoolCycle)
-                    //    return;
+                    if (currentCycle != weakPoolCycle)
+                        return;
 
                     weakPoolCycle++;
 
@@ -471,13 +471,6 @@ namespace FlaxEngine.Interop
                         return;
                     lastWeakPoolCollectionTime = System.Diagnostics.Stopwatch.GetTimestamp();
 
-                    foreach (IntPtr ptr in weakPoolOther)
-                    {
-                        GCHandle gcHandle = GCHandle.FromIntPtr(ptr);
-                        
-                        gcHandle.Free();
-                    }
-
                     // Swap the pools and release the oldest pool for GC
                     (weakPool, weakPoolOther) = (weakPoolOther, weakPool);
                     weakPool.Clear();
@@ -498,24 +491,26 @@ namespace FlaxEngine.Interop
             //[MethodImpl(MethodImplOptions.AggressiveInlining)]
             //private static GCHandleType GetHandleType(IntPtr handle) => (GCHandleType)(((ulong)handle & 0xC000000000000000) >> 62);
 
+            private static bool IsWeakHandle(IntPtr handle) => ((ulong)handle & 0xC000000000000000) >> 62 == 2;
+
             internal static IntPtr AllocateHandle(object value, GCHandleType type)
             {
-                GCHandleType type2 = type;
                 if (type == GCHandleType.Weak || type == GCHandleType.WeakTrackResurrection)
                 {
-                    type2 = GCHandleType.Normal;
+
+                    lock (weakPoolLock)
+                    {
+                        IntPtr handle = (IntPtr)(++weakHandleAccumulator);
+                        weakPool.Add(handle, value);
+                        return handle;
+                    }
                 }
-
-                GCHandle gcHandle = GCHandle.Alloc(value, type2);
-                IntPtr handle = GCHandle.ToIntPtr(gcHandle);
-
-                if (type == GCHandleType.Weak || type == GCHandleType.WeakTrackResurrection)
+                else
                 {
-                    //TryCollectWeakHandles();
-                    weakPool.Add(handle/*, gcHandle.Target*/);
+                    GCHandle gcHandle = GCHandle.Alloc(value, type);
+                    IntPtr handle = GCHandle.ToIntPtr(gcHandle);
+                    return handle;
                 }
-
-                return handle;
                 /*IntPtr handle = NewHandle(type);
                 switch (type)
                 {
@@ -541,8 +536,14 @@ namespace FlaxEngine.Interop
 
             internal static object GetObject(IntPtr handle)
             {
-                //if (weakPool.Contains(handle) || weakPoolOther.Contains(handle))
-                //    handle = handle;
+                if (IsWeakHandle(handle))
+                    lock (weakPoolLock)
+                    {
+                        if (weakPool.TryGetValue(handle, out object weakValue) || weakPoolOther.TryGetValue(handle, out weakValue))
+                        {
+                            return weakValue;
+                        }
+                    }
 
                 GCHandle gcHandle = GCHandle.FromIntPtr(handle);
 
@@ -584,6 +585,21 @@ namespace FlaxEngine.Interop
 
             internal static void SetObject(IntPtr handle, object value)
             {
+                if (IsWeakHandle(handle))
+                    lock (weakPoolLock)
+                    {
+                        if (weakPool.ContainsKey(handle))
+                        {
+                            weakPool[handle] = value;
+                            return;
+                        }
+                        else if (weakPoolOther.ContainsKey(handle))
+                        {
+                            weakPoolOther[handle] = value;
+                            return;
+                        }
+                    }
+
                 GCHandle gcHandle = GCHandle.FromIntPtr(handle);
 
                 if (!gcHandle.IsAllocated)
@@ -644,8 +660,12 @@ namespace FlaxEngine.Interop
 
             internal static void FreeHandle(IntPtr handle)
             {
-                if (weakPool.Contains(handle) || weakPoolOther.Contains(handle))
-                    return;
+                if (IsWeakHandle(handle))
+                    lock (weakPoolLock)
+                    {
+                        if (weakPool.ContainsKey(handle) || weakPoolOther.ContainsKey(handle))
+                            return;
+                    }
 
                 GCHandle gcHandle = GCHandle.FromIntPtr(handle);