Create only one Runspace instance at the FunctionLoad stage (#201)

This is to fix the regression in local debugging.
We create the first `PowerShellManager` instance when processing the first `FunctionLoad` request, so the debugger has the target Runspace to attach to.
However, we delay the initialization of the first `PowerShellManager` instance until the first invocation (when checking out the instance from pool).

It turns out we cannot have the dependency manager to create a separate Runspace for downloading.
The VS Code local debugging assumes the target `Runspace` has `Id = 1`, which was true before as only one Runspace will be created initially.
Now, the dependency manager creates a separate `Runspace` for downloading, and there will be a race condition -- sometimes the `Runspace` created for the downloading has `Id = 1`, and thus the debugger will attach to that Runspace, which will be reclaimed after the download and causing the `Wait-Debugger` in the user's function script to not be hit.

This refactoring makes sure only one `Runspace` is created initially -- the `PowerShell` instance used for the downloading is also used to serve the first invocation request.
So there is no race condition and local debugging can work reliably.

Author: daxian-dbw

src/DependencyManagement/DependencyManager.cs

@@ -79,7 +79,8 @@ internal DependencyManager()
         /// </summary>
         /// <param name="msgStream">The protobuf messaging stream</param>
         /// <param name="request">The StreamingMessage request for function load</param>
-        internal void ProcessDependencyDownload(MessagingStream msgStream, StreamingMessage request)
+        /// <param name="pwsh">The PowerShell instance used to download modules</param>
+        internal void ProcessDependencyDownload(MessagingStream msgStream, StreamingMessage request, PowerShell pwsh)
         {
             if (request.FunctionLoadRequest.ManagedDependencyEnabled)
             {
@@ -100,7 +101,7 @@ internal void ProcessDependencyDownload(MessagingStream msgStream, StreamingMess
                 }
 
                 //Start dependency download on a separate thread
-                _dependencyDownloadTask = Task.Run(() => ProcessDependencies(rpcLogger));
+                _dependencyDownloadTask = Task.Run(() => InstallFunctionAppDependencies(pwsh, rpcLogger));
             }
         }
 
@@ -117,22 +118,6 @@ internal void WaitOnDependencyDownload()
             }
         }
 
-        private void ProcessDependencies(RpcLogger rpcLogger)
-        {
-            try
-            {
-                _dependencyError = null;
-                using (PowerShell pwsh = PowerShell.Create(Utils.SingletonISS.Value))
-                {
-                    InstallFunctionAppDependencies(pwsh, rpcLogger);
-                }
-            }
-            catch (Exception e)
-            {
-                _dependencyError = e;
-            }
-        }
-
         /// <summary>
         /// Initializes the dependency manger and performs the following:
         /// - Parse functionAppRoot\requirements.psd1 file and create a list of dependencies to install.
@@ -251,7 +236,7 @@ internal void InstallFunctionAppDependencies(PowerShell pwsh, ILogger logger)
             catch (Exception e)
             {
                 var errorMsg = string.Format(PowerShellWorkerStrings.FailToInstallFuncAppDependencies, e.Message);
-                throw new DependencyInstallationException(errorMsg, e);
+                _dependencyError = new DependencyInstallationException(errorMsg, e);
             }
         }
 

src/PowerShell/PowerShellManager.cs

@@ -22,6 +22,7 @@ internal class PowerShellManager
     {
         private readonly ILogger _logger;
         private readonly PowerShell _pwsh;
+        private bool _runspaceInited;
 
         /// <summary>
         /// Gets the Runspace InstanceId.
@@ -45,27 +46,64 @@ static PowerShellManager()
             addMethod.Invoke(null, new object[] { "HttpRequestContext", typeof(HttpRequestContext) });
         }
 
