Offload the load context caching to bepinex

BepInEx.Core/Bootstrap/BaseChainloader.cs

@@ -146,7 +146,7 @@ public abstract class BaseChainloader<TPlugin>
     ///     Value holds info of the originating provider. These contexts are used by chainloader to
     ///     load plugins. Plugins that are loaded successfully are added to <see cref="Plugins" />.
     /// </summary>
-    public Dictionary<IPluginLoadContext, PluginInfo> LoadContexts { get; } = new();
+    internal Dictionary<CachedPluginLoadContext, PluginInfo> LoadContexts { get; } = new();
     
     /// <summary>
     ///     Collection of error chainloader messages that occured during plugin loading.
@@ -350,7 +350,7 @@ public abstract class BaseChainloader<TPlugin>
             LoadPlugins(providers, provider: true);
             FinishedProviders?.Invoke();
             
-            var plugins = TypeLoader.GetPluginsFromLoadContexts(LoadContexts.Keys, ToPluginInfo<TPlugin>, HasBepinPluginType<BepInPlugin>, "chainloader");
+            var plugins = TypeLoader.GetPluginsFromLoadContexts(LoadContexts.Keys.Cast<IPluginLoadContext>(), ToPluginInfo<TPlugin>, HasBepinPluginType<BepInPlugin>, "chainloader");
             Logger.Log(LogLevel.Info, $"{plugins.Count} plugin{(plugins.Count == 1 ? "" : "s")} to load");
             LoadPlugins(plugins, provider: false);
             Finished?.Invoke();
@@ -414,23 +414,24 @@ public abstract class BaseChainloader<TPlugin>
                     var pluginLoadContexts = ((BasePluginProvider)pluginInfo.Instance).GetPlugins();
                     foreach (IPluginLoadContext context in pluginLoadContexts)
                     {
-                        LoadContexts[context] = pluginInfo;
+                        var cachedContext = new CachedPluginLoadContext(context);
+                        LoadContexts[cachedContext] = pluginInfo;
                     }
                     ProviderLoaded?.Invoke(pluginInfo);
                 }
                 else
                 {
-                    pluginInfo.Source = LoadContexts[pluginInfo.LoadContext];
+                    pluginInfo.Source = LoadContexts[(CachedPluginLoadContext)pluginInfo.LoadContext];
                     pluginInfo.Instance = LoadPlugin(pluginInfo, ass);
                     PluginLoaded?.Invoke(pluginInfo);
-                    pluginInfo.LoadContext.Dispose();
+                    foreach (CachedPluginLoadContext context in LoadContexts.Keys)
+                        context.Dispose();
                 }
             }
             catch (Exception ex)
             {
                 invalidPlugins.Add(pluginInfo.Metadata.GUID);
                 existingPluginInfos.Remove(pluginInfo.Metadata.GUID);
-                pluginInfo.LoadContext?.Dispose();
 
                 Logger.Log(LogLevel.Error,
                            $"Error loading [{pluginInfo}]: {(ex is ReflectionTypeLoadException re ? TypeLoader.TypeLoadExceptionToString(re) : ex.ToString())}");

BepInEx.Core/Bootstrap/TypeLoader.cs

@@ -10,6 +10,47 @@ using Mono.Cecil;
 
 namespace BepInEx.Bootstrap;
 
+internal class CachedPluginLoadContext : IPluginLoadContext, IDisposable
+{
+    public IPluginLoadContext PluginLoadContext { get; }
+    private byte[] assemblyData;
+    private byte[] assemblySymbolsData;
+    private readonly Dictionary<string, byte[]> files = new();
+
+    public CachedPluginLoadContext(IPluginLoadContext pluginLoadContext)
+    {
+        PluginLoadContext = pluginLoadContext;
+    }
+
+    public string AssemblyIdentifier => PluginLoadContext.AssemblyIdentifier;
+    public string AssemblyHash => PluginLoadContext.AssemblyHash;
+    public byte[] GetAssemblyData()
+    {
+        return assemblyData ??= PluginLoadContext.GetAssemblyData();
+    }
+
+    public byte[] GetAssemblySymbolsData()
+    {
+        return assemblySymbolsData ??= PluginLoadContext.GetAssemblySymbolsData();
+    }
+
+    public byte[] GetFile(string relativePath)
+    {
+        if (!files.ContainsKey(relativePath))
+        {
+            files[relativePath] = PluginLoadContext.GetFile(relativePath);
+        }
+        return files[relativePath];
+    }
+
+    public void Dispose()
+    {
+        assemblyData = null;
+        assemblySymbolsData = null;
+        files.Clear();
+    }
+}
+
 /// <summary>
 ///     A cacheable metadata item. Can be used with <see cref="TypeLoader.LoadAssemblyCache{T}" /> and
 ///     <see cref="TypeLoader.SaveAssemblyCache{T}" /> to cache plugin metadata.

BepInEx.Core/Contract/IPluginLoadContext.cs

@@ -6,7 +6,7 @@ namespace BepInEx;
 ///     An interface that allows to dynamically load and track an assembly
 ///     which can be obtained from a provider
 /// </summary>
-public interface IPluginLoadContext : IDisposable
+public interface IPluginLoadContext
 {
     /// <summary>
     ///     An identifier that uniquely identifies an assembly from a provider no matter its revision

BepInEx.PatcherProvider/BepInExPatcherLoadContext.cs

@@ -7,25 +7,15 @@ public class BepInExPatcherLoadContext : IPluginLoadContext
 {
     public string AssemblyIdentifier { get; internal set; }
     public string AssemblyHash { get; internal set; }
-    private byte[] assemblyData;
-    private byte[] assemblySymbolsData;
     public byte[] GetAssemblyData()
     {
-        if (assemblyData == null)
-        {
-            assemblyData = File.ReadAllBytes(AssemblyIdentifier);
-        }
-
-        return assemblyData;
+        return File.ReadAllBytes(AssemblyIdentifier);
     }
     
     public byte[] GetAssemblySymbolsData()
     {
-        if (assemblySymbolsData == null)
-        {
-            if (!Utility.TryResolveDllSymbols(AssemblyIdentifier, out assemblySymbolsData))
-                assemblySymbolsData = null;
-        }
+        if (Utility.TryResolveAssemblySymbols(AssemblyIdentifier, out var assemblySymbolsData))
+            return assemblySymbolsData;
 
         return assemblySymbolsData;
     }
@@ -39,10 +29,4 @@ public class BepInExPatcherLoadContext : IPluginLoadContext
         string filePath = Path.Combine(assemblyFolder, relativePath);
         return File.ReadAllBytes(filePath);
     }
-
-    public void Dispose()
-    {
-        assemblyData = null;
-        assemblySymbolsData = null;
-    }
 }

BepInEx.PluginProvider/BepInExPluginLoadContext.cs

@@ -7,25 +7,15 @@ public class BepInExPluginLoadContext : IPluginLoadContext
 {
     public string AssemblyIdentifier { get; internal set; }
     public string AssemblyHash { get; internal set; }
-    private byte[] assemblyData;
-    private byte[] assemblySymbolsData;
     public byte[] GetAssemblyData()
     {
-        if (assemblyData == null)
-        {
-            assemblyData = File.ReadAllBytes(AssemblyIdentifier);
-        }
-
-        return assemblyData;
+        return File.ReadAllBytes(AssemblyIdentifier);
     }
-
+    
     public byte[] GetAssemblySymbolsData()
     {
-        if (assemblySymbolsData == null)
-        {
-            if (!Utility.TryResolveDllSymbols(AssemblyIdentifier, out assemblySymbolsData))
-                assemblySymbolsData = null;
-        }
+        if (Utility.TryResolveAssemblySymbols(AssemblyIdentifier, out var assemblySymbolsData))
+            return assemblySymbolsData;
 
         return assemblySymbolsData;
     }
@@ -39,10 +29,4 @@ public class BepInExPluginLoadContext : IPluginLoadContext
         string filePath = Path.Combine(assemblyFolder, relativePath);
         return File.ReadAllBytes(filePath);
     }
-
-    public void Dispose()
-    {
-        assemblyData = null;
-        assemblySymbolsData = null;
-    }
 }

