Properly unload script assembly when reloading

GlitchyEngine/src/Scripting/CoreClrHelper.bf

@@ -29,7 +29,9 @@ static class CoreClrHelper
 
 	//static function void(char8* path) LoadScriptAssembly;
 	static LoadScriptAssemblyFunc _loadScriptAssembly;
-	static function void() _unloadAssemblies;
+	static function void() _unloadScriptAssembly;
+	
+	static function void() _waitUntilOldAssemblyDead;
 
 	// public static unsafe void GetScriptClasses(void** outBuffer, long* length)
 	private function void GetScriptClassesFunc(void** outBuffer, int64* length);
@@ -108,8 +110,10 @@ static class CoreClrHelper
 	{
 		GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "LoadScriptAssembly", out _loadScriptAssembly);
 
-		GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "UnloadAssemblies", out _unloadAssemblies);
+		GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "UnloadScriptAssembly", out _unloadScriptAssembly);
 		
+		GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "WaitUntilOldAssemblyDead", out _waitUntilOldAssemblyDead);
+
 		GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "GetScriptClasses", out _getScriptClasses);
 
 		GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "FreeScriptClassNames", out _freeScriptClassNames);
@@ -130,16 +134,21 @@ static class CoreClrHelper
 		GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "InvokeEntityOnCollisionEnter2D", out _entityScriptFunctions.OnCollisionEnter2D);
 	}
 
-	public static void LoadAppAssembly(Span<uint8> appAssemblyData, Span<uint8> pdbData)
+	public static void LoadScriptAssembly(Span<uint8> appAssemblyData, Span<uint8> pdbData)
 	{
 		_loadScriptAssembly(appAssemblyData.Ptr, appAssemblyData.Length, pdbData.Ptr, pdbData.Length);
 	}
 
-	public static void UnloadAssemblies()
+	public static void UnloadScriptAssembly()
 	{
-		_unloadAssemblies();
+		_unloadScriptAssembly();
 	}	
 
+	public static void WaitUntilOldAssemblyDead()
+	{
+		_waitUntilOldAssemblyDead();
+	}
+
 	/// Gets an array of script class infos. Use FreeScriptClassNames to release the buffer.
 	public static void GetScriptClasses(out void* outBuffer, out int64 entryCount)
 	{

GlitchyEngine/src/Scripting/EntitySerializerWrapper.bf

@@ -16,6 +16,7 @@ class EntitySerializerWrapper : NewScriptClass
 	function void DeserializeMethod(UUID entityId, void* serializedObjPtr, void* serializerPtr);
 	function void SerializeStaticMethod(char8* fullTypeName, void* serializedObjPtr, void* serializerPtr);
 	function void DeserializeStaticMethod(char8* fullTypeName, void* serializedObjPtr, void* serializerPtr);
+	function void FinishSerializationMethod(void* serializerPtr);
 
 	private CreateSerializationContextMethod _createSerializationContextMethod;
 	//private ClearSerializationContextMethod _clearSerializationContextMethod;
@@ -25,16 +26,19 @@ class EntitySerializerWrapper : NewScriptClass
 	private DeserializeMethod _deserializeMethod;
 	private SerializeStaticMethod _serializeStaticMethod;
 	private DeserializeStaticMethod _deserializeStaticMethod;
+	private FinishSerializationMethod _finishSerializationMethod;
 
 	[AllowAppend]
 	public this() : base("GlitchyEngine.Serialization.EntitySerializer", .Empty, .None, false)
 	{
 		CoreClrHelper.GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "CreateSerializationContext", out _createSerializationContextMethod);
 		CoreClrHelper.GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "DestroySerializationContext", out _destroySerializationContextMethod);
+
 		CoreClrHelper.GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "EntitySerializer_Serialize", out _serializeMethod);
 		CoreClrHelper.GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "EntitySerializer_Deserialize", out _deserializeMethod);
 		CoreClrHelper.GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "EntitySerializer_SerializeStaticFields", out _serializeStaticMethod);
 		CoreClrHelper.GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "EntitySerializer_DeserializeStaticFields", out _deserializeStaticMethod);
+		CoreClrHelper.GetFunctionPointerUnmanagedCallersOnly("GlitchyEngine.ScriptGlue, ScriptCore", "EntitySerializer_FinishSerialization", out _finishSerializationMethod);
 	}
 
 	public void CreateSerializationContext(ScriptInstanceSerializer serializer)
@@ -93,4 +97,10 @@ class EntitySerializerWrapper : NewScriptClass
 		
 		_deserializeStaticMethod(@class.FullName.CStr(), serializedObjectPtr, serializerPtr);
 	}
+	
+	public void FinishSerialization(ScriptInstanceSerializer serializer)
+	{
+		void* serializerPtr = Internal.UnsafeCastToPtr(serializer);
+		_finishSerializationMethod(serializerPtr);
+	}
 }

GlitchyEngine/src/Scripting/ScriptEngine.bf

@@ -209,32 +209,6 @@ static class ScriptEngine
 		CoreClrHelper.Init(ScriptCorePath);
 	}	
 
