Fix scripting AssemblyLoadContext not getting unloaded

Source/Engine/Engine/NativeInterop.Managed.cs

@@ -176,6 +176,7 @@ namespace FlaxEngine.Interop
                 _managedHandle.Free();
                 _unmanagedData = IntPtr.Zero;
             }
+            _arrayType = _elementType = null;
             ManagedArrayPool.Put(this);
         }
 
@@ -442,22 +443,25 @@ namespace FlaxEngine.Interop
             /// <summary>
             /// Tries to free all references to old weak handles so GC can collect them.
             /// </summary>
-            internal static void TryCollectWeakHandles()
+            internal static void TryCollectWeakHandles(bool force = false)
             {
-                if (weakHandleAccumulator < nextWeakPoolCollection)
+                if (!force)
-                    return;
+                {
+                    if (weakHandleAccumulator < nextWeakPoolCollection)
+                        return;
 
-                nextWeakPoolCollection = weakHandleAccumulator + 1000;
+                    nextWeakPoolCollection = weakHandleAccumulator + 1000;
 
-                // Try to swap pools after garbage collection or whenever the pool gets too large
+                    // Try to swap pools after garbage collection or whenever the pool gets too large
-                var gc0CollectionCount = GC.CollectionCount(0);
+                    var gc0CollectionCount = GC.CollectionCount(0);
-                if (gc0CollectionCount < nextWeakPoolGCCollection && weakPool.Count < WeakPoolCollectionSizeThreshold)
+                    if (gc0CollectionCount < nextWeakPoolGCCollection && weakPool.Count < WeakPoolCollectionSizeThreshold)
-                    return;
+                        return;
-                nextWeakPoolGCCollection = gc0CollectionCount + 1;
+                    nextWeakPoolGCCollection = gc0CollectionCount + 1;
 
-                // Prevent huge allocations from swapping the pools in the middle of the operation
+                    // Prevent huge allocations from swapping the pools in the middle of the operation
-                if (System.Diagnostics.Stopwatch.GetElapsedTime(lastWeakPoolCollectionTime).TotalMilliseconds < WeakPoolCollectionTimeThreshold)
+                    if (System.Diagnostics.Stopwatch.GetElapsedTime(lastWeakPoolCollectionTime).TotalMilliseconds < WeakPoolCollectionTimeThreshold)
-                    return;
+                        return;
+                }
                 lastWeakPoolCollectionTime = System.Diagnostics.Stopwatch.GetTimestamp();
 
                 // Swap the pools and release the oldest pool for GC

Source/Engine/Engine/NativeInterop.Unmanaged.cs

@@ -1054,7 +1054,49 @@ namespace FlaxEngine.Interop
         }
 
         [UnmanagedCallersOnly]
