Assembly loading improvements

.gitignore

@@ -239,3 +239,6 @@ src/DataFactory/DataFactoryV2.Test/SessionRecords/Microsoft.Azure.Commands.DataF
 
 # GitHub codespaces
 .venv
+
+# Visual studio live unit testing config
+*.lutconfig

build.proj

@@ -205,7 +205,7 @@
       <RuntimeDllsIncludeList>Microsoft.Powershell.*.dll,System*.dll,Microsoft.VisualBasic.dll,Microsoft.CSharp.dll,Microsoft.CodeAnalysis.dll,Microsoft.CodeAnalysis.CSharp.dll</RuntimeDllsIncludeList>
       <RuntimeDllsExcludeList>System.Security.Cryptography.ProtectedData.dll,System.Configuration.ConfigurationManager.dll,System.Runtime.CompilerServices.Unsafe.dll,System.IO.FileSystem.AccessControl.dll,System.Buffers.dll,System.Text.Encodings.Web.dll,System.CodeDom.dll,System.Management.dll,System.Text.Json.dll,System.Threading.Tasks.Extensions.dll,System.IO.Hashing.dll</RuntimeDllsExcludeList>
     </PropertyGroup>
-    <Exec Command="$(PowerShellCoreCommandPrefix) &quot;Get-ChildItem -Path $(RepoArtifacts)/$(Configuration) -Recurse -Include $(RuntimeDllsIncludeList) -Exclude $(RuntimeDllsExcludeList) | Where-Object {$_.FullName -notlike '*PreloadAssemblies*' -and $_.FullName -notlike '*NetCoreAssemblies*' -and $_.FullName -notlike '*AzSharedAlcAssemblies*' -and $_.FullName -notlike '*ModuleAlcAssemblies*'} | Remove-Item -Force&quot;"/>
+    <Exec Command="$(PowerShellCoreCommandPrefix) &quot;Get-ChildItem -Path $(RepoArtifacts)/$(Configuration) -Recurse -Include $(RuntimeDllsIncludeList) -Exclude $(RuntimeDllsExcludeList) | Where-Object {$_.FullName -notlike '*Accounts*lib*' -and $_.FullName -notlike '*ModuleAlcAssemblies*'} | Remove-Item -Force&quot;"/>
     <Exec Command="$(PowerShellCoreCommandPrefix) &quot;Get-ChildItem -Path $(RepoArtifacts)/$(Configuration) -Recurse -Include 'runtimes' | Remove-Item -Recurse -Force&quot;" Condition="'$(CodeSign)' == 'true'" />
 
     <Exec Command="$(PowerShellCoreCommandPrefix) &quot;. $(RepoTools)/UpdateModules.ps1 -BuildConfig $(Configuration) -Scope $(Scope)&quot;"/>

src/Accounts/Accounts.sln

@@ -1,6 +1,6 @@
 Microsoft Visual Studio Solution File, Format Version 12.00
-# Visual Studio Version 16
-VisualStudioVersion = 16.0.30503.244
+# Visual Studio Version 17
+VisualStudioVersion = 17.3.32929.385
 MinimumVisualStudioVersion = 10.0.40219.1
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Accounts", "Accounts\Accounts.csproj", "{142D7B0B-388A-4CEB-A228-7F6D423C5C2E}"
 EndProject
@@ -22,6 +22,10 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Authentication.Test", "Auth
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "AuthenticationAssemblyLoadContext", "AuthenticationAssemblyLoadContext\AuthenticationAssemblyLoadContext.csproj", "{D06EF9FC-4C3E-4A51-A4AD-D39AD426EF8C}"
 EndProject
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "AssemblyLoading", "AssemblyLoading\AssemblyLoading.csproj", "{74C0BD7E-DFFD-4B96-818E-D8660FA43159}"
+EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "AssemblyLoading.Test", "AssemblyLoading.Test\AssemblyLoading.Test.csproj", "{D4540550-9808-4DEB-9D5E-F88E38D58A85}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -64,13 +68,22 @@ Global
 		{D06EF9FC-4C3E-4A51-A4AD-D39AD426EF8C}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{D06EF9FC-4C3E-4A51-A4AD-D39AD426EF8C}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{D06EF9FC-4C3E-4A51-A4AD-D39AD426EF8C}.Release|Any CPU.Build.0 = Release|Any CPU
