Fix writing to streams from profile.ps1 (#411)

AnatoliB · AnatoliB · commit 983a12c590ec · 2020-04-28T16:18:22.000-07:00
* Refactor PowerShellManagerPool to make it more testable

* Store _createLogger instead of _msgStream in PowerShellManagerPool
* Pass Func&lt;ILogger&gt; to PowerShellManagerPool ctor
* Pass requestId and invocationId instead of StreamingMessage to CheckoutIdleWorker
* Pass functionName to CheckoutIdleWorker
* Pass outputBindings instead of AzFunctionInfo to RegisterFunctionMetadata
* Pass outputBindings to CheckoutIdleWorker
* Don't pass AzFunctionInfo to CheckoutIdleWorker

* Add LoggerContextIsSet test

* Set invocation context on the logger before invoking profile.ps1

* Improve the message on errors reported by profile

* Get rid of 'Import-Module profile' output
diff --git a/src/PowerShell/PowerShellManager.cs b/src/PowerShell/PowerShellManager.cs
@@ -162,6 +162,8 @@ internal void InvokeProfile(string profilePath)
                 return;
             }
 
+            var profileExecutionHadErrors = false;
+
             try
             {
                 var stopwatch = new Stopwatch();
@@ -170,8 +172,13 @@ internal void InvokeProfile(string profilePath)
                 // Import-Module on a .ps1 file will evaluate the script in the global scope.
                 _pwsh.AddCommand(Utils.ImportModuleCmdletInfo)
                         .AddParameter("Name", profilePath)
-                        .AddParameter("PassThru", Utils.BoxedTrue)
-                     .AddCommand(Utils.RemoveModuleCmdletInfo)
+                     .AddCommand("Microsoft.Azure.Functions.PowerShellWorker\\Trace-PipelineObject")
+                     .InvokeAndClearCommands();
+
+                profileExecutionHadErrors = _pwsh.HadErrors;
+
+                _pwsh.AddCommand(Utils.RemoveModuleCmdletInfo)
+                        .AddParameter("FullyQualifiedName", profilePath)
                         .AddParameter("Force", Utils.BoxedTrue)
                         .AddParameter("ErrorAction", "SilentlyContinue")
                      .InvokeAndClearCommands();
@@ -185,9 +192,9 @@ internal void InvokeProfile(string profilePath)
             }
             finally
             {
-                if (_pwsh.HadErrors)
+                if (profileExecutionHadErrors || _pwsh.HadErrors)
                 {
-                    string errorMsg = string.Format(PowerShellWorkerStrings.FailToRunProfile, profilePath);
+                    string errorMsg = string.Format(PowerShellWorkerStrings.ErrorsWhileExecutingProfile, profilePath);
                     Logger.Log(isUserOnlyLog: true, LogLevel.Error, errorMsg, exception);
                 }
             }
diff --git a/src/PowerShell/PowerShellManagerPool.cs b/src/PowerShell/PowerShellManagerPool.cs
@@ -3,11 +3,11 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
 using System.Collections.Concurrent;
+using System.Collections.ObjectModel;
 using System.Threading;
-using Microsoft.Azure.Functions.PowerShellWorker.Messaging;
 using Microsoft.Azure.Functions.PowerShellWorker.Utility;
