Trace pipeline output in a streaming way (#111)

In order to trace the pipeline output in a streaming way, the `DataAdding` event won't work, because even though you can get the pipeline object right after it's inserted into the pipeline, you will have to call `Out-String` to get the formatted output. You cannot call `Out-String` within the same Runspace as it's busying running the user's function script, so you will have to create a separate Runspace just for this purpose. This won't work given the concurrency support PS worker needs to have.

The approach I choose is to have a special cmdlet `Trace-PipelineObject` behaves like `Tee-Object`. For each object from the pipeline, `Trace-PipelineObject` pushes the object through a `SteppablePipeline` that consists of `Out-String -Stream | Write-Information -Tags "__PipelineObject__"`, and then writes the object back to pipeline. By using this `SteppablePipeline`, every pipeline object is properly formatted and the resulted strings are written to the information stream.

Author: daxian-dbw

src/Logging/ILogger.cs

@@ -8,8 +8,8 @@
 
 namespace Microsoft.Azure.Functions.PowerShellWorker.Utility
 {
-    internal abstract class ILogger
+    internal interface ILogger
     {
-        public abstract void Log(LogLevel logLevel, string message, Exception exception = null, bool isUserLog = false);
+        void Log(LogLevel logLevel, string message, Exception exception = null, bool isUserLog = false);
     }
 }

src/Logging/RpcLogger.cs

@@ -39,7 +39,7 @@ public void Dispose()
             _invocationId = null;
         }
 
-        public override async void Log(LogLevel logLevel, string message, Exception exception = null, bool isUserLog = false)
+        public async void Log(LogLevel logLevel, string message, Exception exception = null, bool isUserLog = false)
         {
             if (isUserLog)
             {

src/Modules/Microsoft.Azure.Functions.PowerShellWorker/Microsoft.Azure.Functions.PowerShellWorker.psd1

@@ -51,7 +51,7 @@ FormatsToProcess = @()
 NestedModules = @()
 
 # Functions to export from this module, for best performance, do not use wildcards and do not delete the entry, use an empty array if there are no functions to export.
-FunctionsToExport = @('Push-OutputBinding', 'Get-OutputBinding')
+FunctionsToExport = @('Push-OutputBinding', 'Get-OutputBinding', 'Trace-PipelineObject')
 
 # Cmdlets to export from this module, for best performance, do not use wildcards and do not delete the entry, use an empty array if there are no cmdlets to export.
 CmdletsToExport = @()

src/Modules/Microsoft.Azure.Functions.PowerShellWorker/Microsoft.Azure.Functions.PowerShellWorker.psm1

@@ -9,6 +9,10 @@ using namespace Microsoft.Azure.Functions.PowerShellWorker
 
 # This holds the current state of the output bindings.
 $script:_OutputBindings = @{}
+# These variables hold the ScriptBlock and CmdletInfo objects for constructing a SteppablePipeline of 'Out-String | Write-Information'.
+$script:outStringCmd = $ExecutionContext.InvokeCommand.GetCommand("Microsoft.PowerShell.Utility\Out-String", [CommandTypes]::Cmdlet)
+$script:writeInfoCmd = $ExecutionContext.InvokeCommand.GetCommand("Microsoft.PowerShell.Utility\Write-Information", [CommandTypes]::Cmdlet)
+$script:tracingSb = { & $script:outStringCmd -Stream | & $script:writeInfoCmd -Tags "__PipelineObject__" }
 # This loads the resource strings.
 Import-LocalizedData LocalizedData -FileName PowerShellWorker.Resource.psd1
 
@@ -183,3 +187,43 @@ function Push-OutputBinding {
         }
     }
 }