-        internal static void ReloadScriptingAssemblyLoadContext()
+        internal static void CreateScriptingAssemblyLoadContext()
+        {
+#if FLAX_EDITOR
+            if (scriptingAssemblyLoadContext != null)
+            {
+                // Wait for previous ALC to finish unloading, track it without holding strong references to it
+                GCHandle weakRef = GCHandle.Alloc(scriptingAssemblyLoadContext, GCHandleType.WeakTrackResurrection);
+                scriptingAssemblyLoadContext = null;
+#if true
+                // In case the ALC doesn't unload properly: https://learn.microsoft.com/en-us/dotnet/standard/assembly/unloadability#debug-unloading-issues
+                while (true)
+#else
+                for (int attempts = 5; attempts > 0; attempts--)
+#endif
+                {
+                    GC.Collect();
+                    GC.WaitForPendingFinalizers();
+
+                    if (!IsHandleAlive(weakRef))
+                        break;
+                    System.Threading.Thread.Sleep(1);
+                }
+                if (IsHandleAlive(weakRef))
+                    Debug.LogWarning("Scripting AssemblyLoadContext was not unloaded.");
+                weakRef.Free();
+
+                static bool IsHandleAlive(GCHandle weakRef)
+                {
+                    // Checking the target in scope somehow holds a reference to it...?
+                    return weakRef.Target != null;
+                }
+            }
+
+            scriptingAssemblyLoadContext = new AssemblyLoadContext("Flax", isCollectible: true);
+            scriptingAssemblyLoadContext.Resolving += OnScriptingAssemblyLoadContextResolving;
+#else
+            scriptingAssemblyLoadContext = new AssemblyLoadContext("Flax", isCollectible: false);
+#endif
+            DelegateHelpers.InitMethods();
+        }
+
+        [UnmanagedCallersOnly]
+        internal static void UnloadScriptingAssemblyLoadContext()
         {
 #if FLAX_EDITOR
             // Clear all caches which might hold references to assemblies in collectible ALC
@@ -1072,24 +1114,73 @@ namespace FlaxEngine.Interop
                 handle.Free();
             propertyHandleCacheCollectible.Clear();
 
+            foreach (var key in assemblyHandles.Keys.Where(x => x.IsCollectible))
+                assemblyHandles.Remove(key);
+            foreach (var key in assemblyOwnedNativeLibraries.Keys.Where(x => x.IsCollectible))
+                assemblyOwnedNativeLibraries.Remove(key);
+
             _typeSizeCache.Clear();
 
             foreach (var pair in classAttributesCacheCollectible)
                 pair.Value.Free();
             classAttributesCacheCollectible.Clear();
 
+            ArrayFactory.marshalledTypes.Clear();
+            ArrayFactory.arrayTypes.Clear();
+            ArrayFactory.createArrayDelegates.Clear();
+
             FlaxEngine.Json.JsonSerializer.ResetCache();
+            DelegateHelpers.Release();
+
+            // Ensure both pools are empty
+            ManagedHandle.ManagedHandlePool.TryCollectWeakHandles(true);
+            ManagedHandle.ManagedHandlePool.TryCollectWeakHandles(true);
+
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+
+            {
+                // HACK: Workaround for TypeDescriptor holding references to collectible types (https://github.com/dotnet/runtime/issues/30656)
+
+                Type TypeDescriptionProviderType = typeof(System.ComponentModel.TypeDescriptionProvider);
+                MethodInfo clearCacheMethod = TypeDescriptionProviderType?.Assembly.GetType("System.ComponentModel.ReflectionCachesUpdateHandler")?.GetMethod("ClearCache");
+                if (clearCacheMethod != null)
+                    clearCacheMethod.Invoke(null, new object[] { null });
+                else
+                {
+                    MethodInfo beforeUpdateMethod = TypeDescriptionProviderType?.Assembly.GetType("System.ComponentModel.ReflectionCachesUpdateHandler")?.GetMethod("BeforeUpdate");
+                    if (beforeUpdateMethod != null)
+                        beforeUpdateMethod.Invoke(null, new object[] { null });
+                }
+
+                Type TypeDescriptorType = typeof(System.ComponentModel.TypeDescriptor);
+
+                object s_internalSyncObject = TypeDescriptorType?.GetField("s_internalSyncObject", System.Reflection.BindingFlags.Static | System.Reflection.BindingFlags.NonPublic)?.GetValue(null);
+                System.Collections.Hashtable s_defaultProviders = (System.Collections.Hashtable)TypeDescriptorType?.GetField("s_defaultProviders", System.Reflection.BindingFlags.Static | System.Reflection.BindingFlags.NonPublic)?.GetValue(null);
+                if (s_internalSyncObject != null && s_defaultProviders != null)
+                {
+                    lock (s_internalSyncObject)
+                        s_defaultProviders.Clear();
+                }
+
+                object s_providerTable = TypeDescriptorType?.GetField("s_providerTable", System.Reflection.BindingFlags.Static | System.Reflection.BindingFlags.NonPublic)?.GetValue(null);
+                System.Collections.Hashtable s_providerTypeTable = (System.Collections.Hashtable)TypeDescriptorType?.GetField("s_providerTypeTable", System.Reflection.BindingFlags.Static | System.Reflection.BindingFlags.NonPublic)?.GetValue(null);
+                if (s_providerTable != null && s_providerTypeTable != null)
+                {
+                    lock (s_providerTable)
+                        s_providerTypeTable.Clear();
+                    TypeDescriptorType.GetField("s_providerTable", System.Reflection.BindingFlags.Static | System.Reflection.BindingFlags.NonPublic)
+                        ?.FieldType.GetMethods(BindingFlags.Instance | BindingFlags.Public).FirstOrDefault(x => x.Name == "Clear")
+                        ?.Invoke(s_providerTable, new object[] { });
+                }
+            }
 
             // Unload the ALC
-            bool unloading = true;
-            scriptingAssemblyLoadContext.Unloading += (alc) => { unloading = false; };
             scriptingAssemblyLoadContext.Unload();
+            scriptingAssemblyLoadContext.Resolving -= OnScriptingAssemblyLoadContextResolving;
 
-            while (unloading)
+            GC.Collect();
-                System.Threading.Thread.Sleep(1);
+            GC.WaitForPendingFinalizers();
-
-            InitScriptingAssemblyLoadContext();
-            DelegateHelpers.InitMethods();
 #endif
         }
 

