Skip to content

Commit 66dcd00

Browse files
daxian-dbwandyleejordan
authored andcommitted
Make AssemblyLoadContext a singleton and refactor
1 parent fdfd02c commit 66dcd00

5 files changed

Lines changed: 213 additions & 74 deletions

File tree

src/Microsoft.PowerShell.CoreCLR.AssemblyLoadContext/CoreCLR/CorePsAssemblyLoadContext.cs

Lines changed: 196 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ namespace System.Management.Automation
1919
/// <summary>
2020
/// The powershell custom AssemblyLoadContext implementation
2121
/// </summary>
22-
public partial class PowerShellAssemblyLoadContext : AssemblyLoadContext
22+
internal partial class PowerShellAssemblyLoadContext : AssemblyLoadContext
2323
{
2424
#region Resource_Strings
2525

@@ -30,21 +30,30 @@ public partial class PowerShellAssemblyLoadContext : AssemblyLoadContext
3030
// 2. Load assembly with culture 'en' (Microsoft.PowerShell.CoreCLR.AssemblyLoadContext.resources, Version=3.0.0.0, Culture=en, PublicKeyToken=31bf3856ad364e35)
3131
// When the first attempt fails, we again need to retrieve the resouce string to construct another exception, which ends up with an infinite loop.
3232
private const string BaseFolderDoesNotExist = "The base directory '{0}' does not exist.";
33-
private const string CannotFindFileBasedOnAssemblyName = "Could not load file or assembly '{0}' or one of its dependencies. The system cannot find the file specified under any probing paths.";
3433
private const string ManifestDefinitionDoesNotMatch = "Could not load file or assembly '{0}' or one of its dependencies. The located assembly's manifest definition does not match the assembly reference.";
3534
private const string AssemblyPathDoesNotExist = "Could not load file or assembly '{0}' or one of its dependencies. The system cannot find the file specified.";
3635
private const string InvalidAssemblyExtensionName = "Could not load file or assembly '{0}' or one of its dependencies. The file specified is not a DLL file.";
3736
private const string AbsolutePathRequired = "Absolute path information is required.";
37+
private const string SingletonAlreadyInitialized = "The singleton of PowerShellAssemblyLoadContext has already been initialized.";
38+
private const string UseResolvingEventHandlerOnly = "PowerShellAssemblyLoadContext was initialized to use its 'Resolving' event handler only.";
3839

3940
#endregion Resource_Strings
4041

4142
#region Constructor
4243

4344
/// <summary>
44-
/// This constructor is for testability purpose only
45+
/// Initialize a singleton of PowerShellAssemblyLoadContext
4546
/// </summary>
46-
protected PowerShellAssemblyLoadContext()
47+
internal static PowerShellAssemblyLoadContext InitializeSingleton(string basePaths, bool useResolvingHandlerOnly)
4748
{
49+
lock (syncObj)
50+
{
51+
if (Instance != null)
52+
throw new InvalidOperationException(SingletonAlreadyInitialized);
53+
54+
Instance = new PowerShellAssemblyLoadContext(basePaths, useResolvingHandlerOnly);
55+
return Instance;
56+
}
4857
}
4958

5059
/// <summary>
@@ -54,7 +63,25 @@ protected PowerShellAssemblyLoadContext()
5463
/// Base directory paths that are separated by semicolon ';'.
5564
/// They will be the default paths to probe assemblies.
5665
/// </param>
57-
internal PowerShellAssemblyLoadContext(string basePaths)
66+
/// <param name="useResolvingHandlerOnly">
67+
/// Indicate whether this instance is going to be used as a
68+
/// full fledged ALC, or only its 'Resolve' handler is going
69+
/// to be used.
70+
/// </param>
71+
/// <remarks>
72+
/// When <paramref name="useResolvingHandlerOnly"/> is true, we will register to the 'Resolving' event of the default
73+
/// load context with our 'Resolve' method, and depend on the default load context to resolve/load assemblies for PS.
74+
/// This mode is used when TPA list of the native host only contains .NET Core libraries.
75+
/// In this case, TPA binder will be consulted before hitting our resolving logic. The binding order of Assembly.Load is:
76+
/// TPA binder --> Resolving event
77+
///
78+
/// When <paramref name="useResolvingHandlerOnly"/> is false, we will use this instance as a full fledged load context
79+
/// to resolve/load assemblies for PS. This mode is used when TPA list of the native host contains both .NET Core libraries
80+
/// and PS assemblies.
81+
/// In this case, our Load override will kick in before consulting the TPA binder. The binding order of Assembly.Load is:
82+
/// Load override --> TPA binder --> Resolving event
83+
/// </remarks>
84+
private PowerShellAssemblyLoadContext(string basePaths, bool useResolvingHandlerOnly)
5885
{
5986
#region Validation
6087
if (string.IsNullOrEmpty(basePaths))
@@ -82,25 +109,32 @@ internal PowerShellAssemblyLoadContext(string basePaths)
82109
this.probingPaths = new List<string>(this.basePaths);
83110

84111
// NEXT: Initialize the CoreCLR type catalog dictionary [OrdinalIgnoreCase]
85-
// - Key: namespace qualified type name (FullName)
86-
// - Value: strong name of the TPA that contains the type represented by Key.
87112
coreClrTypeCatalog = InitializeTypeCatalog();
113+
114+
// LAST: Handle useResolvingHandlerOnly flag
115+
this.useResolvingHandlerOnly = useResolvingHandlerOnly;
116+
this.activeLoadContext = useResolvingHandlerOnly ? AssemblyLoadContext.Default : this;
117+
if (useResolvingHandlerOnly)
118+
AssemblyLoadContext.Default.Resolving += Resolve;
88119
}
89120

90121
#endregion Constructor
91122

92123
#region Fields
93-
94-
// Serialized type catalog file
95-
private readonly object syncObj = new object();
124+
125+
private readonly bool useResolvingHandlerOnly;
126+
private readonly AssemblyLoadContext activeLoadContext;
127+
private readonly static object syncObj = new object();
96128
private readonly string[] basePaths;
97129
// Initially, 'probingPaths' only contains psbase path. But every time we load an assembly through 'LoadFrom(string AssemblyPath)', we
98130
// add its parent path to 'probingPaths', so that we are able to support implicit loading of an assembly from the same place where the
99131
// requesting assembly is located.
100132
// We don't need to worry about removing any paths from 'probingPaths', because once an assembly is loaded, it won't be unloaded until
101133
// the current process exits, and thus the assembly itself and its parent folder cannot be deleted or renamed.
102134
private readonly List<string> probingPaths;
103-
// We use dictionary because the generated binary file by DataContractSerializer is about 39% smaller in size than using Hashtable.
135+
// CoreCLR type catalog dictionary
136+
// - Key: namespace qualified type name (FullName)
137+
// - Value: strong name of the TPA that contains the type represented by Key.
104138
private readonly Dictionary<string, string> coreClrTypeCatalog;
105139
private readonly string[] extensions = new string[] { ".ni.dll", ".dll" };
106140

@@ -123,6 +157,18 @@ internal PowerShellAssemblyLoadContext(string basePaths)
123157

124158
#endregion Fields
125159

160+
#region Properties
161+
162+
/// <summary>
163+
/// Singleton instance of PowerShellAssemblyLoadContext
164+
/// </summary>
165+
public static PowerShellAssemblyLoadContext Instance
166+
{
167+
get; private set;
168+
}
169+
170+
#endregion Properties
171+
126172
#region Events
127173

128174
/// <summary>
@@ -139,6 +185,17 @@ internal PowerShellAssemblyLoadContext(string basePaths)
139185
/// Search the file "[assemblyName.Name][.ni].dll" in probing paths. If the file is found and it matches the requested AssemblyName, load it with LoadFromAssemblyPath.
140186
/// </summary>
141187
protected override Assembly Load(AssemblyName assemblyName)
188+
{
189+
if (useResolvingHandlerOnly)
190+
throw new NotSupportedException(UseResolvingEventHandlerOnly);
191+
192+
return Resolve(this, assemblyName);
193+
}
194+
195+
/// <summary>
196+
/// The handler for the Resolving event
197+
/// </summary>
198+
private Assembly Resolve(AssemblyLoadContext loadContext, AssemblyName assemblyName)
142199
{
143200
// Probe the assembly cache
144201
Assembly asmLoaded;
@@ -185,25 +242,16 @@ protected override Assembly Load(AssemblyName assemblyName)
185242
}
186243
}
187244