+
+<#
+.SYNOPSIS
+    Writes the formatted output of the pipeline object to the information stream before passing the object down to the pipeline.
+.DESCRIPTION
+    INTERNAL POWERSHELL WORKER USE ONLY. Writes the formatted output of the pipeline object to the information stream before passing the object down to the pipeline.
+.PARAMETER InputObject
+    The object from pipeline.
+#>
+function Trace-PipelineObject {
+
+    [CmdletBinding()]
+    param(
+        [Parameter(Mandatory = $true, ValueFromPipeline = $true)]
+        [object]
+        $InputObject
+    )
+
+    <#
+        This function behaves like 'Tee-Object'.
+        An input pipeline object is first pushed through a steppable pipeline that consists of 'Out-String | Write-Information -Tags "__PipelineObject__"',
+        and then it's written out back to the pipeline without change. In this approach, we can intercept and trace the pipeline objects in a streaming way
+        and keep the objects in pipeline at the same time.
+    #>
+
+    Begin {
+        # A micro-optimization: we use the cached 'CmdletInfo' objects to avoid command resolution every time this cmdlet is called.
+        $stepPipeline = $script:tracingSb.GetSteppablePipeline([CommandOrigin]::Internal)
+        $stepPipeline.Begin($PSCmdlet)
+    }
+
+    Process {
+        $stepPipeline.Process($InputObject)
+        $InputObject
+    }
+
+    End {
+        $stepPipeline.End()
+    }
+}

src/PowerShell/PowerShellExtensions.cs

@@ -37,17 +37,5 @@ public static Collection<T> InvokeAndClearCommands<T>(this PowerShell pwsh)
                 pwsh.Commands.Clear();
             }
         }
-
-        public static string FormatObjectToString(this PowerShell pwsh, object inputObject)
-        {
-            // PowerShell's `Out-String -InputObject` handles collections differently
-            // than when receiving InputObjects from the pipeline. (i.e. `$collection | Out-String`).
-            // That is why we need `Write-Output` here. See related GitHub issue here:
-            // https://github.com/PowerShell/PowerShell/issues/8246
-            return pwsh.AddCommand("Write-Output")
-                .AddParameter("InputObject", inputObject)
-                .AddCommand("Out-String")
-                .InvokeAndClearCommands<string>()[0];
-        }
     }
 }

src/PowerShell/PowerShellManager.cs