-using Microsoft.Azure.WebJobs.Script.Grpc.Messages;
 
 namespace Microsoft.Azure.Functions.PowerShellWorker.PowerShell
 {
@@ -19,7 +19,7 @@ namespace Microsoft.Azure.Functions.PowerShellWorker.PowerShell
     /// </summary>
     internal class PowerShellManagerPool
     {
-        private readonly MessagingStream _msgStream;
+        private readonly Func<ILogger> _createLogger;
         private readonly BlockingCollection<PowerShellManager> _pool;
         private int _poolSize;
 
@@ -31,9 +31,9 @@ internal class PowerShellManagerPool
         /// <summary>
         /// Constructor of the pool.
         /// </summary>
-        internal PowerShellManagerPool(MessagingStream msgStream)
+        internal PowerShellManagerPool(Func<ILogger> createLogger)
         {
-            _msgStream = msgStream;
+            _createLogger = createLogger;
             _pool = new BlockingCollection<PowerShellManager>(UpperBound);
             RpcLogger.WriteSystemLog(LogLevel.Information, string.Format(PowerShellWorkerStrings.LogConcurrencyUpperBound, UpperBound.ToString()));
         }
@@ -44,7 +44,7 @@ internal PowerShellManagerPool(MessagingStream msgStream)
         /// </summary>
         internal void Initialize(PowerShell pwsh)
         {
-            var logger = new RpcLogger(_msgStream);
+            var logger = _createLogger();
             var psManager = new PowerShellManager(logger, pwsh);
             _pool.Add(psManager);
             _poolSize = 1;
@@ -53,11 +53,13 @@ internal void Initialize(PowerShell pwsh)
         /// <summary>
         /// Checkout an idle PowerShellManager instance in a non-blocking asynchronous way.
         /// </summary>
-        internal PowerShellManager CheckoutIdleWorker(StreamingMessage request, AzFunctionInfo functionInfo)
+        internal PowerShellManager CheckoutIdleWorker(
+            string requestId,
+            string invocationId,
+            string functionName,
+            ReadOnlyDictionary<string, ReadOnlyBindingInfo> outputBindings)
         {
             PowerShellManager psManager = null;
-            string requestId = request.RequestId;
-            string invocationId = request.InvocationRequest?.InvocationId;
 
             // If the pool has an idle one, just use it.
             if (!_pool.TryTake(out psManager))
@@ -80,21 +82,22 @@ internal PowerShellManager CheckoutIdleWorker(StreamingMessage request, AzFuncti
                 if (psManager == null)
                 {
                     var logger = CreateLoggerWithContext(requestId, invocationId);
-                    logger.Log(isUserOnlyLog: true, LogLevel.Warning, string.Format(PowerShellWorkerStrings.FunctionQueuingRequest, functionInfo.FuncName));
+                    logger.Log(isUserOnlyLog: true, LogLevel.Warning, string.Format(PowerShellWorkerStrings.FunctionQueuingRequest, functionName));
 
                     // If the pool has reached its bounded capacity, then the thread
                     // should be blocked until an idle one becomes available.
                     psManager = _pool.Take();
                 }
             }
 
+            psManager.Logger.SetContext(requestId, invocationId);
+
             // 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);
+            FunctionMetadata.RegisterFunctionMetadata(psManager.InstanceId, outputBindings);
 
             return psManager;
         }
@@ -113,9 +116,9 @@ internal void ReclaimUsedWorker(PowerShellManager psManager)
             }
         }
 
-        private RpcLogger CreateLoggerWithContext(string requestId, string invocationId)
+        private ILogger CreateLoggerWithContext(string requestId, string invocationId)
         {
-            var logger = new RpcLogger(_msgStream);
+            var logger = _createLogger();
             logger.SetContext(requestId, invocationId);
             return logger;
         }
diff --git a/src/Public/FunctionMetadata.cs b/src/Public/FunctionMetadata.cs
@@ -42,9 +42,8 @@ internal static Hashtable GetOutputBindingHashtable(Guid runspaceInstanceId)
         /// <summary>
         /// Helper method to set the output binding metadata for the function that is about to run.
         /// </summary>
-        internal static void RegisterFunctionMetadata(Guid instanceId, AzFunctionInfo functionInfo)
+        internal static void RegisterFunctionMetadata(Guid instanceId, ReadOnlyDictionary<string, ReadOnlyBindingInfo> outputBindings)
         {
-            var outputBindings = functionInfo.OutputBindings;
             OutputBindingCache.AddOrUpdate(instanceId, outputBindings, (key, value) => outputBindings);
         }
 
diff --git a/src/RequestProcessor.cs b/src/RequestProcessor.cs
@@ -40,7 +40,7 @@ internal class RequestProcessor
         internal RequestProcessor(MessagingStream msgStream)
         {
             _msgStream = msgStream;
-            _powershellPool = new PowerShellManagerPool(msgStream);
+            _powershellPool = new PowerShellManagerPool(() => new RpcLogger(msgStream));
 
             // Host sends capabilities/init data to worker
             _requestHandlers.Add(StreamingMessage.ContentOneofCase.WorkerInitRequest, ProcessWorkerInitRequest);
@@ -255,7 +255,13 @@ internal StreamingMessage ProcessInvocationRequest(StreamingMessage request)
                 stopwatch.OnCheckpoint(FunctionInvocationPerformanceStopwatch.Checkpoint.DependenciesAvailable);
 
                 AzFunctionInfo functionInfo = FunctionLoader.GetFunctionInfo(request.InvocationRequest.FunctionId);
-                PowerShellManager psManager = _powershellPool.CheckoutIdleWorker(request, functionInfo);
+
+                PowerShellManager psManager = _powershellPool.CheckoutIdleWorker(
+                                                request.RequestId,
+                                                request.InvocationRequest?.InvocationId,
+                                                functionInfo.FuncName,
+                                                functionInfo.OutputBindings);
+
                 stopwatch.OnCheckpoint(FunctionInvocationPerformanceStopwatch.Checkpoint.RunspaceAvailable);
 
                 // When the concurrency upper bound is more than 1, we have to handle the invocation in a worker
diff --git a/src/resources/PowerShellWorkerStrings.resx b/src/resources/PowerShellWorkerStrings.resx
@@ -154,8 +154,8 @@
   <data name="FileNotFound" xml:space="preserve">
     <value>No '{0}' is found at the function app root folder: {1}.</value>
   </data>
-  <data name="FailToRunProfile" xml:space="preserve">
-    <value>Failed to run profile.ps1. See logs for detailed errors. Profile location: {0}.</value>
+  <data name="ErrorsWhileExecutingProfile" xml:space="preserve">
+    <value>Errors reported while executing profile.ps1. See logs for detailed errors. Profile location: {0}.</value>
   </data>
   <data name="UnsupportedMessage" xml:space="preserve">
     <value>Unsupported message type: {0}.</value>
diff --git a/test/Unit/Modules/HelperModuleTests.cs b/test/Unit/Modules/HelperModuleTests.cs
@@ -56,7 +56,7 @@ static HelperModuleTests()
 
         public HelperModuleTests()
         {
-            FunctionMetadata.RegisterFunctionMetadata(s_pwsh.Runspace.InstanceId, s_funcInfo);
+            FunctionMetadata.RegisterFunctionMetadata(s_pwsh.Runspace.InstanceId, s_funcInfo.OutputBindings);
         }
 
         public void Dispose()
diff --git a/test/Unit/PowerShell/PowerShellManagerTests.cs b/test/Unit/PowerShell/PowerShellManagerTests.cs
@@ -105,7 +105,7 @@ public void InvokeBasicFunctionWorks()
 
             try
             {
-                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo);
+                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo.OutputBindings);
                 Hashtable result = InvokeFunction(testManager, functionInfo);
 
                 // The outputBinding hashtable for the runspace should be cleared after 'InvokeFunction'
@@ -130,7 +130,7 @@ public void InvokeFunctionWithSpecialVariableWorks()
 
             try
             {
-                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo);
+                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo.OutputBindings);
                 Hashtable result = InvokeFunction(testManager, functionInfo);
 
                 // The outputBinding hashtable for the runspace should be cleared after 'InvokeFunction'