-        internal PowerShellManager(ILogger logger)
+        /// <summary>
+        /// Constructor for setting the basic fields.
+        /// </summary>
+        private PowerShellManager(ILogger logger, PowerShell pwsh, int id)
+        {
+            _logger = logger;
+            _pwsh = pwsh;
+            _pwsh.Runspace.Name = $"PowerShellManager{id}";
+        }
+
+        /// <summary>
+        /// Create a PowerShellManager instance but defer the Initialization.
+        /// </summary>
+        /// <remarks>
+        /// This constructor is only for creating the very first PowerShellManager instance.
+        /// The initialization work is deferred until all prerequisites are ready, such as
+        /// the dependent modules are downloaded and all Az functions are loaded.
+        /// </remarks>
+        internal PowerShellManager(ILogger logger, PowerShell pwsh)
+            : this(logger, pwsh, id: 1)
         {
-            if (FunctionLoader.FunctionAppRootPath == null)
+        }
+
+        /// <summary>
+        /// Create a PowerShellManager instance and initialize it.
+        /// </summary>
+        internal PowerShellManager(ILogger logger, int id)
+            : this(logger, Utils.NewPwshInstance(), id)
+        {
+            // Initialize the Runspace
+            Initialize();
+        }
+
+        /// <summary>
+        /// Extra initialization of the Runspace.
+        /// </summary>
+        internal void Initialize()
+        {
+            if (!_runspaceInited)
             {
-                throw new InvalidOperationException(PowerShellWorkerStrings.FunctionAppRootNotResolved);
+                RegisterStreamEvents();
+                InvokeProfile(FunctionLoader.FunctionAppProfilePath);
+                _runspaceInited = true;
             }
+        }
 
-            _logger = logger;
-            _pwsh = PowerShell.Create(Utils.SingletonISS.Value);
-
-            // Setup Stream event listeners
-            var streamHandler = new StreamHandler(logger);
+        /// <summary>
+        /// Setup Stream event listeners.
+        /// </summary>
+        private void RegisterStreamEvents()
+        {
+            var streamHandler = new StreamHandler(_logger);
             _pwsh.Streams.Debug.DataAdding += streamHandler.DebugDataAdding;
             _pwsh.Streams.Error.DataAdding += streamHandler.ErrorDataAdding;
             _pwsh.Streams.Information.DataAdding += streamHandler.InformationDataAdding;
             _pwsh.Streams.Progress.DataAdding += streamHandler.ProgressDataAdding;
             _pwsh.Streams.Verbose.DataAdding += streamHandler.VerboseDataAdding;
             _pwsh.Streams.Warning.DataAdding += streamHandler.WarningDataAdding;
-
-            // Initialize the Runspace
-            InvokeProfile(FunctionLoader.FunctionAppProfilePath);
         }
 
         /// <summary>
@@ -76,7 +114,8 @@ internal void InvokeProfile(string profilePath)
             Exception exception = null;
             if (profilePath == null)
             {
-                RpcLogger.WriteSystemLog(string.Format(PowerShellWorkerStrings.FileNotFound, "profile.ps1", FunctionLoader.FunctionAppRootPath));
+                string noProfileMsg = string.Format(PowerShellWorkerStrings.FileNotFound, "profile.ps1", FunctionLoader.FunctionAppRootPath);
+                _logger.Log(LogLevel.Trace, noProfileMsg);
                 return;
             }
 

src/PowerShell/PowerShellManagerPool.cs

@@ -46,6 +46,18 @@ internal PowerShellManagerPool(MessagingStream msgStream)
             RpcLogger.WriteSystemLog(string.Format(PowerShellWorkerStrings.LogConcurrencyUpperBound, _upperBound.ToString()));
         }
 
+        /// <summary>
+        /// Populate the pool with the very first PowerShellManager instance.
+        /// We instantiate PowerShellManager instances in a lazy way, starting from size 1 and increase the number of workers as needed.
+        /// </summary>
+        internal void Initialize(PowerShell pwsh)
+        {
+            var logger = new RpcLogger(_msgStream);
+            var psManager = new PowerShellManager(logger, pwsh);
+            _pool.Add(psManager);
+            _poolSize = 1;
+        }
+
         /// <summary>
         /// Checkout an idle PowerShellManager instance in a non-blocking asynchronous way.
         /// </summary>
