diff --git a/src/System.Management.Automation/engine/CommandBase.cs b/src/System.Management.Automation/engine/CommandBase.cs index 4d90a7c249..3708611df0 100644 --- a/src/System.Management.Automation/engine/CommandBase.cs +++ b/src/System.Management.Automation/engine/CommandBase.cs @@ -8,8 +8,7 @@ using System.Management.Automation.Internal; using System.Management.Automation.Internal.Host; using System.Management.Automation.Language; using System.Management.Automation.Runspaces; - -using Dbg = System.Management.Automation.Diagnostics; +using System.Threading; namespace System.Management.Automation.Internal { @@ -132,6 +131,13 @@ namespace System.Management.Automation.Internal } } + /// + /// Gets the CancellationToken that is signaled when the pipeline is stopping. + /// + internal CancellationToken StopToken => commandRuntime is MshCommandRuntime mcr + ? mcr.PipelineProcessor.PipelineStopToken + : default; + /// /// The information about the command. /// diff --git a/src/System.Management.Automation/engine/cmdlet.cs b/src/System.Management.Automation/engine/cmdlet.cs index 742224c21d..ffe13b7304 100644 --- a/src/System.Management.Automation/engine/cmdlet.cs +++ b/src/System.Management.Automation/engine/cmdlet.cs @@ -10,7 +10,7 @@ using System.Globalization; using System.Reflection; using System.Resources; using System.Management.Automation.Internal; -using Dbg = System.Management.Automation.Diagnostics; +using System.Threading; namespace System.Management.Automation { @@ -100,6 +100,11 @@ namespace System.Management.Automation } } + /// + /// Gets the CancellationToken that is signaled when the pipeline is stopping. + /// + public CancellationToken PipelineStopToken => StopToken; + /// /// The name of the parameter set in effect. /// diff --git a/src/System.Management.Automation/engine/pipeline.cs b/src/System.Management.Automation/engine/pipeline.cs index 91bf435918..191f80e1d8 100644 --- a/src/System.Management.Automation/engine/pipeline.cs +++ b/src/System.Management.Automation/engine/pipeline.cs @@ -7,6 +7,7 @@ using System.Management.Automation.Runspaces; using System.Management.Automation.Tracing; using System.Reflection; using System.Runtime.ExceptionServices; +using System.Threading; using Microsoft.PowerShell.Telemetry; using Dbg = System.Management.Automation.Diagnostics; @@ -29,6 +30,7 @@ namespace System.Management.Automation.Internal { #region private_members + private readonly CancellationTokenSource _pipelineStopTokenSource = new CancellationTokenSource(); private List _commands = new List(); private List _redirectionPipes; private PipelineReader _externalInputPipe; @@ -85,6 +87,7 @@ namespace System.Management.Automation.Internal _externalErrorOutput = null; _executionScope = null; _eventLogBuffer = null; + _pipelineStopTokenSource.Dispose(); #if !CORECLR // Impersonation Not Supported On CSS SecurityContext.Dispose(); SecurityContext = null; @@ -118,6 +121,11 @@ namespace System.Management.Automation.Internal } } + /// + /// Gets the CancellationToken that is signaled when the pipeline is stopping. + /// + internal CancellationToken PipelineStopToken => _pipelineStopTokenSource.Token; + internal void LogExecutionInfo(InvocationInfo invocationInfo, string text) { string message = StringUtil.Format(PipelineStrings.PipelineExecutionInformation, GetCommand(invocationInfo), text); @@ -896,6 +904,8 @@ namespace System.Management.Automation.Internal return; } + _pipelineStopTokenSource.Cancel(); + // Call StopProcessing() for all the commands. foreach (CommandProcessorBase commandProcessor in commands) { diff --git a/test/powershell/Language/Scripting/PipelineStoppedToken.Tests.ps1 b/test/powershell/Language/Scripting/PipelineStoppedToken.Tests.ps1 new file mode 100644 index 0000000000..54d1510a20 --- /dev/null +++ b/test/powershell/Language/Scripting/PipelineStoppedToken.Tests.ps1 @@ -0,0 +1,97 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +Describe 'PipelineStopToken tests' -Tags 'CI' { + + BeforeAll { + Function Invoke-WithStop { + [CmdletBinding()] + param ( + [Parameter(Mandatory)] + [ScriptBlock] + $ScriptBlock, + + [Parameter(ValueFromRemainingArguments)] + [object[]] + $ArgumentList + ) + + $ps = [PowerShell]::Create() + $null = $ps.AddScript("'start'`n" + $ScriptBlock.ToString()) + foreach ($arg in $ArgumentList) { + $null = $ps.AddArgument($arg) + } + + $inPipe = [System.Management.Automation.PSDataCollection[object]]::new() + $inPipe.Complete() + $outPipe = [System.Management.Automation.PSDataCollection[object]]::new() + + # Use an event to make sure Stop is called once the pipeline has started + # and not before. + $eventId = [Guid]::NewGuid().ToString() + Register-ObjectEvent -InputObject $outPipe -EventName DataAdded -SourceIdentifier $eventId + try { + $task = $ps.BeginInvoke($inPipe, $outPipe) + Wait-Event -SourceIdentifier $eventId | Remove-Event + } + finally { + Remove-Event -SourceIdentifier $eventId -ErrorAction SilentlyContinue + Unregister-Event -SourceIdentifier $eventId + } + + $ps.Stop() + $ps.Streams.Error | Write-Error + { $ps.EndInvoke($task) } | Should -Throw -ErrorId PipelineStoppedException + } + } + + It 'Signal advanced function to stop' { + $start = Get-Date + + Invoke-WithStop -ScriptBlock { + Function Test-FunctionWithStop { + [CmdletBinding()] + param ([Parameter()][int]$Timeout) + + [System.Threading.Tasks.Task]::Delay($Timeout * 1000, $PSCmdlet.PipelineStopToken).GetAwaiter().GetResult() + } + + Test-FunctionWithStop -Timeout 10 + } + + $end = (Get-Date) - $start + $end.TotalSeconds | Should -BeLessThan 10 + } + + It 'Signals compiled cmdlet to stop' { + $binaryAssembly = Add-Type @' +using System; +using System.Management.Automation; +using System.Threading.Tasks; + +namespace PipelineStoppedToken.Tests; + +[Cmdlet(VerbsDiagnostic.Test, "CmdletWithStop")] +public sealed class TestCmdletWithStop : Cmdlet +{ + [Parameter] + public int Timeout { get; set; } + + protected override void EndProcessing() + { + Task.Delay(Timeout * 1000, PipelineStopToken).GetAwaiter().GetResult(); + } +} +'@ -PassThru | ForEach-Object Assembly | Select-Object -First 1 + + $start = Get-Date + Invoke-WithStop -ScriptBlock { + Import-Module -Assembly $args[0] + + Test-CmdletWithStop -Timeout 10 + } -ArgumentList $binaryAssembly + + $end = (Get-Date) - $start + $end.TotalSeconds | Should -BeLessThan 10 + } +}