Skip to content

Commit

Permalink
Fix AssemblyLoadContext assembly loading behavior
Browse files Browse the repository at this point in the history
The behavior should now be equivalent to that of the AppDomain codepath:
assemblies are loaded from referenced assembly paths only after
failing to resolve them from loaded assemblies in the host context
and after the runtime fails to resolve them from the default context.
  • Loading branch information
mhutch committed Jan 10, 2024
1 parent 6ec6c31 commit 2aea692
Showing 1 changed file with 48 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
using System;
using System.CodeDom.Compiler;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Runtime.Loader;

Expand All @@ -15,6 +16,8 @@ namespace Mono.TextTemplating
{
class TemplateAssemblyLoadContext : AssemblyLoadContext
{
readonly AssemblyLoadContext hostAssemblyLoadContext;

readonly string[] templateAssemblyFiles;
readonly ITextTemplatingEngineHost host;
readonly Assembly hostAssembly;
Expand All @@ -28,31 +31,71 @@ public TemplateAssemblyLoadContext (string[] templateAssemblyFiles, ITextTemplat
: base (isCollectible: true)
#endif
{
this.templateAssemblyFiles = templateAssemblyFiles;
this.host = host;
hostAssembly = host.GetType ().Assembly;

codeDomAsmName = typeof (CompilerErrorCollection).Assembly.GetName ();
textTemplatingAsmName = typeof (TemplateGenerator).Assembly.GetName ();
hostAsmName = hostAssembly.GetName ();

// the parent load context is the context that loaded the host assembly
this.hostAssemblyLoadContext = AssemblyLoadContext.GetLoadContext (hostAssembly);

this.templateAssemblyFiles = templateAssemblyFiles;
this.host = host;

this.Resolving += ResolveAssembly;
}

// Load order is as follows:
//
// First, the Load(AssemblyName) override is called. Our impl of this ensures that the CodeDom and TextTemplating
// and other host assemblies are loaded from the host AssemblyLoadContext, so that we can interchange types.
//
// For assemblies that are not handled by Load(AssemblyName), the runtime next attempts to resolve them
// from AssemblyLoadContext.Default, which may load assemblies into AssemblyLoadContext.Default via
// assembly probing. This is where runtime assemblies wil be loaded.
//
// Finally, if the runtime fails to resolve the assembly, the Resolving event is raised, which we handle
// to resolve assemblies explicitly referenced by the template. The priority of this event is equivalent
// to AppDomain.AssemblyResolve, so using this matches the behavior of the AppDomain codepath.
//
// See https://learn.microsoft.com/en-us/dotnet/core/dependency-loading/loading-managed#algorithm

protected override Assembly Load (AssemblyName assemblyName)
{
// CodeDom and TextTemplating MUST come from the same context as the host as we need to be able to reflect and cast
// this is an issue with MSBuild which loads tasks in their own load contexts
// The CodeDom and TextTemplating assemblies in the template load context MUST be the same as the
// ones used by the host as we need to be able to reflect and cast and interchange types.
// We cannot rely on falling back to resolving them from AssemblyLoadContext.Default, as this fails
// in cases such as running the templating engine in an MSBuild task, as MSBuild loads tasks in
// their own load contexts.
if (assemblyName.Name == codeDomAsmName.Name) {
return typeof (CompilerErrorCollection).Assembly;
}
if (assemblyName.Name == textTemplatingAsmName.Name) {
return typeof (TemplateGenerator).Assembly;
}

// the host may be a custom host, and must also be in the same context
// The host may be a custom host, and must also be in the same context, so that host-specific
// templates can access the host instance.
if (assemblyName.Name == hostAsmName.Name) {
return hostAssembly;
}

// Resolve any more assemblies from the parent context that we can. There may be a custom host,
// and it may expose types from other assemblies that may also need to be interchangeable.
// Technically this loops makes the explicit checks above redundant but it's better
// to be absolutely clear about what we're doing and why.
var fromParent = hostAssemblyLoadContext.Assemblies.FirstOrDefault (a => a.GetName ().Name == assemblyName.Name);
if (fromParent is not null) {
return fromParent;
}

// let the runtime resolve from AssemblyLoadContext.Default
return null;
}

Assembly ResolveAssembly (AssemblyLoadContext context, AssemblyName assemblyName)
{
for (int i = 0; i < templateAssemblyFiles.Length; i++) {
var asmFile = templateAssemblyFiles[i];
if (asmFile is null) {
Expand Down

0 comments on commit 2aea692

Please sign in to comment.