@@ -59,28 +71,37 @@ internal PowerShellManager CheckoutIdleWorker(StreamingMessage request, AzFuncti
             if (!_pool.TryTake(out psManager))
             {
                 // The pool doesn't have an idle one.
-                if (_poolSize < _upperBound &&
-                    Interlocked.Increment(ref _poolSize) <= _upperBound)
+                if (_poolSize < _upperBound)
                 {
-                    // If the pool hasn't reached its bounded capacity yet, then
-                    // we create a new item and return it.
-                    var logger = new RpcLogger(_msgStream);
-                    logger.SetContext(requestId, invocationId);
-                    psManager = new PowerShellManager(logger);
+                    int id = Interlocked.Increment(ref _poolSize);
+                    if (id <= _upperBound)
+                    {
+                        // If the pool hasn't reached its bounded capacity yet, then
+                        // we create a new item and return it.
+                        var logger = new RpcLogger(_msgStream);
+                        logger.SetContext(requestId, invocationId);
+                        psManager = new PowerShellManager(logger, id);
 
-                    RpcLogger.WriteSystemLog(string.Format(PowerShellWorkerStrings.LogNewPowerShellManagerCreated, _poolSize.ToString()));
+                        RpcLogger.WriteSystemLog(string.Format(PowerShellWorkerStrings.LogNewPowerShellManagerCreated, id.ToString()));
+                    }
                 }
-                else
+
+                if (psManager == null)
                 {
                     // If the pool has reached its bounded capacity, then the thread
                     // should be blocked until an idle one becomes available.
                     psManager = _pool.Take();
                 }
             }
 
+            // Finish the initialization if not yet.
+            // This applies only to the very first PowerShellManager instance, whose initialization was deferred.
+            psManager.Initialize();
+
             // Register the function with the Runspace before returning the idle PowerShellManager.
             FunctionMetadata.RegisterFunctionMetadata(psManager.InstanceId, functionInfo);
             psManager.Logger.SetContext(requestId, invocationId);
+
             return psManager;
         }
 

src/RequestProcessor.cs

@@ -170,7 +170,15 @@ internal StreamingMessage ProcessFunctionLoadRequest(StreamingMessage request)
 
                     // Setup the FunctionApp root path and module path.
                     FunctionLoader.SetupWellKnownPaths(functionLoadRequest);