@@ -27,7 +27,7 @@ internal class PowerShellManager
         internal PowerShellManager(ILogger logger)
         {
             var initialSessionState = InitialSessionState.CreateDefault();
-            
+
             // Setting the execution policy on macOS and Linux throws an exception so only update it on Windows
             if(Platform.IsWindows)
             {
@@ -161,30 +161,27 @@ internal Hashtable InvokeFunction(
                 Collection<object> pipelineItems = null;
                 using (ExecutionTimer.Start(_logger, "Execution of the user's function completed."))
                 {
-                    pipelineItems = _pwsh.InvokeAndClearCommands<object>();
+                    pipelineItems = _pwsh.AddCommand("Microsoft.Azure.Functions.PowerShellWorker\\Trace-PipelineObject")
+                                         .InvokeAndClearCommands<object>();
                 }
 
                 var result = _pwsh.AddCommand("Microsoft.Azure.Functions.PowerShellWorker\\Get-OutputBinding")
-                                  .AddParameter("Purge")
+                                  .AddParameter("Purge", true)
                                   .InvokeAndClearCommands<Hashtable>()[0];
 
+                /*
+                 * TODO: See GitHub issue #82. We are not settled on how to handle the Azure Functions concept of the $returns Output Binding
                 if (pipelineItems != null && pipelineItems.Count > 0)
                 {
-                    // Log everything we received from the pipeline
-                    foreach (var item in pipelineItems)
-                    {
-                        _logger.Log(LogLevel.Information, $"OUTPUT: {_pwsh.FormatObjectToString(item)}", isUserLog: true);
-                    }
-                    
-                    // TODO: See GitHub issue #82. We are not settled on how to handle the Azure Functions concept of the $returns Output Binding
-                    // If we would like to support Option 1 from #82, uncomment the following code:
-                    // object[] items = new object[pipelineItems.Count];
-                    // pipelineItems.CopyTo(items, 0);
-                    // result.Add(AzFunctionInfo.DollarReturn, items);
-
-                    // If we would like to support Option 2 from #82, uncomment this line:
-                    // result.Add(AzFunctionInfo.DollarReturn, pipelineItems[pipelineItems.Count - 1]);
+                    // If we would like to support Option 1 from #82, use the following 3 lines of code:                    
+                    object[] items = new object[pipelineItems.Count];
+                    pipelineItems.CopyTo(items, 0);
+                    result.Add(AzFunctionInfo.DollarReturn, items);
+
+                    // If we would like to support Option 2 from #82, use this line:
+                    result.Add(AzFunctionInfo.DollarReturn, pipelineItems[pipelineItems.Count - 1]);
                 }
+                */
 
                 return result;
             }

src/PowerShell/StreamHandler.cs

@@ -39,7 +39,8 @@ public void InformationDataAdding(object sender, DataAddingEventArgs e)
         {
             if(e.ItemAdded is InformationRecord record)
             {
-                _logger.Log(LogLevel.Information, $"INFORMATION: {record.MessageData}", isUserLog: true);
+                string prefix = (record.Tags.Count == 1 && record.Tags[0] == "__PipelineObject__") ? "OUTPUT:" : "INFORMATION:";
+                _logger.Log(LogLevel.Information, $"{prefix} {record.MessageData}", isUserLog: true);
             }
         }
 

test/Unit/Logging/ConsoleLogger.cs

@@ -15,7 +15,7 @@ internal class ConsoleLogger : ILogger
     {
         public List<string> FullLog = new List<string>();
 
-        public override void Log(LogLevel logLevel, string message, Exception exception = null, bool isUserLog = false)
+        public void Log(LogLevel logLevel, string message, Exception exception = null, bool isUserLog = false)
         {
             var log = $"{logLevel}: {message}";
             Console.WriteLine(log);

test/Unit/Modules/Microsoft.Azure.Functions.PowerShellWorker.Tests.ps1

@@ -154,4 +154,60 @@ Describe 'Azure Functions PowerShell Langauge Worker Helper Module Tests' {
                 -Because 'The OutputBindings should be empty'
         }
     }
+
+    Context 'Trace-PipelineObject tests' {
+        BeforeAll {
+            $scriptToRun = @'
+    param($cmd, $modulePath)
+    Import-Module $modulePath
+    function Write-TestObject {
+        foreach ($i in 1..20) {
+            Write-Output $cmd
+        }
+        Write-Information '__LAST_INFO_MSG__'
+    }
+'@
+            $cmd = Get-Command Get-Command
+            $ps = [powershell]::Create()
+            $ps.AddScript($scriptToRun).AddParameter("cmd", $cmd).AddParameter("modulePath", $modulePath).Invoke()
+            $ps.Commands.Clear()
+            $ps.Streams.ClearStreams()
+
+            function Write-TestObject {
+                foreach ($i in 1..20) {
+                    Write-Output $cmd
+                }
+                Write-Information '__LAST_INFO_MSG__'
+            }
+        }
+
+        AfterAll {
+            $ps.Dispose()
+        }
+
+        AfterEach {
+            $ps.Commands.Clear()
+            $ps.Streams.ClearStreams()
+        }
+
+        It "Can write tracing to information stream while keeps input object in pipeline" {
+            $results = $ps.AddCommand("Write-TestObject").AddCommand("Trace-PipelineObject").Invoke()
+
+            $results.Count | Should -BeExactly 20
+            for ($i = 0; $i -lt 20; $i++) {
+                $results[0].Name | Should -BeExactly $cmd.Name
+            }
+
+            $outStringResults = Write-TestObject | Out-String -Stream
+            $ps.Streams.Information.Count | Should -BeExactly ($outStringResults.Count + 1)
+            $countWithoutTrailingNewLines = $outStringResults.Count - 2
+            for ($i = 0; $i -lt $countWithoutTrailingNewLines; $i++) {
+                $ps.Streams.Information[$i].MessageData | Should -BeExactly $outStringResults[$i]
+                $ps.Streams.Information[$i].Tags | Should -BeExactly "__PipelineObject__"
+            }
+
+            $ps.Streams.Information[$i].MessageData | Should -BeExactly "__LAST_INFO_MSG__"
+            $ps.Streams.Information[$i].Tags | Should -BeNullOrEmpty
+        }
+    }
 }

test/Unit/PowerShell/PowerShellExtensionsTests.cs

@@ -38,24 +38,5 @@ public void CanInvokeAndClearCommandsWithReturnValue()
 
             Assert.Single(results, data);
         }
-
-        [Fact]
-        public void CanFormatObjectsToString()
-        {
-            string[] expectedResult = {
-                "",
-                "Name                           Value",
-                "----                           -----",
-                "Foo                            bar",
-                "",
-                "",
-                ""
-            };
-
-            Hashtable data = new Hashtable {{ "Foo", "bar"}};   
-            Assert.Equal(
-                string.Join(Environment.NewLine, expectedResult), 
-                PowerShell.Create().FormatObjectToString(data));
-        }
     }
 }