BepInEx.Preloader.Core/Patching/AssemblyPatcher.cs

@@ -25,7 +25,7 @@ public class AssemblyPatcher : IDisposable
 
     private Func<byte[], string, Assembly> assemblyLoader;
 
-    private readonly List<IPluginLoadContext> loadContexts = new();
+    private readonly List<CachedPluginLoadContext> loadContexts = new();
     private readonly List<BasePatcherProvider> providers = new();
     
     public AssemblyPatcher(Func<byte[], string, Assembly> assemblyLoader)
@@ -164,7 +164,7 @@ public class AssemblyPatcher : IDisposable
 
                     var instance = (BasePatcherProvider) Activator.CreateInstance(type);
                     providers.Add(instance);
-                    loadContexts.AddRange(instance.GetPatchers());
+                    loadContexts.AddRange(instance.GetPatchers().Select(x => new CachedPluginLoadContext(x)));
                 }
                 catch (Exception e)
                 {
@@ -174,7 +174,7 @@ public class AssemblyPatcher : IDisposable
             }
         }
 
-        var patchers = TypeLoader.GetPluginsFromLoadContexts(loadContexts, ToPatcherPlugin<BasePatcher>, HasPatcherType<BasePatcher>);
+        var patchers = TypeLoader.GetPluginsFromLoadContexts(loadContexts.Cast<IPluginLoadContext>(), ToPatcherPlugin<BasePatcher>, HasPatcherType<BasePatcher>);
         
         var sortedPatchers = new List<PatchDefinition>();
         
@@ -188,11 +188,15 @@ public class AssemblyPatcher : IDisposable
 
             var assName = ass.GetName();
             Logger.Log(LogLevel.Debug, $"Loaded patcher {patcherPlugin.TypeName} from [{assName.Name} {assName.Version}]");
-            patcherPlugin.LoadContext.Dispose();
         }
 
         PatcherContext.PatchDefinitions.AddRange(sortedPatchers);
         Logger.Log(LogLevel.Info, $"Loaded {sortedPatchers.Count} patcher{(sortedPatchers.Count > 0 ? "s" : "")}");
+
+        foreach (CachedPluginLoadContext loadContext in loadContexts)
+        {
+            loadContext.Dispose();
+        }
     }
 
     private Assembly PatcherProvidersAssemblyResolver(object sender, ResolveEventArgs args)