PowerShell · rjmholt · Oct 7, 2021 · Oct 5, 2021 · Oct 6, 2021 · Oct 6, 2021
diff --git a/src/System.Management.Automation/engine/MshCommandRuntime.cs b/src/System.Management.Automation/engine/MshCommandRuntime.cs
@@ -928,7 +928,8 @@ private bool InitShouldLogPipelineExecutionDetail()
         /// </summary>
         internal string PipelineVariable { get; set; }
 
-        private PSVariable _pipelineVarReference = null;
+        private PSVariable _pipelineVarReference;
+        private bool _shouldRemovePipelineVariable;
 
         internal void SetupOutVariable()
         {
@@ -972,7 +973,7 @@ internal void SetupPipelineVariable()
             // This can't use the common SetupVariable implementation, as this needs to persist for an entire
             // pipeline.
 
-            if (string.IsNullOrEmpty(this.PipelineVariable))
+            if (string.IsNullOrEmpty(PipelineVariable))
             {
                 return;
             }
@@ -983,19 +984,40 @@ internal void SetupPipelineVariable()
                 _state = new SessionState(Context.EngineSessionState);
 
             // Create the pipeline variable
-            _pipelineVarReference = new PSVariable(this.PipelineVariable);
-            _state.PSVariable.Set(_pipelineVarReference);
+            _pipelineVarReference = new PSVariable(PipelineVariable);
+            object varToUse = _state.Internal.SetVariable(
+                _pipelineVarReference,
+                force: false,
+                CommandOrigin.Internal);
 
-            // Get the reference again in case we re-used one from the
-            // same scope.
-            _pipelineVarReference = _state.PSVariable.Get(this.PipelineVariable);
+            if (ReferenceEquals(_pipelineVarReference, varToUse))
+            {
+                // The returned variable is the exact same instance, which means we set a new variable.
+                // In this case, we will try removing the pipeline variable in the end.
+                _shouldRemovePipelineVariable = true;
+            }
+            else
+            {
+                // A variable with the same name already exists in the same scope and it was returned.
+                // In this case, we update the reference and don't remove the variable in the end.
+                _pipelineVarReference = (PSVariable)varToUse;
+            }
 
             if (_thisCommand is not PSScriptCmdlet)
             {
                 this.OutputPipe.SetPipelineVariable(_pipelineVarReference);
             }
         }
 
+        internal void RemovePipelineVariable()
+        {
+            if (_shouldRemovePipelineVariable)
+            {
+                // Remove pipeline variable when a pipeline is being torn down.
+                _state.PSVariable.Remove(PipelineVariable);
+            }
+        }
+
         /// <summary>
         /// Configures the number of objects to buffer before calling the downstream Cmdlet.
         /// </summary>
@@ -3765,25 +3787,12 @@ internal void SetVariableListsInPipe()
 
             if (this.PipelineVariable != null)
             {
-                // _state can be null if the current script block is dynamicparam, etc.
-                if (_state != null)
-                {
-                    // Create the pipeline variable
-                    _state.PSVariable.Set(_pipelineVarReference);
-
-                    // Get the reference again in case we re-used one from the
-                    // same scope.
-                    _pipelineVarReference = _state.PSVariable.Get(this.PipelineVariable);
-                }
-
                 this.OutputPipe.SetPipelineVariable(_pipelineVarReference);
             }
         }
 
         internal void RemoveVariableListsInPipe()
         {
-            // Diagnostics.Assert(thisCommand is PSScriptCmdlet, "this is only done for script cmdlets");
-
             if (_outVarList != null)
             {
                 this.OutputPipe.RemoveVariableList(VariableStreamKind.Output, _outVarList);
@@ -3807,9 +3816,6 @@ internal void RemoveVariableListsInPipe()
             if (this.PipelineVariable != null)
             {
                 this.OutputPipe.RemovePipelineVariable();
-                // '_state' could be null when a 'DynamicParam' block runs because the 'DynamicParam' block runs in 'DoPrepare',
-                // before 'PipelineProcessor.SetupParameterVariables' is called, where '_state' is initialized.
-                _state?.PSVariable.Remove(this.PipelineVariable);
             }
         }
     }

diff --git a/src/System.Management.Automation/engine/pipeline.cs b/src/System.Management.Automation/engine/pipeline.cs
@@ -1298,7 +1298,21 @@ private void DisposeCommands()
                         // pipeline failure and continue disposing cmdlets.
                         try
                         {
-                            commandProcessor.CommandRuntime.RemoveVariableListsInPipe();
+                            // Only cmdlets can have variables defined via the common parameters.
+                            // We handle the cleanup of those variables only if we need to.
+                            if (commandProcessor is CommandProcessor)
+                            {
+                                if (commandProcessor.Command is not PSScriptCmdlet)
+                                {
+                                    // For script cmdlets, the variable lists were already removed when exiting a scope.
+                                    // So we only need to take care of binary cmdlets here.
+                                    commandProcessor.CommandRuntime.RemoveVariableListsInPipe();
+                                }
+
+                                // Remove the pipeline variable if we need to.
+                                commandProcessor.CommandRuntime.RemovePipelineVariable();
+                            }
+
                             commandProcessor.Dispose();
                         }
                         // 2005/04/13-JonN: The only vaguely plausible reason

diff --git a/test/powershell/engine/ParameterBinding/ParameterBinding.Tests.ps1 b/test/powershell/engine/ParameterBinding/ParameterBinding.Tests.ps1
@@ -264,6 +264,40 @@ Describe "Parameter Binding Tests" -Tags "CI" {
         { Copy-Item "~\$guid*" -Destination ~ -ToSession $null } | Should -Throw -ErrorId 'ParameterArgumentValidationError'
     }
 
+    It 'PipelineVariable should not cause variable-removal exception (issue #16155)' {
+        function Invoke-AddOne {
+            param (
+                [Parameter(ValueFromPipeline)]
+                [int]$Number
+            )
+
+            Begin {
+                $testValue = 'prefix-'
+            }
+
+            Process {
+                $testValue + $Number
+            }
+        }
+
+        1,2,3 | Invoke-AddOne -PipelineVariable testValue | ForEach-Object { $testValue } | Should -Be @('prefix-1', 'prefix-2', 'prefix-3')
+        { Get-Variable -Name testValue -ErrorAction Stop } | Should -Throw -ErrorId 'VariableNotFound,Microsoft.PowerShell.Commands.GetVariableCommand'
+
+        $results = & { $test = 'str'; 1,2,3 | Invoke-AddOne -PipelineVariable test | ForEach-Object { $test }; Get-Variable test }
+        $results | Should -HaveCount 4
+        $results[0] | Should -BeExactly 'prefix-1'
+        $results[1] | Should -BeExactly 'prefix-2'
+        $results[2] | Should -BeExactly 'prefix-3'
+        $results[3] -is [psvariable] | Should -BeTrue
+
+        $results = & { Set-Variable -Name test -Value 'str'; 1,2,3 | Invoke-AddOne -PipelineVariable test | ForEach-Object { $test }; Get-Variable test }
+        $results | Should -HaveCount 4
+        $results[0] | Should -BeExactly 'prefix-1'
+        $results[1] | Should -BeExactly 'prefix-2'
+        $results[2] | Should -BeExactly 'prefix-3'
+        $results[3] -is [psvariable] | Should -BeTrue
+    }
+
     Context "PipelineVariable Behaviour" {
 
         BeforeAll {