+		{74C0BD7E-DFFD-4B96-818E-D8660FA43159}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{74C0BD7E-DFFD-4B96-818E-D8660FA43159}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{74C0BD7E-DFFD-4B96-818E-D8660FA43159}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{74C0BD7E-DFFD-4B96-818E-D8660FA43159}.Release|Any CPU.Build.0 = Release|Any CPU
+		{D4540550-9808-4DEB-9D5E-F88E38D58A85}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{D4540550-9808-4DEB-9D5E-F88E38D58A85}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{D4540550-9808-4DEB-9D5E-F88E38D58A85}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{D4540550-9808-4DEB-9D5E-F88E38D58A85}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
 	EndGlobalSection
 	GlobalSection(NestedProjects) = preSolution
 		{152D78F0-A642-4D0E-B3A8-2FC64FFA9714} = {95C16AED-FD57-42A0-86C3-2CF4300A4817}
 		{43BE9983-BD21-4474-92E5-1FFC0220B2AD} = {95C16AED-FD57-42A0-86C3-2CF4300A4817}
+		{D4540550-9808-4DEB-9D5E-F88E38D58A85} = {95C16AED-FD57-42A0-86C3-2CF4300A4817}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {AA51E4F8-AA75-429D-9626-12C7B4305D72}

src/Accounts/Accounts/Accounts.csproj

@@ -11,8 +11,10 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PreLoadAssemblies Include="$(RepoSrc)lib\NetFxPreloadAssemblies\*.dll" />
-    <NetCoreAssemblies Include="$(RepoSrc)lib\NetCorePreloadAssemblies\*.dll" />
+    <NetFxAssemblies Include="$(RepoSrc)lib\netfx\*.dll" />
+    <NetCore21Assemblies Include="$(RepoSrc)lib\netcoreapp2.1\*.dll" />
+    <NetCore31Assemblies Include="$(RepoSrc)lib\netcoreapp3.1\*.dll" />
+    <NetStandardAssemblies Include="$(RepoSrc)lib\netstandard2.0\*.dll" />
     <StorageDependencies Include="$(RepoSrc)lib\WindowsAzure.Storage\9.3.0\Microsoft.WindowsAzure.Storage.dll" />
     <StorageDependencies Include="$(RepoSrc)lib\WindowsAzure.Storage\9.3.0\Microsoft.WindowsAzure.Storage.DataMovement.dll" />
   </ItemGroup>
@@ -22,10 +24,12 @@
   </ItemGroup>
 
   <Target Name="CopyFiles" AfterTargets="Build">
-    <Copy SourceFiles="@(PreLoadAssemblies)" DestinationFolder="$(TargetDir)PreloadAssemblies" />
-    <Copy SourceFiles="@(NetCoreAssemblies)" DestinationFolder="$(TargetDir)AzSharedAlcAssemblies" />
+    <Copy SourceFiles="@(NetFxAssemblies)" DestinationFolder="$(TargetDir)lib\netfx" />
+    <Copy SourceFiles="@(NetCore21Assemblies)" DestinationFolder="$(TargetDir)lib\netcoreapp2.1" />
+    <Copy SourceFiles="@(NetCore31Assemblies)" DestinationFolder="$(TargetDir)lib\netcoreapp3.1" />
+    <Copy SourceFiles="@(NetStandardAssemblies)" DestinationFolder="$(TargetDir)lib\netstandard2.0" />
     <Copy SourceFiles="@(StorageDependencies)" DestinationFolder="$(TargetDir)" />
-    <!--To automatically lauch pwsh for debugging when hitting F5, set below and UNCOMMENT Copy:
+    <!--To automatically launch pwsh for debugging when hitting F5, set below and UNCOMMENT Copy:
       Executable: pwsh
       Application arguments: -NoExit -Command "Import-Module C:\AME\azure-powershell\artifacts\Debug\Az.Accounts\Az.Accounts.psd1"-->
     <!--<Copy SourceFiles="@(NetCoreAssemblies)" DestinationFolder="$(TargetDir)" />-->
@@ -43,27 +47,27 @@
     <EmbeddedResource Include="AzureRmAlias\Mappings.json" />
     <EmbeddedResource Include="Utilities\CommandMappings.json" />
   </ItemGroup>
-   
+
   <ItemGroup>
     <ProjectReference Include="..\Authentication.ResourceManager\Authentication.ResourceManager.csproj" />
     <ProjectReference Include="..\Authentication\Authentication.csproj" />