-	/*static void InitMono()
-	{
-		Mono.mono_set_assemblies_path("mono/lib/4.5");
-
-		if (_debuggingEnabled)
-		{
-			char8*[2] options = .(
-				  "--debugger-agent=transport=dt_socket,address=127.0.0.1:2550,server=y,suspend=n,loglevel=3,logfile=MonoDebugger.log",
-				  "--soft-breakpoints"
-				);
-
-			Mono.mono_jit_parse_options(options.Count, &options);
-			Mono.mono_debug_init(.Mono);
-		}
-
-		s_RootDomain = Mono.mono_jit_init("GlitchyEngineJITRuntime");
-		Log.EngineLogger.Assert(s_RootDomain != null, "Failed to initialize mono root domain");
-
-		if (_debuggingEnabled)
-		{
-			Mono.mono_debug_domain_create(s_RootDomain);
-		}
-
-		Mono.mono_thread_set_main(Mono.mono_thread_current());
-	}*/
-	
 	static bool _requestingReload = false;
 
 	static void InitAssemblyWatcher()
@@ -301,7 +275,7 @@ static class ScriptEngine
 
 			File.ReadAll(pdbPath, pdbData);
 
-			CoreClrHelper.LoadAppAssembly(data, pdbData);
+			CoreClrHelper.LoadScriptAssembly(data, pdbData);
 			
 			GetEntitiesFromAssemblies();
 			
@@ -445,14 +419,15 @@ static class ScriptEngine
 	{
 		Debug.Profiler.ProfileFunction!();
 
-		Log.EngineLogger.Info("Reloading script assemblies.");
+		// TODO: Assembly reload might actually take a second. We should consider doing it in a separate thread in the editor.
 
+		Log.EngineLogger.Info("Reloading script assemblies.");
 
 		ScriptInstanceSerializer contextSerializer = scope .();
 
 		contextSerializer.SerializeScriptInstances();
 
-		// TODO: Unload script assemblies, remove class handles, fire unload events?
+		UnloadScriptAssemblies();
 
 		LoadScriptAssemblies();
 
@@ -479,15 +454,28 @@ static class ScriptEngine
 		Log.EngineLogger.Info("Script assemblies reloaded!");
 	}
 
+	private static void UnloadScriptAssemblies()
+	{
+		for (let (id, scriptInstance) in _entityScriptInstances)
+		{
+			if (Context.GetEntityByID(id) case .Ok(let entity))
+			{
+				ScriptComponent* script = entity.GetComponent<ScriptComponent>();
+				DestroyInstance(entity, script);
+			}
+			else
+			{
+				Log.EngineLogger.AssertDebug(false, "Entity script instance exists but no entity?!");
+			}
+		}
+
+		CoreClrHelper.UnloadScriptAssembly();
+		CoreClrHelper.WaitUntilOldAssemblyDead();
+	}
+
 	public static void Shutdown()
 	{
-		/*Mono.mono_domain_set(s_RootDomain, false);
-
-		Mono.mono_domain_unload(s_AppDomain);
-		s_AppDomain = null;
-
-		Mono.mono_jit_cleanup(s_RootDomain);
-		s_RootDomain = null;*/
+		UnloadScriptAssemblies();
 	}
 
 	/// Returns the script class with the given name, or null, if no such class exists.

GlitchyEngine/src/Scripting/ScriptInstanceSerializer.bf

@@ -31,6 +31,15 @@ public class ScriptInstanceSerializer
 		ScriptEngine.Classes.EntitySerializer.CreateSerializationContext(this);
 	}
 
+	/// Signals, that no more objects will be serialized into this context.
+	/// This allows the managed serialization context to do cleanup:
+	/// i. e. for C# we can release references to types which allows us to cleanup the old script assembly.
+	public void FinishSerialization()
+	{
+		ScriptEngine.Classes.EntitySerializer.FinishSerialization(this);
+	}
+
+
 	/// Serializes all script instances that are currently managed by the ScriptEngine.
 	public void SerializeScriptInstances()
 	{
@@ -47,6 +56,8 @@ public class ScriptInstanceSerializer
 		{
 			SerializeStaticScriptClassFields(scriptClass);
 		}
+
+		FinishSerialization();
 	}
 
 	/// Serializes the given script instance.
@@ -79,6 +90,8 @@ public class ScriptInstanceSerializer
 		{
 			DeserializeStaticScriptClassFields(scriptClass);
 		}
+
+		FinishSerialization();
 	}
 
 	/// Deserializes the data into the given script instance, if there is data available.

ScriptCore/ScriptGlue.cs

@@ -8,6 +8,7 @@ using ImGuiNET;
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Diagnostics.Metrics;
 using System.IO;
 using System.Reflection;
 using System.Runtime.CompilerServices;
@@ -15,6 +16,7 @@ using System.Runtime.InteropServices;
 using System.Runtime.Loader;
 using System.Text;
 using System.Threading;
+using System.Xml.Linq;
 
 namespace GlitchyEngine;
 
@@ -112,25 +114,53 @@ internal static unsafe partial class ScriptGlue
         }
     }
 
