diff --git a/src/Microsoft.PowerShell.Commands.Management/commands/management/Process.cs b/src/Microsoft.PowerShell.Commands.Management/commands/management/Process.cs index 73f18cea9e..c5df76c044 100644 --- a/src/Microsoft.PowerShell.Commands.Management/commands/management/Process.cs +++ b/src/Microsoft.PowerShell.Commands.Management/commands/management/Process.cs @@ -822,6 +822,7 @@ namespace Microsoft.PowerShell.Commands /// This class implements the Wait-process command. /// [Cmdlet(VerbsLifecycle.Wait, "Process", DefaultParameterSetName = "Name", HelpUri = "https://go.microsoft.com/fwlink/?LinkID=2097146")] + [OutputType(typeof(Process))] public sealed class WaitProcessCommand : ProcessBaseCommand { #region Parameters @@ -895,6 +896,18 @@ namespace Microsoft.PowerShell.Commands _timeOutSpecified = true; } } + + /// + /// Gets or sets a value indicating whether to return after any one process exits. + /// + [Parameter] + public SwitchParameter Any { get; set; } + + /// + /// Gets or sets a value indicating whether to return the Process objects after waiting. + /// + [Parameter] + public SwitchParameter PassThru { get; set; } private int _timeout = 0; private bool _timeOutSpecified; @@ -927,7 +940,7 @@ namespace Microsoft.PowerShell.Commands // Handle Exited event and display process information. private void myProcess_Exited(object sender, System.EventArgs e) { - if (System.Threading.Interlocked.Decrement(ref _numberOfProcessesToWaitFor) == 0) + if (Any || (Interlocked.Decrement(ref _numberOfProcessesToWaitFor) == 0)) { _waitHandle?.Set(); } @@ -979,7 +992,12 @@ namespace Microsoft.PowerShell.Commands { try { - if (!process.HasExited) + // Check for processes that exit too soon for us to add an event. + if (Any && process.HasExited) + { + _waitHandle.Set(); + } + else if (!process.HasExited) { process.EnableRaisingEvents = true; process.Exited += myProcess_Exited; @@ -995,11 +1013,12 @@ namespace Microsoft.PowerShell.Commands } } + bool hasTimedOut = false; if (_numberOfProcessesToWaitFor > 0) { if (_timeOutSpecified) { - _waitHandle.WaitOne(_timeout * 1000); + hasTimedOut = !_waitHandle.WaitOne(_timeout * 1000); } else { @@ -1007,21 +1026,29 @@ namespace Microsoft.PowerShell.Commands } } - foreach (Process process in _processList) + if (hasTimedOut || (!Any && _numberOfProcessesToWaitFor > 0)) { - try + foreach (Process process in _processList) { - if (!process.HasExited) + try { - string message = StringUtil.Format(ProcessResources.ProcessNotTerminated, new object[] { process.ProcessName, process.Id }); - ErrorRecord errorRecord = new(new TimeoutException(message), "ProcessNotTerminated", ErrorCategory.CloseError, process); - WriteError(errorRecord); + if (!process.HasExited) + { + string message = StringUtil.Format(ProcessResources.ProcessNotTerminated, new object[] { process.ProcessName, process.Id }); + ErrorRecord errorRecord = new(new TimeoutException(message), "ProcessNotTerminated", ErrorCategory.CloseError, process); + WriteError(errorRecord); + } + } + catch (Win32Exception exception) + { + WriteNonTerminatingError(process, exception, ProcessResources.ProcessIsNotTerminated, "ProcessNotTerminated", ErrorCategory.CloseError); } } - catch (Win32Exception exception) - { - WriteNonTerminatingError(process, exception, ProcessResources.ProcessIsNotTerminated, "ProcessNotTerminated", ErrorCategory.CloseError); - } + } + + if (PassThru) + { + WriteObject(_processList, enumerateCollection: true); } } diff --git a/test/powershell/Modules/Microsoft.PowerShell.Management/Wait-Process.Tests.ps1 b/test/powershell/Modules/Microsoft.PowerShell.Management/Wait-Process.Tests.ps1 new file mode 100644 index 0000000000..829b9d919f --- /dev/null +++ b/test/powershell/Modules/Microsoft.PowerShell.Management/Wait-Process.Tests.ps1 @@ -0,0 +1,74 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +Describe "Wait-Process" { + + BeforeAll { + $pingCommandPath = (Get-Command -CommandType Application ping)[0].Definition + + $startProcessArgs = @{ + FilePath = $pingCommandPath + PassThru = $true + } + $nc = $IsWindows ? "-n" : "-c" + + function longPing { Start-Process @startProcessArgs -ArgumentList "$nc 10 localhost" } + function shortPing { Start-Process @startProcessArgs -ArgumentList "$nc 2 localhost" } + } + + BeforeEach { + $Processes = @( 1..3 | ForEach-Object { longPing } ) + ($shortPing = shortPing) + } + + AfterEach { + Stop-Process -InputObject $Processes + } + + It "Should wait until all processes have exited" { + Wait-Process -InputObject $Processes + + $Processes.Where({$_.HasExited -eq $true}).Count | Should -Be $Processes.Count + } + + It "Should return after all processes have exited, even if some exit before the wait starts." { + Wait-UntilTrue -sb { $shortPing.HasExited -eq $true } -IntervalInMilliseconds 100 + Wait-Process -InputObject $Processes + + $Processes.Where({$_.HasExited -eq $true}).Count | Should -Be $Processes.Count + } + + It "Should return immediately if all processes have exited before the wait starts" { + Wait-UntilTrue -sb { $Processes.HasExited -NotContains $false } -IntervalInMilliseconds 100 + Wait-Process -InputObject $Processes + + $Processes.Where({$_.HasExited -eq $true}).Count | Should -Be $Processes.Count + } + + It "Should return immediately if at least one process has exited before the wait starts" { + Wait-UntilTrue -sb { $shortPing.HasExited -eq $true } -IntervalInMilliseconds 100 + Wait-Process -InputObject $Processes -Any + + $Processes.Where({$_.HasExited -eq $true}).Count | Should -Be 1 + $Processes.Where({$_.HasExited -eq $false}).Count | Should -Be ($Processes.Count - 1) + } + + It "Should wait until any one process has exited" { + Wait-Process -InputObject $Processes -Any + + $Processes.Where({$_.HasExited -eq $true}).Count | Should -Be 1 + $Processes.Where({$_.HasExited -eq $false}).Count | Should -Be ($Processes.Count - 1) + } + + It "Should passthru all processes when all processes have exited" { + $PassThruProcesses = Wait-Process -InputObject $Processes -PassThru + + $PassThruProcesses.Where({$_.HasExited -eq $true}).Count | Should -Be $Processes.Count + } + + It "Should passthru all processes when any one process has exited" { + $PassThruProcesses = Wait-Process -InputObject $Processes -Any -PassThru + + $PassThruProcesses.Where({$_.HasExited -eq $true}).Count | Should -Be 1 + $PassThruProcesses.Where({$_.HasExited -eq $false}).Count | Should -Be ($Processes.Count - 1) + } +}