-    <ProjectReference Include="..\Authenticators\Authenticators.csproj" />  
+    <ProjectReference Include="..\Authenticators\Authenticators.csproj" />
   </ItemGroup>
-   
+
   <ItemGroup>
     <Reference Include="FuzzySharp">
       <HintPath>..\..\lib\FuzzySharp.dll</HintPath>
     </Reference>
   </ItemGroup>
-   
+
   <ItemGroup>
     <Compile Update="Properties\Resources.Designer.cs">
       <DesignTime>True</DesignTime>
       <AutoGen>True</AutoGen>
       <DependentUpon>Resources.resx</DependentUpon>
     </Compile>
   </ItemGroup>
-   
+
   <ItemGroup>
     <EmbeddedResource Update="Properties\Resources.resx">
       <Generator>ResXFileCodeGenerator</Generator>

src/Accounts/Accounts/Az.Accounts.psd1

@@ -56,7 +56,8 @@ DotNetFrameworkVersion = '4.7.2'
 # RequiredModules = @()
 
 # Assemblies that must be loaded prior to importing this module
-RequiredAssemblies = 'Microsoft.Azure.PowerShell.Authentication.Abstractions.dll', 
+RequiredAssemblies = 'Microsoft.Azure.PowerShell.AssemblyLoading.dll',
+               'Microsoft.Azure.PowerShell.Authentication.Abstractions.dll', 
                'Microsoft.Azure.PowerShell.Authentication.dll', 
                'Microsoft.Azure.PowerShell.Authenticators.dll', 
                'Microsoft.Azure.PowerShell.Authentication.ResourceManager.dll', 

src/Accounts/Accounts/ChangeLog.md

@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Optimized the mechanism for assembly loading.
 * Enabled AzKeyStore with keyring in Linux.
 
 ## Version 2.10.4

src/Accounts/Accounts/StartupScripts/InitializeAssemblyResolver.ps1