Source/Engine/Engine/NativeInterop.cs

@@ -73,19 +73,6 @@ namespace FlaxEngine.Interop
             return nativeLibrary;
         }
 
-        private static void InitScriptingAssemblyLoadContext()
-        {
-#if FLAX_EDITOR
-            var isCollectible = true;
-#else
-            var isCollectible = false;
-#endif
-            scriptingAssemblyLoadContext = new AssemblyLoadContext("Flax", isCollectible);
-#if FLAX_EDITOR
-            scriptingAssemblyLoadContext.Resolving += OnScriptingAssemblyLoadContextResolving;
-#endif
-        }
-
         [UnmanagedCallersOnly]
         internal static unsafe void Init()
         {
@@ -97,8 +84,6 @@ namespace FlaxEngine.Interop
             System.Threading.Thread.CurrentThread.CurrentCulture = CultureInfo.InvariantCulture;
             System.Threading.Thread.CurrentThread.CurrentUICulture = CultureInfo.InvariantCulture;
 
-            InitScriptingAssemblyLoadContext();
-            DelegateHelpers.InitMethods();
         }
 
 #if FLAX_EDITOR
@@ -1475,11 +1460,11 @@ namespace FlaxEngine.Interop
 
         internal static class ArrayFactory
         {
-            private delegate Array CreateArrayDelegate(long size);
+            internal delegate Array CreateArrayDelegate(long size);
 
-            private static ConcurrentDictionary<Type, Type> marshalledTypes = new ConcurrentDictionary<Type, Type>(1, 3);
+            internal static ConcurrentDictionary<Type, Type> marshalledTypes = new ConcurrentDictionary<Type, Type>(1, 3);
-            private static ConcurrentDictionary<Type, Type> arrayTypes = new ConcurrentDictionary<Type, Type>(1, 3);
+            internal static ConcurrentDictionary<Type, Type> arrayTypes = new ConcurrentDictionary<Type, Type>(1, 3);
-            private static ConcurrentDictionary<Type, CreateArrayDelegate> createArrayDelegates = new ConcurrentDictionary<Type, CreateArrayDelegate>(1, 3);
+            internal static ConcurrentDictionary<Type, CreateArrayDelegate> createArrayDelegates = new ConcurrentDictionary<Type, CreateArrayDelegate>(1, 3);
 
             internal static Type GetMarshalledType(Type elementType)
             {
@@ -1645,17 +1630,6 @@ namespace FlaxEngine.Interop
             return RegisterType(type, true).typeHolder;
         }
 
-        internal static (TypeHolder typeHolder, ManagedHandle handle) GetTypeHolderAndManagedHandle(Type type)
-        {
-            if (managedTypes.TryGetValue(type, out (TypeHolder typeHolder, ManagedHandle handle) tuple))
-                return tuple;
-#if FLAX_EDITOR
-            if (managedTypesCollectible.TryGetValue(type, out tuple))
-                return tuple;
-#endif
-            return RegisterType(type, true);
-        }
-
         /// <summary>
         /// Returns a static ManagedHandle to TypeHolder for given Type, and caches it if needed.
         /// </summary>
@@ -1781,6 +1755,14 @@ namespace FlaxEngine.Interop
 #endif
             }
 
