Resetting shared load context during specialization to ensure dependencies are properly resolved

Author: fabiocav

src/WebJobs.Script.WebHost/Standby/IStandbyManager.cs

@@ -3,7 +3,6 @@
 
 using System;
 using System.Threading.Tasks;
-using Microsoft.Extensions.Primitives;
 
 namespace Microsoft.Azure.WebJobs.Script.WebHost
 {

src/WebJobs.Script.WebHost/Standby/StandbyManager.cs

@@ -7,6 +7,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Hosting;
+using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Rpc;
 using Microsoft.Azure.WebJobs.Script.WebHost.Properties;
 using Microsoft.Extensions.Configuration;
@@ -97,9 +98,14 @@ public async Task SpecializeHostCoreAsync()
 
             _hostNameProvider.Reset();
 
+            // Reset the shared load context to ensure we're reloading
+            // user dependencies
+            FunctionAssemblyLoadContext.ResetSharedContext();
+
             await _languageWorkerChannelManager.SpecializeAsync();
 
             NotifyChange();
+
             await _scriptHostManager.RestartHostAsync();
             await _scriptHostManager.DelayUntilHostReady();
         }

src/WebJobs.Script/Description/DotNet/FunctionAssemblyLoadContext.cs

@@ -9,6 +9,7 @@
 using System.Reflection;
 using System.Runtime.InteropServices;
 using System.Runtime.Loader;
+using System.Threading;
 using Microsoft.Azure.WebJobs.Host;
 using Microsoft.Azure.WebJobs.Script.Config;
 using Microsoft.Extensions.DependencyModel;
@@ -24,7 +25,8 @@ public partial class FunctionAssemblyLoadContext : AssemblyLoadContext
         private static readonly Lazy<Dictionary<string, ScriptRuntimeAssembly>> _runtimeAssemblies = new Lazy<Dictionary<string, ScriptRuntimeAssembly>>(DependencyHelper.GetRuntimeAssemblies);
         private static readonly Lazy<Dictionary<string, ResolutionPolicyEvaluator>> _resolutionPolicyEvaluators = new Lazy<Dictionary<string, ResolutionPolicyEvaluator>>(InitializeLoadPolicyEvaluators);
         private static readonly ConcurrentDictionary<string, object> _sharedContextAssembliesInFallbackLoad = new ConcurrentDictionary<string, object>();
-        private static readonly Lazy<FunctionAssemblyLoadContext> _defaultContext = new Lazy<FunctionAssemblyLoadContext>(CreateSharedContext, true);
+
+        private static Lazy<FunctionAssemblyLoadContext> _defaultContext = new Lazy<FunctionAssemblyLoadContext>(CreateSharedContext, true);
 
         private readonly List<string> _probingPaths = new List<string>();
         private readonly IDictionary<string, string> _depsAssemblies;
@@ -43,6 +45,11 @@ public FunctionAssemblyLoadContext(string basePath)
 
         public static FunctionAssemblyLoadContext Shared => _defaultContext.Value;
 
+        internal static void ResetSharedContext()
+        {
+            _defaultContext = new Lazy<FunctionAssemblyLoadContext>(CreateSharedContext, true);
+        }
+
         internal static IDictionary<string, string> InitializeDeps(string basePath)
         {
             string depsFilePath = Path.Combine(basePath, DotNetConstants.FunctionsDepsFileName);

src/WebJobs.Script/ScriptHostBuilderExtensions.cs

@@ -204,7 +204,8 @@ public static void AddCommonServices(IServiceCollection services)
 
         public static IWebJobsBuilder UseScriptExternalStartup(this IWebJobsBuilder builder, string rootScriptPath, ILoggerFactory loggerFactory, IExtensionBundleManager extensionBundleManager)
         {
-            var logger = loggerFactory.CreateLogger<ScriptStartupTypeLocator>() ?? throw new ArgumentNullException(nameof(loggerFactory));
+            var logger = loggerFactory?.CreateLogger<ScriptStartupTypeLocator>() ?? throw new ArgumentNullException(nameof(loggerFactory));
+
             return builder.UseExternalStartup(new ScriptStartupTypeLocator(rootScriptPath, logger, extensionBundleManager));
         }
 

test/WebJobs.Script.Tests.Integration/WebHostEndToEnd/SpecializationE2ETests.cs

@@ -16,6 +16,7 @@
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Azure.WebJobs.Logging.ApplicationInsights;
 using Microsoft.Azure.WebJobs.Script.Configuration;
+using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.Rpc;
 using Microsoft.Azure.WebJobs.Script.WebHost;
 using Microsoft.Extensions.Configuration;
@@ -187,6 +188,34 @@ public async Task Specialization_ThreadUtilization()
             }
         }
 
+        [Fact]
+        public async Task Specialization_ResetsSharedLoadContext()
+        {
+            var builder = CreateStandbyHostBuilder("FunctionExecutionContext");
+
+            using (var testServer = new TestServer(builder))
+            {
+                var client = testServer.CreateClient();
+
+                var response = await client.GetAsync("api/warmup");
+                response.EnsureSuccessStatusCode();
+
+                var placeholderContext = FunctionAssemblyLoadContext.Shared;
+
+                _environment.SetEnvironmentVariable(EnvironmentSettingNames.AzureWebsiteContainerReady, "1");
+                _environment.SetEnvironmentVariable(EnvironmentSettingNames.AzureWebsitePlaceholderMode, "0");
+
+                //await _pauseBeforeHostBuild.WaitAsync(10000);
+
+                response = await client.GetAsync("api/functionexecutioncontext");
+                response.EnsureSuccessStatusCode();
+
+                var specializedContext = FunctionAssemblyLoadContext.Shared;
+
+                Assert.NotSame(placeholderContext, specializedContext);
+            }
+        }
+
         [Fact]
         public async Task StartAsync_SetsCorrectActiveHost()
         {