@@ -155,7 +155,7 @@ public void InvokeBasicFunctionWithRequiresWorks()
 
             try
             {
-                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo);
+                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo.OutputBindings);
                 Hashtable result = InvokeFunction(testManager, functionInfo);
 
                 // The outputBinding hashtable for the runspace should be cleared after 'InvokeFunction'
@@ -186,7 +186,7 @@ public void InvokeBasicFunctionWithTriggerMetadataAndTraceContextWorks()
 
             try
             {
-                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo);
+                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo.OutputBindings);
 
                 Hashtable result = InvokeFunction(testManager, functionInfo, triggerMetadata);
 
@@ -212,7 +212,7 @@ public void InvokeFunctionWithEntryPointWorks()
 
             try
             {
-                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo);
+                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo.OutputBindings);
                 Hashtable result = InvokeFunction(testManager, functionInfo);
 
                 // The outputBinding hashtable for the runspace should be cleared after 'InvokeFunction'
@@ -236,7 +236,7 @@ public void FunctionShouldCleanupVariableTable()
 
             try
             {
-                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo);
+                FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo.OutputBindings);
 
                 Hashtable result1 = InvokeFunction(testManager, functionInfo);
                 Assert.Equal("is not set", result1[TestOutputBindingName]);
@@ -266,7 +266,7 @@ public void RegisterAndUnregisterFunctionMetadataShouldWork()
             string path = Path.Join(s_funcDirectory, "testBasicFunction.ps1");
             var (functionInfo, testManager) = PrepareFunction(path, string.Empty);
 
