Nerfed/Nerfed.Editor/Project/EditorAssemblyLoader.cs

using Nerfed.Runtime;
using System.Diagnostics;
using System.Reflection;
using System.Runtime.CompilerServices;

namespace Nerfed.Editor.Project;

// https://github.com/godotengine/godot/blob/master/modules/mono/glue/GodotSharp/GodotPlugins/Main.c
// https://gitlab.com/robertk92/assemblyreloadtest/-/blob/main/AppContextTest/Program.cs

internal static class EditorAssemblyLoader
{
    private sealed class EditorAssemblyLoadContextWrapper
    {
        private EditorAssemblyLoadContext assemblyLoadContext;
        private readonly WeakReference weakReference;

        private EditorAssemblyLoadContextWrapper(EditorAssemblyLoadContext assemblyLoadContext, WeakReference weakReference)
        {
            this.assemblyLoadContext = assemblyLoadContext;
            this.weakReference = weakReference;
        }

        public bool IsCollectible
        {
            [MethodImpl(MethodImplOptions.NoInlining)]
            // If assemblyLoadContext is null we already started unloading, so it was collectible.
            get => assemblyLoadContext?.IsCollectible ?? true;
        }

        public bool IsAlive
        {
            [MethodImpl(MethodImplOptions.NoInlining)]
            get => weakReference.IsAlive;
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        public static (Assembly, EditorAssemblyLoadContextWrapper) CreateAndLoad(AssemblyName assemblyName)
        {
            EditorAssemblyLoadContext context = new EditorAssemblyLoadContext();
            WeakReference reference = new WeakReference(context, trackResurrection: true);
            EditorAssemblyLoadContextWrapper wrapper = new EditorAssemblyLoadContextWrapper(context, reference);
            Assembly assembly = context.LoadFromAssemblyName(assemblyName);
            return (assembly, wrapper);
        }

        [MethodImpl(MethodImplOptions.NoInlining)]
        internal void Unload()
        {
            assemblyLoadContext?.Unload();
            assemblyLoadContext = null;
        }
    }

    private static (Assembly, EditorAssemblyLoadContextWrapper) Load(string assemblyFilePath)
    {
        string assemblyFileName = Path.GetFileNameWithoutExtension(assemblyFilePath);

        AssemblyName assemblyName = new AssemblyName(assemblyFileName);

        return EditorAssemblyLoadContextWrapper.CreateAndLoad(assemblyName);
    }

    private static bool Unload(EditorAssemblyLoadContextWrapper assemblyLoadContextWrapper)
    {
        if (assemblyLoadContextWrapper == null)
        {
            return true;
        }

        if (!assemblyLoadContextWrapper.IsCollectible) 
        {
            Log.Error($"{assemblyLoadContextWrapper} is not collectable!");
            return false;
        }

        assemblyLoadContextWrapper.Unload();

        GC.Collect();
        GC.WaitForPendingFinalizers();

        TimeSpan timeout = TimeSpan.FromSeconds(30);
        Stopwatch stopwatch = Stopwatch.StartNew();

        while (assemblyLoadContextWrapper.IsAlive) 
        {
            GC.Collect(GC.MaxGeneration, GCCollectionMode.Forced);
            GC.WaitForPendingFinalizers();

            if (!assemblyLoadContextWrapper.IsAlive)
            {
                break;
            }

            if (stopwatch.Elapsed.TotalSeconds % 10 == 0)
            {
                Log.Info("Tring to unload assembly...");
            }

            if (stopwatch.Elapsed >= timeout)
            {
                Log.Error("Failed to unload assembly!");
                return false;
            }
        }

        return true;
    }
}