Dongbo Wang (daxian-dbw) · Mar 24, 2025 · Nov 26, 2024 · Jan 28, 2025 · Mar 7, 2025 · Mar 7, 2025
@@ -8,8 +8,7 @@
 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 @@ internal bool IsStopping
             }
         }
 
+        /// <summary>
+        /// Gets the CancellationToken that is signaled when the pipeline is stopping.
+        /// </summary>
+        internal CancellationToken StopToken => commandRuntime is MshCommandRuntime mcr
+            ? mcr.PipelineProcessor.PipelineStopToken
+            : default;
+
         /// <summary>
         /// The information about the command.
         /// </summary>

@@ -10,7 +10,7 @@
 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 @@ public bool Stopping
             }
         }
 
+        /// <summary>
+        /// Gets the CancellationToken that is signaled when the pipeline is stopping.
+        /// </summary>
+        public CancellationToken PipelineStopToken => StopToken;
+
         /// <summary>
         /// The name of the parameter set in effect.
         /// </summary>

@@ -7,6 +7,7 @@
 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 @@ internal class PipelineProcessor : IDisposable
     {
         #region private_members
 
+        private readonly CancellationTokenSource _pipelineStopTokenSource = new CancellationTokenSource();
         private List<CommandProcessorBase> _commands = new List<CommandProcessorBase>();
         private List<PipelineProcessor> _redirectionPipes;
         private PipelineReader<object> _externalInputPipe;
@@ -85,6 +87,7 @@ private void Dispose(bool disposing)
                 _externalErrorOutput = null;
                 _executionScope = null;
                 _eventLogBuffer = null;
+                _pipelineStopTokenSource.Dispose();
 #if !CORECLR // Impersonation Not Supported On CSS
                 SecurityContext.Dispose();
                 SecurityContext = null;
@@ -118,6 +121,11 @@ internal bool ExecutionFailed
             }
         }
 
+        /// <summary>
+        /// Gets the CancellationToken that is signaled when the pipeline is stopping.
+        /// </summary>
+        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 @@ internal void Stop()
                 return;
             }
 
+            _pipelineStopTokenSource.Cancel();
+
             // Call StopProcessing() for all the commands.
             foreach (CommandProcessorBase commandProcessor in commands)
             {

@@ -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
+    }
+}