+            internal static void Release()
+            {
+                MakeNewCustomDelegateFunc = null;
+#if FLAX_EDITOR
+                MakeNewCustomDelegateFuncCollectible = null;
+#endif
+            }
+
             internal static Type MakeNewCustomDelegate(Type[] parameters)
             {
 #if FLAX_EDITOR

Source/Engine/Scripting/ManagedCLR/MCore.h

@@ -45,9 +45,16 @@ public:
     /// </summary>
     static void UnloadEngine();
 
+    /// <summary>
+    /// Creates the assembly load context for assemblies used by Scripting.
+    /// </summary>
+    static void CreateScriptingAssemblyLoadContext();
+
 #if USE_EDITOR
-    // Called by Scripting in a middle of hot-reload (after unloading modules but before loading them again).
+    /// <summary>
-    static void ReloadScriptingAssemblyLoadContext();
+    /// Called by Scripting in a middle of hot-reload (after unloading modules but before loading them again).
+    /// </summary>
+    static void UnloadScriptingAssemblyLoadContext();
 #endif
 
 public:

Source/Engine/Scripting/Runtime/DotNet.cpp

@@ -330,9 +330,15 @@ void MCore::UnloadEngine()
     ShutdownHostfxr();
 }
 
+void MCore::CreateScriptingAssemblyLoadContext()
+{
+    static void* CreateScriptingAssemblyLoadContextPtr = GetStaticMethodPointer(TEXT("CreateScriptingAssemblyLoadContext"));
+    CallStaticMethod<void>(CreateScriptingAssemblyLoadContextPtr);
+}
+
 #if USE_EDITOR
 
-void MCore::ReloadScriptingAssemblyLoadContext()
+void MCore::UnloadScriptingAssemblyLoadContext()
 {
     // Clear any cached class attributes (see https://github.com/FlaxEngine/FlaxEngine/issues/1108)
     for (auto e : CachedClassHandles)
@@ -377,8 +383,8 @@ void MCore::ReloadScriptingAssemblyLoadContext()
         }
     }
 
-    static void* ReloadScriptingAssemblyLoadContextPtr = GetStaticMethodPointer(TEXT("ReloadScriptingAssemblyLoadContext"));
+    static void* UnloadScriptingAssemblyLoadContextPtr = GetStaticMethodPointer(TEXT("UnloadScriptingAssemblyLoadContext"));
-    CallStaticMethod<void>(ReloadScriptingAssemblyLoadContextPtr);
+    CallStaticMethod<void>(UnloadScriptingAssemblyLoadContextPtr);
 }
 
 #endif

Source/Engine/Scripting/Runtime/Mono.cpp

@@ -715,9 +715,13 @@ void MCore::UnloadEngine()
 #endif
 }
 
+void MCore::CreateScriptingAssemblyLoadContext()
+{
+}
+
 #if USE_EDITOR
 
-void MCore::ReloadScriptingAssemblyLoadContext()
+void MCore::UnloadScriptingAssemblyLoadContext()
 {
 }
 

Source/Engine/Scripting/Runtime/None.cpp

@@ -58,9 +58,13 @@ void MCore::UnloadEngine()
     MRootDomain = nullptr;
 }
 
+void MCore::CreateScriptingAssemblyLoadContext()
+{
+}
+
 #if USE_EDITOR
 
-void MCore::ReloadScriptingAssemblyLoadContext()
+void MCore::UnloadScriptingAssemblyLoadContext()
 {
 }
 

Source/Engine/Scripting/Scripting.cpp

@@ -182,6 +182,8 @@ bool ScriptingService::Init()
         return true;
     }
 
+    MCore::CreateScriptingAssemblyLoadContext();
+
     // Cache root domain
     _rootDomain = MCore::GetRootDomain();
 
@@ -710,7 +712,8 @@ void Scripting::Reload(bool canTriggerSceneReload)
     _hasGameModulesLoaded = false;
 
     // Release and create a new assembly load context for user assemblies
-    MCore::ReloadScriptingAssemblyLoadContext();
+    MCore::UnloadScriptingAssemblyLoadContext();
+    MCore::CreateScriptingAssemblyLoadContext();
 
     // Give GC a try to cleanup old user objects and the other mess
     MCore::GC::Collect();