-                    _dependencyManager.ProcessDependencyDownload(_msgStream, request);
+
+                    // Create the very first Runspace so the debugger has the target to attach to.
+                    // This PowerShell instance is shared by the first PowerShellManager instance created in the pool,
+                    // and the dependency manager (used to download dependent modules if needed).
+                    var pwsh = Utils.NewPwshInstance();
+                    _powershellPool.Initialize(pwsh);
+
+                    // Start the download asynchronously if needed.
+                    _dependencyManager.ProcessDependencyDownload(_msgStream, request, pwsh);
                 }
                 catch (Exception e)
                 {

src/Utility/Utils.cs

@@ -5,43 +5,54 @@
 
 using System;
 using System.IO;
-using System.Management.Automation;
-using System.Management.Automation.Runspaces;
 using System.Text;
 using System.Threading;
 using Microsoft.PowerShell.Commands;
 
 namespace Microsoft.Azure.Functions.PowerShellWorker.Utility
 {
+    using System.Management.Automation;
+    using System.Management.Automation.Runspaces;
+
     internal class Utils
     {
         internal readonly static CmdletInfo ImportModuleCmdletInfo = new CmdletInfo("Import-Module", typeof(ImportModuleCommand));
         internal readonly static CmdletInfo RemoveModuleCmdletInfo = new CmdletInfo("Remove-Module", typeof(RemoveModuleCommand));
         internal readonly static CmdletInfo GetJobCmdletInfo = new CmdletInfo("Get-Job", typeof(GetJobCommand));
         internal readonly static CmdletInfo RemoveJobCmdletInfo = new CmdletInfo("Remove-Job", typeof(RemoveJobCommand));
 
-        internal readonly static Lazy<InitialSessionState> SingletonISS
-            = new Lazy<InitialSessionState>(NewInitialSessionState, LazyThreadSafetyMode.PublicationOnly);
+        private static InitialSessionState s_iss;
 
-        private static InitialSessionState NewInitialSessionState()
+        /// <summary>
+        /// Create a new PowerShell instance using our singleton InitialSessionState instance.
+        /// </summary>
+        internal static PowerShell NewPwshInstance()
         {
-            var iss = InitialSessionState.CreateDefault();
-            iss.ThreadOptions = PSThreadOptions.UseCurrentThread;
-            iss.EnvironmentVariables.Add(
-                new SessionStateVariableEntry(
-                    "PSModulePath",
-                    FunctionLoader.FunctionModulePath,
-                    description: null));
-
-            // Setting the execution policy on macOS and Linux throws an exception so only update it on Windows
-            if(Platform.IsWindows)
+            if (s_iss == null)
             {
-                // This sets the execution policy on Windows to Unrestricted which is required to run the user's function scripts on
-                // Windows client versions. This is needed if a user is testing their function locally with the func CLI.
-                iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Unrestricted;
+                if (FunctionLoader.FunctionAppRootPath == null)
+                {
+                    throw new InvalidOperationException(PowerShellWorkerStrings.FunctionAppRootNotResolved);
+                }
+
+                s_iss = InitialSessionState.CreateDefault();
+                s_iss.ThreadOptions = PSThreadOptions.UseCurrentThread;
+                s_iss.EnvironmentVariables.Add(
+                    new SessionStateVariableEntry(
+                        "PSModulePath",
+                        FunctionLoader.FunctionModulePath,
+                        description: null));
+
+                // Setting the execution policy on macOS and Linux throws an exception so only update it on Windows
+                if(Platform.IsWindows)
+                {
+                    // This sets the execution policy on Windows to Unrestricted which is required to run the user's function scripts on
+                    // Windows client versions. This is needed if a user is testing their function locally with the func CLI.
+                    s_iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Unrestricted;
+                }
             }
 
-            return iss;
+            return PowerShell.Create(s_iss);
         }
 
         /// <summary>

test/Unit/PowerShell/PowerShellManagerTests.cs

@@ -7,13 +7,15 @@
 using System.Collections;
 using System.Collections.Generic;
 using System.IO;
-using System.Management.Automation;
 using Microsoft.Azure.Functions.PowerShellWorker.PowerShell;
+using Microsoft.Azure.Functions.PowerShellWorker.Utility;
 using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
 using Xunit;
 
 namespace Microsoft.Azure.Functions.PowerShellWorker.Test
 {
+    using System.Management.Automation;
+
     internal class TestUtils
     {
         internal const string TestInputBindingName = "req";
@@ -43,9 +45,9 @@ static TestUtils()
 
         // Have a single place to get a PowerShellManager for testing.
         // This is to guarantee that the well known paths are setup before calling the constructor of PowerShellManager.
-        internal static PowerShellManager NewTestPowerShellManager(ConsoleLogger logger)
+        internal static PowerShellManager NewTestPowerShellManager(ConsoleLogger logger, PowerShell pwsh = null)
         {
-            return new PowerShellManager(logger);
+            return pwsh != null ? new PowerShellManager(logger, pwsh) : new PowerShellManager(logger, id: 2);
         }
 
         internal static AzFunctionInfo NewAzFunctionInfo(string scriptFile, string entryPoint)
@@ -235,5 +237,26 @@ public void ProfileWithNonTerminatingError()
             Assert.Equal("Error: ERROR: help me!", _testLogger.FullLog[0]);
             Assert.Matches("Error: Fail to run profile.ps1. See logs for detailed errors. Profile location: ", _testLogger.FullLog[1]);
         }
+
+        [Fact]
+        public void PSManagerCtorRunsProfileByDefault()
+        {
+            //initialize fresh log
+            _testLogger.FullLog.Clear();
+            TestUtils.NewTestPowerShellManager(_testLogger);
+
+            Assert.Single(_testLogger.FullLog);
+            Assert.Equal($"Trace: No 'profile.ps1' is found at the FunctionApp root folder: {FunctionLoader.FunctionAppRootPath}.", _testLogger.FullLog[0]);
+        }
+
+        [Fact]
+        public void PSManagerCtorDoesNotRunProfileIfDelayInit()
+        {
+            //initialize fresh log
+            _testLogger.FullLog.Clear();
+            TestUtils.NewTestPowerShellManager(_testLogger, Utils.NewPwshInstance());
+
+            Assert.Empty(_testLogger.FullLog);
+        }
     }
 }