Warm up PowerShell Worker before receiving function load request (#535)

Author: AnatoliB

src/RequestProcessor.cs

@@ -19,12 +19,12 @@
 namespace Microsoft.Azure.Functions.PowerShellWorker
 {
     using System.Diagnostics;
-    using System.IO;
     using LogLevel = Microsoft.Azure.WebJobs.Script.Grpc.Messages.RpcLog.Types.Level;
 
     internal class RequestProcessor
     {
         private readonly MessagingStream _msgStream;
+        private readonly System.Management.Automation.PowerShell _firstPwshInstance;
         private readonly PowerShellManagerPool _powershellPool;
         private DependencyManager _dependencyManager;
 
@@ -37,9 +37,10 @@ internal class RequestProcessor
         private Dictionary<StreamingMessage.ContentOneofCase, Func<StreamingMessage, StreamingMessage>> _requestHandlers =
             new Dictionary<StreamingMessage.ContentOneofCase, Func<StreamingMessage, StreamingMessage>>();
 
-        internal RequestProcessor(MessagingStream msgStream)
+        internal RequestProcessor(MessagingStream msgStream, System.Management.Automation.PowerShell firstPwshInstance)
         {
             _msgStream = msgStream;
+            _firstPwshInstance = firstPwshInstance;
             _powershellPool = new PowerShellManagerPool(() => new RpcLogger(msgStream));
 
             // Host sends capabilities/init data to worker
@@ -194,18 +195,12 @@ internal StreamingMessage ProcessFunctionLoadRequest(StreamingMessage request)
                     _dependencyManager = new DependencyManager(request.FunctionLoadRequest.Metadata.Directory, logger: rpcLogger);
                     var managedDependenciesPath = _dependencyManager.Initialize(request, rpcLogger);
 
-                    // Setup the FunctionApp root path and module path.
-                    FunctionLoader.SetupWellKnownPaths(functionLoadRequest, managedDependenciesPath);
+                    SetupAppRootPathAndModulePath(functionLoadRequest, managedDependenciesPath);
 
-                    // 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();
-                    LogPowerShellVersion(rpcLogger, pwsh);
-                    _powershellPool.Initialize(pwsh);
+                    _powershellPool.Initialize(_firstPwshInstance);
 
                     // Start the download asynchronously if needed.
-                    _dependencyManager.StartDependencyInstallationIfNeeded(request, pwsh, rpcLogger);
+                    _dependencyManager.StartDependencyInstallationIfNeeded(request, _firstPwshInstance, rpcLogger);
 
                     rpcLogger.Log(isUserOnlyLog: false, LogLevel.Trace, string.Format(PowerShellWorkerStrings.FirstFunctionLoadCompleted, stopwatch.ElapsedMilliseconds));
                 }
@@ -493,10 +488,19 @@ private static void BindOutputFromResult(InvocationResponse response, AzFunction
             }
         }
 
-        private static void LogPowerShellVersion(RpcLogger rpcLogger, System.Management.Automation.PowerShell pwsh)
+        private void SetupAppRootPathAndModulePath(FunctionLoadRequest functionLoadRequest, string managedDependenciesPath)
         {
-            var message = string.Format(PowerShellWorkerStrings.PowerShellVersion, Utils.GetPowerShellVersion(pwsh));
-            rpcLogger.Log(isUserOnlyLog: false, LogLevel.Information, message);
+            FunctionLoader.SetupWellKnownPaths(functionLoadRequest, managedDependenciesPath);
+
+            if (FunctionLoader.FunctionAppRootPath == null)
+            {
+                throw new InvalidOperationException(PowerShellWorkerStrings.FunctionAppRootNotResolved);
+            }
+
+            _firstPwshInstance.AddCommand("Microsoft.PowerShell.Management\\Set-Content")
+                .AddParameter("Path", "env:PSModulePath")
+                .AddParameter("Value", FunctionLoader.FunctionModulePath)
+                .InvokeAndClearCommands();
         }
 
         #endregion

src/Utility/Utils.cs

@@ -37,11 +37,6 @@ internal static PowerShell NewPwshInstance()
         {
             if (s_iss == null)
             {
-                if (FunctionLoader.FunctionAppRootPath == null)
-                {
-                    throw new InvalidOperationException(PowerShellWorkerStrings.FunctionAppRootNotResolved);
-                }
-
                 s_iss = InitialSessionState.CreateDefault();
 
                 if (!AreDurableFunctionsEnabled())
@@ -51,11 +46,14 @@ internal static PowerShell NewPwshInstance()
                     s_iss.ThreadOptions = PSThreadOptions.UseCurrentThread;
                 }
 
-                s_iss.EnvironmentVariables.Add(
-                    new SessionStateVariableEntry(
-                        "PSModulePath",
-                        FunctionLoader.FunctionModulePath,
-                        description: null));
+                if (FunctionLoader.FunctionAppRootPath != null)
+                {
+                    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)

src/Worker.cs

@@ -4,10 +4,12 @@
 //
 
 using System;
+using System.Management.Automation;
 using System.Threading.Tasks;
 
 using CommandLine;
 using Microsoft.Azure.Functions.PowerShellWorker.Messaging;
+using Microsoft.Azure.Functions.PowerShellWorker.PowerShell;
 using Microsoft.Azure.Functions.PowerShellWorker.Utility;
 using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
 
@@ -34,18 +36,50 @@ public async static Task Main(string[] args)
                 .WithParsed(ops => arguments = ops)
                 .WithNotParsed(err => Environment.Exit(1));
 
+            // 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 firstPowerShellInstance = Utils.NewPwshInstance();
+            LogPowerShellVersion(firstPowerShellInstance);
+            WarmUpPowerShell(firstPowerShellInstance);
+
             var msgStream = new MessagingStream(arguments.Host, arguments.Port);
-            var requestProcessor = new RequestProcessor(msgStream);
+            var requestProcessor = new RequestProcessor(msgStream, firstPowerShellInstance);
 
             // Send StartStream message
-            var startedMessage = new StreamingMessage() {
+            var startedMessage = new StreamingMessage()
+            {
                 RequestId = arguments.RequestId,
                 StartStream = new StartStream() { WorkerId = arguments.WorkerId }
             };
 
             msgStream.Write(startedMessage);
             await requestProcessor.ProcessRequestLoop();
         }
+
+        // Warm up the PowerShell instance so that the subsequent function load and invocation requests are faster
+        private static void WarmUpPowerShell(System.Management.Automation.PowerShell firstPowerShellInstance)
+        {
+            // It turns out that creating/removing a function warms up the runspace enough.
+            // We just need this name to be unique, so that it does not coincide with an actual function.
+            const string DummyFunctionName = "DummyFunction-71b09c92-6bce-42d0-aba1-7b985b8c3563";
+
+            firstPowerShellInstance.AddCommand("Microsoft.PowerShell.Management\\New-Item")
+                .AddParameter("Path", "Function:")
+                .AddParameter("Name", DummyFunctionName)
+                .AddParameter("Value", ScriptBlock.Create(string.Empty))
+                .InvokeAndClearCommands();
+
+            firstPowerShellInstance.AddCommand("Microsoft.PowerShell.Management\\Remove-Item")
+                .AddParameter("Path", $"Function:{DummyFunctionName}")
+                .InvokeAndClearCommands();
+        }
+
+        private static void LogPowerShellVersion(System.Management.Automation.PowerShell pwsh)
+        {
+            var message = string.Format(PowerShellWorkerStrings.PowerShellVersion, Utils.GetPowerShellVersion(pwsh));
+            RpcLogger.WriteSystemLog(LogLevel.Information, message);
+        }
     }
 
     internal class WorkerArguments