diff --git a/src/System.Management.Automation/engine/MshCommandRuntime.cs b/src/System.Management.Automation/engine/MshCommandRuntime.cs index 7b693ce0e..209652584 100644 --- a/src/System.Management.Automation/engine/MshCommandRuntime.cs +++ b/src/System.Management.Automation/engine/MshCommandRuntime.cs @@ -928,7 +928,8 @@ namespace System.Management.Automation /// internal string PipelineVariable { get; set; } - private PSVariable _pipelineVarReference = null; + private PSVariable _pipelineVarReference; + private bool _shouldRemovePipelineVariable; internal void SetupOutVariable() { @@ -972,7 +973,7 @@ namespace System.Management.Automation // 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,12 +984,24 @@ namespace System.Management.Automation _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) { @@ -996,6 +1009,15 @@ namespace System.Management.Automation } } + internal void RemovePipelineVariable() + { + if (_shouldRemovePipelineVariable) + { + // Remove pipeline variable when a pipeline is being torn down. + _state.PSVariable.Remove(PipelineVariable); + } + } + /// /// Configures the number of objects to buffer before calling the downstream Cmdlet. /// @@ -3765,25 +3787,12 @@ namespace System.Management.Automation 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 @@ namespace System.Management.Automation 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 index 58ea07de6..78fb19dbc 100644 --- a/src/System.Management.Automation/engine/pipeline.cs +++ b/src/System.Management.Automation/engine/pipeline.cs @@ -1298,7 +1298,21 @@ namespace System.Management.Automation.Internal // 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 index 5bacbae7b..bf48dd346 100644 --- 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 {