Fix assembly load priorities when using AssemblyLoadContext

Mono.TextTemplating/Mono.TextTemplating/CurrentDomainAssemblyResolver.cs

@@ -26,17 +26,20 @@ Assembly ResolveReferencedAssemblies (object sender, ResolveEventArgs args)
 		{
 			var asmName = new AssemblyName (args.Name);
 
+			// The list of assembly files referenced by the template may contain reference assemblies,
+			// which will fail to load. Letting the host attempt to resolve the assembly first
+			// gives it an opportunity to resolve runtime assemblies.
+			var path = resolveAssemblyReference (asmName.Name + ".dll");
+			if (File.Exists (path)) {
+				return Assembly.LoadFrom (path);
+			}
+
 			foreach (var asmFile in assemblyFiles) {
 				if (asmName.Name == Path.GetFileNameWithoutExtension (asmFile)) {
 					return Assembly.LoadFrom (asmFile);
 				}
 			}
 
-			var path = resolveAssemblyReference (asmName.Name + ".dll");
-			if (File.Exists (path)) {
-				return Assembly.LoadFrom (path);
-			}
-
 			return null;
 		}
 

Mono.TextTemplating/Mono.TextTemplating/TemplateAssemblyLoadContext.cs

@@ -6,6 +6,7 @@
 using System;
 using System.CodeDom.Compiler;
 using System.IO;
+using System.Linq;
 using System.Reflection;
 using System.Runtime.Loader;
 
@@ -15,6 +16,8 @@ namespace Mono.TextTemplating
 {
 	class TemplateAssemblyLoadContext : AssemblyLoadContext
 	{
+		readonly AssemblyLoadContext hostAssemblyLoadContext;
+
 		readonly string[] templateAssemblyFiles;
 		readonly ITextTemplatingEngineHost host;
 		readonly Assembly hostAssembly;
@@ -28,31 +31,79 @@ 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)
+		{
+			// The list of assembly files referenced by the template may contain reference assemblies,
+			// which will fail to load. Letting the host attempt to resolve the assembly first
+			// gives it an opportunity to resolve runtime assemblies.
+			var path = host.ResolveAssemblyReference (assemblyName.Name + ".dll");
+			if (File.Exists (path)) {
+				return LoadFromAssemblyPath (path);
+			}
+
 			for (int i = 0; i < templateAssemblyFiles.Length; i++) {
 				var asmFile = templateAssemblyFiles[i];
 				if (asmFile is null) {
@@ -70,11 +121,6 @@ protected override Assembly Load (AssemblyName assemblyName)
 				}
 			}
 
-			var path = host.ResolveAssemblyReference (assemblyName.Name + ".dll");
-			if (File.Exists (path)) {
-				return LoadFromAssemblyPath (path);
-			}
-
 			return null;
 		}
 	}