188-
// We failed to find the file specified
189-
if (!isAssemblyFileFound)
190-
{
191-
ThrowFileNotFoundException(
192-
CannotFindFileBasedOnAssemblyName,
193-
assemblyName.FullName);
194-
}
195-
196-
// We found the file specified, but the found assembly doesn't match the request
197-
if (!isAssemblyFileMatching)
245+
// We failed to find the assembly file; or we found the file, but the assembly file doesn't match the request.
246+
// In this case, return null so that other Resolving event handlers can kick in to resolve the request.
247+
if (!isAssemblyFileFound || !isAssemblyFileMatching)
198248
{
199-
ThrowFileLoadException(
200-
ManifestDefinitionDoesNotMatch,
201-
assemblyName.FullName);
249+
return null;
202250
}
203251

204252
asmLoaded = asmFilePath.EndsWith(".ni.dll", StringComparison.OrdinalIgnoreCase)
205-
? base.LoadFromNativeImagePath(asmFilePath, null)
206-
: base.LoadFromAssemblyPath(asmFilePath);
253+
? loadContext.LoadFromNativeImagePath(asmFilePath, null)
254+
: loadContext.LoadFromAssemblyPath(asmFilePath);
207255
if (asmLoaded != null)
208256
{
209257
// Add the loaded assembly to the cache
@@ -239,14 +287,14 @@ internal Assembly LoadFrom(string assemblyPath)
239287

240288
// Load the assembly through 'LoadFromNativeImagePath' or 'LoadFromAssemblyPath'
241289
asmLoaded = assemblyPath.EndsWith(".ni.dll", StringComparison.OrdinalIgnoreCase)
242-
? base.LoadFromNativeImagePath(assemblyPath, null)
243-
: base.LoadFromAssemblyPath(assemblyPath);
290+
? activeLoadContext.LoadFromNativeImagePath(assemblyPath, null)
291+
: activeLoadContext.LoadFromAssemblyPath(assemblyPath);
244292

245293
if (asmLoaded != null)
246294
{
247295
// Add the loaded assembly to the cache
248296
AssemblyCache.TryAdd(assemblyName.Name, asmLoaded);
249-
// Add the its parent path to our probing paths
297+
// Add its parent path to our probing paths
250298
string parentPath = Path.GetDirectoryName(assemblyPath);
251299
if (!probingPaths.Contains(parentPath))
252300
{
@@ -282,8 +330,8 @@ internal Assembly LoadFrom(Stream assembly)
282330
if (TryGetAssemblyFromCache(assemblyName, out asmLoaded))
283331
return asmLoaded;
284332

285-
// Load the assembly through 'base.LoadFromStream'
286-
asmLoaded = base.LoadFromStream(assembly);
333+
// Load the assembly through 'LoadFromStream'
334+
asmLoaded = activeLoadContext.LoadFromStream(assembly);
287335
if (asmLoaded != null)
288336
{
289337
// Add the loaded assembly to the cache
@@ -534,12 +582,38 @@ private bool IsAssemblyMatching(AssemblyName requestedAssembly, AssemblyName loa
534582
/// </param>
535583
private Assembly GetTrustedPlatformAssembly(string tpaStrongName)
536584
{
537-
// Load the specified TPA. If the TPA is already loaded, it will be somehow
538-
// cached in CoreCLR runtime, and thus calling 'Assembly.Load' again won't
539-
// cause any overhead.
585+
Assembly asmLoaded;
540586
AssemblyName assemblyName = new AssemblyName(tpaStrongName);
541-
Assembly asmLoaded = Assembly.Load(assemblyName);
542-
return asmLoaded;
587+
588+
// With the current standalone-app model of OPS, .NET Core libraries and PS assemblies are mixed together in one folder.
589+
// So when using PSALC as a full fledged ALC in OPS, some TPAs might be loaded by our Load override. In that case, if we
590+
// alwasy call Assembly.Load here to get a TPA, we might end up with a different Assembly instance of the the same TPA
591+
// loaded in the default load context. We want to use the same assembly instance for type resolution in PS to avoid creating
592+
// types and running .NET code from different assembly instances of the same DLL. Therefore, we try our cache first to see
593+
// if the requested TPA is already loaded. If so, we use that one. If not, we load it in default context using Assembly.Load.
594+
// Once a TPA is loaded in the default context, the same Assembly instance will always be used by custom ALC's when they attempt
595+
// to resolve an "Assembly.Load" request for the same TPA.
596+
//
597+
// For in-box PS of NanoServer/IoT and the share-framework host model of OPS, we don't have the mixed libraries/assemblies
598+
// problem, and TPAs are always resolved/loaded by the default context. In those cases, checking our cache would be unnecessary,
599+
// but it won't cause any problems.
600+
601+
// Probe the assembly cache
602+
if (TryGetAssemblyFromCache(assemblyName, out asmLoaded))
603+
return asmLoaded;
604+
605+
// Prepare to load the assembly
606+
lock (syncObj)
607+
{
608+
// Probe the cache again in case it's already loaded
609+
if (TryGetAssemblyFromCache(assemblyName, out asmLoaded))
610+
return asmLoaded;
611+
612+
// The requested TPA is not loaded by PS ALC, so load it in the default load context using Assembly.Load.
613+
// There is no need to add it to our cache. It's cached in the default context.
614+
asmLoaded = Assembly.Load(assemblyName);
615+
return asmLoaded;
616+
}
543617
}
544618

545619
/// <summary>
@@ -564,36 +638,103 @@ private void ThrowFileNotFoundException(string errorTemplate, params object[] ar
564638
}
565639

566640
/// <summary>
567-
/// Set an instance of PowerShellAssemblyLoadContext to be the default Assembly Load Context.
568641
/// This is the managed entry point for Microsoft.PowerShell.CoreCLR.AssemblyLoadContext.dll.
569642
/// </summary>
570-
public static class PowerShellAssemblyLoadContextInitializer
643+
public class PowerShellAssemblyLoadContextInitializer
571644
{
572-
// Porting note: it's much easier to send an LPStr on Linux
573-
private const UnmanagedType stringType =
574-
#if LINUX
575-
UnmanagedType.LPStr
576-
#else
577-
UnmanagedType.LPWStr
578-
#endif
579-
;
645+
private static object[] EmptyArray = new object[0];
580646

581-
public static PowerShellAssemblyLoadContext PSAsmLoadContext;
647+
/// <summary>
648+
/// Create a singleton of PowerShellAssemblyLoadContext.
649+
/// Then register to the Resolving event of the load context that loads this assembly.
650+
/// </summary>
651+
/// <remarks>
652+
/// This method is to be used by native host whose TPA list doesn't include PS assemblies, such as the
653+
/// in-box Nano powershell.exe, the PS remote WinRM plugin, in-box Nano DSC and in-box Nano SCOM agent.
654+
/// </remarks>
655+
/// <param name="basePaths">
656+
/// Base directory paths that are separated by semicolon ';'.
657+
/// They will be the default paths to probe assemblies.
658+
/// </param>
659+
public static void SetPowerShellAssemblyLoadContext([MarshalAs(UnmanagedType.LPWStr)]string basePaths)
660+
{
661+
if (string.IsNullOrEmpty(basePaths))
662+
throw new ArgumentNullException("basePaths");
582663

664+
PowerShellAssemblyLoadContext.InitializeSingleton(basePaths, useResolvingHandlerOnly: true);
665+
}
666+
583667
/// <summary>
584-
/// Set the default Assembly Load Context
668+
/// Create a singleton of PowerShellAssemblyLoadContext.
669+
/// Then load the assembly containing the actual entry point using it.
585670
/// </summary>
586-
public static void SetPowerShellAssemblyLoadContext([MarshalAs(stringType)]string basePaths)
671+
/// <param name="basePaths">
672+
/// Base directory paths that are separated by semicolon ';'.
673+
/// They will be the default paths to probe assemblies.
674+
/// </param>
675+
/// <param name="entryAssemblyName">
676+
/// Name of the assembly that contains the actual entry point.
677+
/// </param>
678+
/// <returns>
679+
/// The assembly that contains the actual entry point.
680+
/// </returns>
681+
public static Assembly InitializeAndLoadEntryAssembly(string basePaths, AssemblyName entryAssemblyName)
587682
{
588683
if (string.IsNullOrEmpty(basePaths))
589-
{
590684
throw new ArgumentNullException("basePaths");
591-
}
592685

593-
if (PSAsmLoadContext == null)
594-
{
595-
PSAsmLoadContext = new PowerShellAssemblyLoadContext(basePaths);
596-
}
686+
if (entryAssemblyName == null)
687+
throw new ArgumentNullException("entryAssemblyName");
688+
689+
var psLoadContext = PowerShellAssemblyLoadContext.InitializeSingleton(basePaths, useResolvingHandlerOnly: false);
690+
return psLoadContext.LoadFromAssemblyName(entryAssemblyName);
691+
}
692+
693+
/// <summary>
694+
/// Create a singleton of PowerShellAssemblyLoadContext.
695+
/// Then call into the actual entry point based on the given assembly name, type name, method name and arguments.
696+
/// </summary>
697+
/// <param name="basePaths">
698+
/// Base directory paths that are separated by semicolon ';'.
699+
/// They will be the default paths to probe assemblies.
700+
/// </param>
701+
/// <param name="entryAssemblyName">
702+
/// Name of the assembly that contains the actual entry point.
703+
/// </param>
704+
/// <param name="entryTypeName">
705+
/// Name of the type that contains the actual entry point.
706+
/// </param>
707+
/// <param name="entryMethodName">
708+
/// Name of the actual entry point method.
709+
/// </param>
710+
/// <param name="args">
711+
/// An array of arguments passed to the entry point method.
712+
/// </param>
713+
/// <returns>
714+
/// The return value of running the entry point method.
715+
/// </returns>
716+
public static object InitializeAndCallEntryMethod(string basePaths, AssemblyName entryAssemblyName, string entryTypeName, string entryMethodName, object[] args)
717+
{
718+
if (string.IsNullOrEmpty(basePaths))
719+
throw new ArgumentNullException("basePaths");
720+
721+
if (entryAssemblyName == null)
722+
throw new ArgumentNullException("entryAssemblyName");
723+
724+
if (string.IsNullOrEmpty(entryTypeName))
725+
throw new ArgumentNullException("entryTypeName");
726+
727+
if (string.IsNullOrEmpty(entryMethodName))
728+
throw new ArgumentNullException("entryMethodName");
729+
730+
args = args ?? EmptyArray;
731+
732+
var psLoadContext = PowerShellAssemblyLoadContext.InitializeSingleton(basePaths, useResolvingHandlerOnly: false);
733+
var entryAssembly = psLoadContext.LoadFromAssemblyName(entryAssemblyName);
734+
var entryType = entryAssembly.GetType(entryTypeName, throwOnError: true, ignoreCase: true);
735+
var methodInfo = entryType.GetMethod(entryMethodName, BindingFlags.Static | BindingFlags.Public | BindingFlags.IgnoreCase);
736+
737+
return methodInfo.Invoke(null, args);
597738
}
598739
}
599740
}

src/Microsoft.PowerShell.CoreCLR.AssemblyLoadContext/project.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
"NETStandard.Library": "1.6.0",
2828
"System.Runtime.Loader": "4.0.0",
2929
"System.Reflection.Metadata": "1.1.0",
30+
"System.Reflection.TypeExtensions": "4.1.0",
3031
"System.Security.Cryptography.Algorithms": "4.2.0"
3132
}
3233
}

0 commit comments

Comments
 (0)