-            FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo);
+            FunctionMetadata.RegisterFunctionMetadata(testManager.InstanceId, functionInfo.OutputBindings);
             var outBindingMap = FunctionMetadata.GetOutputBindingInfo(testManager.InstanceId);
             Assert.Single(outBindingMap);
             Assert.Equal(TestOutputBindingName, outBindingMap.First().Key);
@@ -317,7 +317,7 @@ public void ProfileWithTerminatingError()
             Assert.Throws<CmdletInvocationException>(() => testManager.InvokeProfile(profilePath));
             var relevantLogs = s_testLogger.FullLog.Where(message => !message.StartsWith("Trace:")).ToList();
             Assert.Single(relevantLogs);
-            Assert.Matches("Error: Failed to run profile.ps1. See logs for detailed errors. Profile location: ", relevantLogs[0]);
+            Assert.Matches("Error: Errors reported while executing profile.ps1. See logs for detailed errors. Profile location: ", relevantLogs[0]);
         }
 
         [Fact]
@@ -334,7 +334,7 @@ public void ProfileWithNonTerminatingError()
             Assert.Equal(2, relevantLogs.Count);
             Assert.StartsWith("Error: ERROR: ", relevantLogs[0]);
             Assert.Contains("help me!", relevantLogs[0]);
-            Assert.Matches("Error: Failed to run profile.ps1. See logs for detailed errors. Profile location: ", relevantLogs[1]);
+            Assert.Matches("Error: Errors reported while executing profile.ps1. See logs for detailed errors. Profile location: ", relevantLogs[1]);
         }
 
         [Fact]
@@ -358,9 +358,44 @@ public void PSManagerCtorDoesNotRunProfileIfDelayInit()
             Assert.Empty(s_testLogger.FullLog);
         }
 
+        [Fact]
+        public void LoggerContextIsSet()
+        {
+            var dummyBindingInfo = new Dictionary<string, ReadOnlyBindingInfo>();
+            var outputBindings = new ReadOnlyDictionary<string, ReadOnlyBindingInfo>(dummyBindingInfo);
+
+            var powerShellManagerPool = new PowerShellManagerPool(() => new ContextValidatingLogger());
+            var pwsh = Utils.NewPwshInstance();
+            powerShellManagerPool.Initialize(pwsh);
+
+            var worker = powerShellManagerPool.CheckoutIdleWorker("requestId", "invocationId", "FuncName", outputBindings);
+
+            powerShellManagerPool.ReclaimUsedWorker(worker);
+        }
+
         private static Hashtable InvokeFunction(PowerShellManager powerShellManager, AzFunctionInfo functionInfo, Hashtable triggerMetadata = null)
         {
             return powerShellManager.InvokeFunction(functionInfo, triggerMetadata, null, s_testInputData, new FunctionInvocationPerformanceStopwatch());
         }