+    /// <summary>
+    /// Counter for the number of script contexts that have been created. This allows us to debug and track the number of script contexts that have been created and destroyed.
+    /// </summary>
+    private static int _contextCounter = 0;
+
     private static void LoadAssembly(Stream assemblyStream, Stream? pdbStream)
     {
-        _scriptAssemblyContext ??= new AssemblyLoadContext("ScriptContext", true);
+        Debug.Assert(_scriptAssemblyContext == null);
+        Debug.Assert(_appAssembly == null);
+
+        // A new AssemblyLoadContext for the script assembly allows us to unload the assembly later.
+        _scriptAssemblyContext = new AssemblyLoadContext($"ScriptContext_{Interlocked.Increment(ref _contextCounter)}", true);
 
         _appAssembly = _scriptAssemblyContext.LoadFromStream(assemblyStream, pdbStream);
 
         Debug.Assert(_appAssembly != null);
     }
-    
+
+    private static WeakReference? _weakAssemblyRef;
+    private static string? _lastAssemblyName;
+
     /// <summary>
-    /// Called by the engine to unload the assembly (containing user scripts).
+    /// Called by the engine to unload the assembly containing user scripts.
     /// </summary>
     [UnmanagedCallersOnly]
-    public static void UnloadAssemblies()
+    public static void UnloadScriptAssembly()
     {
         try
         {
-            _scriptAssemblyContext?.Unload();
+            Debug.Assert(EntityScriptInstances.Count == 0, "Can not unload assemblies while entities are still active!");
+
+            // TODO: Should we allow calling unload when we never loaded?
+            if (_scriptAssemblyContext == null)
+                return;
+
+            _weakAssemblyRef = new WeakReference(_scriptAssemblyContext);
+            _lastAssemblyName = _scriptAssemblyContext.Name;
+            Debug.WriteLine($"--- Killing {_lastAssemblyName}...");
+
+            // TODO: we should probably have a way to collect and terminate threads. (like a EngineThread override or something)
+
+            _scriptAssemblyContext.Unload();
             _scriptAssemblyContext = null;
+            _appAssembly = null;
+
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
         }
         catch (Exception e)
         {
@@ -138,6 +168,42 @@ internal static unsafe partial class ScriptGlue
         }
     }
 
+    /// <summary>
+    /// Waits until the old assembly is dead.
+    /// </summary>
+    [UnmanagedCallersOnly]
+    public static void WaitUntilOldAssemblyDead()
+    {
+        try
+        {
+            if (_weakAssemblyRef == null)
+                return;
+
+            int iterations = 0;
+
+            while (_weakAssemblyRef.IsAlive)
+            {
+                GC.Collect();
+                GC.WaitForPendingFinalizers();
+
+                // We must sleep a bit, otherwise we absolutely wreck the debugger.
+                Thread.Sleep(10);
+                iterations++;
+            }
+            Debug.WriteLine($"--- {_lastAssemblyName} is now dead after {iterations} garbage collections!");
+
+            _weakAssemblyRef = null;
+            _lastAssemblyName = null;
+        }
+        catch (Exception e)
+        {
+            Log.Exception(e);
+        }
+    }
+    
+    /// <summary>
+    /// Waits until the old assembly is dead.
+    /// </summary>
     private static NativeScriptClassInfo[]? _unsafeClasses;
 
     [StructLayout(LayoutKind.Sequential, Pack = 0)]
@@ -563,6 +629,19 @@ internal static unsafe partial class ScriptGlue
         }
     }
 
+    [UnmanagedCallersOnly]
+    public static void EntitySerializer_FinishSerialization(IntPtr engineSerializer)
+    {
+        try
+        {
+            EntitySerializer.FinishSerialization(engineSerializer);
+        }
+        catch (Exception e)
+        {
+            Log.Exception(e);
+        }
+    }
+
     #endregion EntitySerializer
 
     #endregion

ScriptCore/Serialization/EntitySerializer.cs

@@ -33,10 +33,21 @@ internal static class EntitySerializer
 
     internal static void DestroySerializationContext(IntPtr engineSerializer)
     {
-        SerializationObjects.Remove(engineSerializer);
-        DeserializationObjects.Remove(engineSerializer);
+        if (SerializationObjects.Remove(engineSerializer, out var serializedObjects))
+        {
+            serializedObjects.Clear();
+        }
+
+        if (DeserializationObjects.Remove(engineSerializer, out var deserializationObjects))
+        {
+            deserializationObjects.Clear();
+        }
     }
-    
+    internal static void FinishSerialization(IntPtr engineSerializer)
+    {
+        ClearSerializationContext(engineSerializer);
+    }
+
     internal static void Serialize(Entity entity, IntPtr engineObject, IntPtr engineSerializer)
     {
         SerializedObject obj = new SerializedObject(engineObject, false, UUID.Zero, SerializationObjects[engineSerializer]);

ScriptCore/Serialization/SerializedObject.cs

@@ -107,6 +107,7 @@ public class SerializedObject
 
     public void Serialize(Entity entity)
     {
+        _serializedClasses.Add(entity, this);
         SerializeInstanceFields(entity);
     }