PowerShell · daxian-dbw · Jul 24, 2023 · Mar 29, 2023 · Mar 29, 2023 · Mar 29, 2023
diff --git a/src/Microsoft.PowerShell.Commands.Management/commands/management/Process.cs b/src/Microsoft.PowerShell.Commands.Management/commands/management/Process.cs
@@ -822,6 +822,7 @@ private static string RetrieveProcessUserName(Process process)
     /// This class implements the Wait-process command.
     /// </summary>
     [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 @@ public int Timeout
                 _timeOutSpecified = true;
             }
         }
+
+        /// <summary>
+        /// Gets or sets a value indicating whether to return after any one process exits.
+        /// </summary>
+        [Parameter]
+        public SwitchParameter Any { get; set; }
+
+        /// <summary>
+        /// Gets or sets a value indicating whether to return the Process objects after waiting.
+        /// </summary>
+        [Parameter]
+        public SwitchParameter PassThru { get; set; }
 
         private int _timeout = 0;
         private bool _timeOutSpecified;
@@ -927,7 +940,7 @@ public void Dispose()
         // 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 @@ protected override void EndProcessing()
             {
                 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,33 +1013,42 @@ protected override void EndProcessing()
                 }
             }
 
+            bool hasTimedOut = false;
             if (_numberOfProcessesToWaitFor > 0)
             {
                 if (_timeOutSpecified)
                 {
-                    _waitHandle.WaitOne(_timeout * 1000);
+                    hasTimedOut = !_waitHandle.WaitOne(_timeout * 1000);
                 }
                 else
                 {
                     _waitHandle.WaitOne();
                 }
             }
 
-            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
@@ -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)
+    }
+}