+
+        private class ContextValidatingLogger : ILogger
+        {
+            private bool _isContextSet = false;
+
+            public void Log(bool isUserOnlyLog, RpcLog.Types.Level logLevel, string message, Exception exception = null)
+            {
+                Assert.True(_isContextSet);
+            }
+
+            public void SetContext(string requestId, string invocationId)
+            {
+                _isContextSet = true;
+            }
+
+            public void ResetContext()
+            {
+                _isContextSet = false;
+            }
+        }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -162,6 +162,8 @@ internal void InvokeProfile(string profilePath)`
`162`	`162`	`return;`
`163`	`163`	`}`
`164`	`164`
	`165`	`+ var profileExecutionHadErrors = false;`
	`166`	`+`
`165`	`167`	`try`
`166`	`168`	`{`
`167`	`169`	`var stopwatch = new Stopwatch();`
`@@ -170,8 +172,13 @@ internal void InvokeProfile(string profilePath)`
`170`	`172`	`// Import-Module on a .ps1 file will evaluate the script in the global scope.`
`171`	`173`	`_pwsh.AddCommand(Utils.ImportModuleCmdletInfo)`
`172`	`174`	`.AddParameter("Name", profilePath)`
`173`		`- .AddParameter("PassThru", Utils.BoxedTrue)`
`174`		`- .AddCommand(Utils.RemoveModuleCmdletInfo)`
	`175`	`+ .AddCommand("Microsoft.Azure.Functions.PowerShellWorker\\Trace-PipelineObject")`
	`176`	`+ .InvokeAndClearCommands();`
	`177`	`+`
	`178`	`+ profileExecutionHadErrors = _pwsh.HadErrors;`
	`179`	`+`
	`180`	`+ _pwsh.AddCommand(Utils.RemoveModuleCmdletInfo)`
	`181`	`+ .AddParameter("FullyQualifiedName", profilePath)`
`175`	`182`	`.AddParameter("Force", Utils.BoxedTrue)`
`176`	`183`	`.AddParameter("ErrorAction", "SilentlyContinue")`
`177`	`184`	`.InvokeAndClearCommands();`
`@@ -185,9 +192,9 @@ internal void InvokeProfile(string profilePath)`
`185`	`192`	`}`
`186`	`193`	`finally`
`187`	`194`	`{`
`188`		`- if (_pwsh.HadErrors)`
	`195`	`+ if (profileExecutionHadErrors \|\| _pwsh.HadErrors)`
`189`	`196`	`{`
`190`		`- string errorMsg = string.Format(PowerShellWorkerStrings.FailToRunProfile, profilePath);`
	`197`	`+ string errorMsg = string.Format(PowerShellWorkerStrings.ErrorsWhileExecutingProfile, profilePath);`
`191`	`198`	`Logger.Log(isUserOnlyLog: true, LogLevel.Error, errorMsg, exception);`
`192`	`199`	`}`
`193`	`200`	`}`
Original file line number	Diff line number	Diff line change
`@@ -42,9 +42,8 @@ internal static Hashtable GetOutputBindingHashtable(Guid runspaceInstanceId)`
`42`	`42`	`/// <summary>`
`43`	`43`	`/// Helper method to set the output binding metadata for the function that is about to run.`
`44`	`44`	`/// </summary>`
`45`		`- internal static void RegisterFunctionMetadata(Guid instanceId, AzFunctionInfo functionInfo)`
	`45`	`+ internal static void RegisterFunctionMetadata(Guid instanceId, ReadOnlyDictionary<string, ReadOnlyBindingInfo> outputBindings)`
`46`	`46`	`{`
`47`		`- var outputBindings = functionInfo.OutputBindings;`
`48`	`47`	`OutputBindingCache.AddOrUpdate(instanceId, outputBindings, (key, value) => outputBindings);`
`49`	`48`	`}`
`50`	`49`
Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@ static HelperModuleTests()`
`56`	`56`
`57`	`57`	`public HelperModuleTests()`
`58`	`58`	`{`
`59`		`- FunctionMetadata.RegisterFunctionMetadata(s_pwsh.Runspace.InstanceId, s_funcInfo);`
	`59`	`+ FunctionMetadata.RegisterFunctionMetadata(s_pwsh.Runspace.InstanceId, s_funcInfo.OutputBindings);`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`public void Dispose()`