@@ -1,4 +1,8 @@
-if ($PSEdition -eq 'Desktop') {
+$assemblyRootPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "lib")
+$conditionalAssemblyContext = [Microsoft.Azure.PowerShell.AssemblyLoading.ConditionalAssemblyContext]::new($Host.Version)
+[Microsoft.Azure.PowerShell.AssemblyLoading.ConditionalAssemblyProvider]::Initialize($assemblyRootPath, $conditionalAssemblyContext)
+
+if ($PSEdition -eq 'Desktop') {
   try {
     [Microsoft.Azure.Commands.Profile.Utilities.CustomAssemblyResolver]::Initialize()
   }
@@ -9,9 +13,8 @@
 else {
   try {
     Add-Type -Path ([System.IO.Path]::Combine($PSScriptRoot, "..", "Microsoft.Azure.PowerShell.AuthenticationAssemblyLoadContext.dll")) | Out-Null
-    $assemblyLoadContextFolder = [System.IO.Path]::Combine($PSScriptRoot, "..", "AzSharedAlcAssemblies")
-    Write-Debug "Registering Az shared AssemblyLoadContext for path: '$assemblyLoadContextFolder'."
-    [Microsoft.Azure.PowerShell.AuthenticationAssemblyLoadContext.AzAssemblyLoadContextInitializer]::RegisterAzSharedAssemblyLoadContext($assemblyLoadContextFolder)
+    Write-Debug "Registering Az shared AssemblyLoadContext."
+    [Microsoft.Azure.PowerShell.AuthenticationAssemblyLoadContext.AzAssemblyLoadContextInitializer]::RegisterAzSharedAssemblyLoadContext()
     Write-Debug "AssemblyLoadContext registered."
   }
   catch {

src/Accounts/AssemblyLoading.Test/AssemblyLoading.Test.csproj

@@ -0,0 +1,17 @@
+<Project Sdk="Microsoft.NET.Sdk">
+	<PropertyGroup>
+		<PsModuleName>Accounts</PsModuleName>
+	</PropertyGroup>
+
+	<Import Project="$(MSBuildThisFileDirectory)..\..\Az.Test.props" />
+
+	<PropertyGroup>
+		<AssemblyName>Microsoft.Azure.PowerShell.AssemblyLoading.Test</AssemblyName>
+		<RootNamespace>Microsoft.Azure.PowerShell.AssemblyLoading.Test</RootNamespace>
+	</PropertyGroup>
+
+	<ItemGroup>
+	  <ProjectReference Include="..\AssemblyLoading\AssemblyLoading.csproj" />
+	</ItemGroup>
+
+</Project>

src/Accounts/AssemblyLoading.Test/Mocks/MockConditionalAssembly.cs

@@ -0,0 +1,40 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using System;
+
+namespace Microsoft.Azure.PowerShell.AssemblyLoading.Test.Mocks
+{
+    internal class MockConditionalAssembly : IConditionalAssembly
+    {
+        public MockConditionalAssembly(IConditionalAssemblyContext context)
+        {
+            Context = context;
+        }
+        public bool ShouldLoad { get; set; } = true;
+
+        public string Name { get; set; }
+
+        public string Path { get; set; }
+
+        public Version Version { get; set; }
+
+        public IConditionalAssemblyContext Context { get; set; }
+
+        public void UpdateShouldLoad(bool shouldLoad)
+        {
+            ShouldLoad = ShouldLoad && shouldLoad;
+        }
+    }
+}

src/Accounts/AssemblyLoading.Test/Mocks/MockConditionalAssemblyContext.cs

@@ -0,0 +1,31 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using System;
+using System.Runtime.InteropServices;
+
+namespace Microsoft.Azure.PowerShell.AssemblyLoading.Test.Mocks
+{
+    internal class MockConditionalAssemblyContext : IConditionalAssemblyContext
+    {
+        public Version PSVersion { get; set; }
+        public Architecture OSArchitecture { get; set; }
+        public OSPlatform OS { get; set; }
+
+        public bool IsOSPlatform(OSPlatform os)
+        {
+            return OS.Equals(os);
+        }
+    }
+}

src/Accounts/AssemblyLoading.Test/UnitTests/ConditionalAssemblyExtensionsTests.cs

@@ -0,0 +1,82 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Microsoft.Azure.PowerShell.AssemblyLoading.Test.Mocks;
+using Microsoft.WindowsAzure.Commands.ScenarioTest;
+using System;
+using System.Runtime.InteropServices;
+using Xunit;
+
+namespace Microsoft.Azure.PowerShell.AssemblyLoading.Test.UnitTests
+{
+    public class ConditionalAssemblyExtensionsTests
+    {
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void CanWorkWithPSVersion()
+        {
+            var windowsPSContext = new MockConditionalAssemblyContext()
+            {
+                PSVersion = Version.Parse("5.1.22621.608")
+            };
+            var windowsPSAssembly = new MockConditionalAssembly(windowsPSContext)
+                .WithWindowsPowerShell();
+            var psCoreAssembly = new MockConditionalAssembly(
+                windowsPSContext)
+                .WithPowerShellCore();
+            Assert.True(windowsPSAssembly.ShouldLoad);
+            Assert.False(psCoreAssembly.ShouldLoad);
+
+            var ps7Context = new MockConditionalAssemblyContext()
+            {
+                PSVersion = Version.Parse("7.3.0")
+            };
+            windowsPSAssembly = new MockConditionalAssembly(
+                ps7Context)
+                .WithWindowsPowerShell();
+            psCoreAssembly = new MockConditionalAssembly(
+                ps7Context)
+                .WithPowerShellCore();
+            Assert.True(psCoreAssembly.ShouldLoad);
+            Assert.False(windowsPSAssembly.ShouldLoad);
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void CanWorkWithOS()
+        {
+            var windowsContext = new MockConditionalAssemblyContext()
+            {
+                OS = OSPlatform.Windows
+            };
+            var windowsAssembly = new MockConditionalAssembly(windowsContext)
+                .WithWindows();
+            var linuxAssembly = new MockConditionalAssembly(windowsContext)
+                .WithLinux();
+            Assert.True(windowsAssembly.ShouldLoad);
+            Assert.False(linuxAssembly.ShouldLoad);
+
+            var linuxContext = new MockConditionalAssemblyContext()
+            {
+                OS = OSPlatform.Linux
+            };
+            windowsAssembly = new MockConditionalAssembly(linuxContext)
+                .WithWindows();
+            linuxAssembly = new MockConditionalAssembly(linuxContext)
+                .WithLinux();
+            Assert.False(windowsAssembly.ShouldLoad);
+            Assert.True(linuxAssembly.ShouldLoad);
+        }
+    }
+}