Implement WaitForExitAsync for System.Diagnostics.Process

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/TaskCompletionSourceWithCancellation.cs → src/libraries/Common/src/System/Threading/Tasks/TaskCompletionSourceWithCancellation.cs

@@ -2,11 +2,13 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System.Threading;
-using System.Threading.Tasks;
-
-namespace System.Net.Http
+namespace System.Threading.Tasks
 {
+    /// <summary>
+    /// A <see cref="TaskCompletionSource{TResult}"/> that supports cancellation registration so that any
+    /// <seealso cref="OperationCanceledException"/>s contain the relevant <see cref="CancellationToken"/>,
+    /// while also avoiding unnecessary allocations for closure captures.
+    /// </summary>
     internal class TaskCompletionSourceWithCancellation<T> : TaskCompletionSource<T>
     {
         private CancellationToken _cancellationToken;

src/libraries/System.Diagnostics.Process/ref/System.Diagnostics.Process.cs

@@ -124,6 +124,7 @@ public void Refresh() { }
         public override string ToString() { throw null; }
         public void WaitForExit() { }
         public bool WaitForExit(int milliseconds) { throw null; }
+        public System.Threading.Tasks.Task WaitForExitAsync(System.Threading.CancellationToken cancellationToken = default) { throw null; }
         public bool WaitForInputIdle() { throw null; }
         public bool WaitForInputIdle(int milliseconds) { throw null; }
     }

src/libraries/System.Diagnostics.Process/src/System.Diagnostics.Process.csproj

@@ -44,6 +44,9 @@
     <Compile Include="$(CommonPath)Interop\Windows\Interop.Errors.cs">
       <Link>Common\Interop\Windows\Interop.Errors.cs</Link>
     </Compile>
+    <Compile Include="$(CommonPath)System\Threading\Tasks\TaskCompletionSourceWithCancellation.cs">
+      <Link>Common\System\Threading\Tasks\TaskCompletionSourceWithCancellation.cs</Link>
+    </Compile>
   </ItemGroup>
   <ItemGroup Condition=" '$(TargetsWindows)' == 'true'">
     <Compile Include="$(CommonPath)Interop\Windows\Kernel32\Interop.EnumProcessModules.cs">

src/libraries/System.Diagnostics.Process/src/System/Diagnostics/Process.cs

@@ -10,6 +10,7 @@
 using System.Runtime.Serialization;
 using System.Text;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace System.Diagnostics
 {
@@ -1340,6 +1341,107 @@ public bool WaitForExit(int milliseconds)
             return exited;
         }
 
+        /// <summary>
+        /// Instructs the Process component to wait for the associated process to exit, or
+        /// for the <paramref name="cancellationToken"/> to be canceled.
+        /// </summary>
+        /// <remarks>
+        /// Calling this method will set <see cref="EnableRaisingEvents"/> to <see langword="true" />.
+        /// </remarks>
+        /// <returns>
+        /// A task that will complete when the process has exited, cancellation has been requested,
+        /// or an error occurs.
+        /// </returns>
+        public async Task WaitForExitAsync(CancellationToken cancellationToken = default)
+        {
+            // Because the process has already started by the time this method is called,
+            // we're in a race against the process to set up our exit handlers before the process
+            // exits. As a result, there are several different flows that must be handled:
+            //
+            // CASE 1: WE ENABLE EVENTS
+            // This is the "happy path". In this case we enable events.
+            //
+            // CASE 1.1: PROCESS EXITS OR IS CANCELED AFTER REGISTERING HANDLER
+            // This case continues the "happy path". The process exits or waiting is canceled after
+            // registering the handler and no special cases are needed.
+            //
+            // CASE 1.2: PROCESS EXITS BEFORE REGISTERING HANDLER
+            // It's possible that the process can exit after we enable events but before we reigster
+            // the handler. In that case we must check for exit after registering the handler.
+            //
+            //
+            // CASE 2: PROCESS EXITS BEFORE ENABLING EVENTS
+            // The process may exit before we attempt to enable events. In that case EnableRaisingEvents
+            // will throw an exception like this:
+            //     System.InvalidOperationException : Cannot process request because the process (42) has exited.
+            // In this case we catch the InvalidOperationException. If the process has exited, our work
+            // is done and we return. If for any reason (now or in the future) enabling events fails
+            // and the process has not exited, bubble the exception up to the user.
+            //
+            //
+            // CASE 3: USER ALREADY ENABLED EVENTS
+            // In this case the user has already enabled raising events. Re-enabling events is a no-op
+            // as the value hasn't changed. However, no-op also means that if the process has already
+            // exited, EnableRaisingEvents won't throw an exception.
+            //
+            // CASE 3.1: PROCESS EXITS OR IS CANCELED AFTER REGISTERING HANDLER
+            // (See CASE 1.1)
+            //
+            // CASE 3.2: PROCESS EXITS BEFORE REGISTERING HANDLER
+            // (See CASE 1.2)
+
+            if (!Associated)
+            {
+                throw new InvalidOperationException(SR.NoAssociatedProcess);
+            }
+
+            if (!HasExited)
+            {
+                // Early out for cancellation before doing more expensive work
+                cancellationToken.ThrowIfCancellationRequested();
+            }
+
+            try
+            {
+                // CASE 1: We enable events
+                // CASE 2: Process exits before enabling events (and throws an exception)
+                // CASE 3: User already enabled events (no-op)
+                EnableRaisingEvents = true;
+            }
+            catch (InvalidOperationException)
+            {
+                // CASE 2: If the process has exited, our work is done, otherwise bubble the
+                // exception up to the user
+                if (HasExited)
+                {
+                    return;
+                }
+
+                throw;
+            }
+
+            var tcs = new TaskCompletionSourceWithCancellation<object>();
+
+            EventHandler handler = (s, e) => tcs.TrySetResult(null);
+            Exited += handler;
+
+            try
+            {
+                if (HasExited)
+                {
+                    // CASE 1.2 & CASE 3.2: Handle race where the process exits before registering the handler
+                    return;
+                }
+
+                // CASE 1.1 & CASE 3.1: Process exits or is canceled here
+                await tcs.WaitWithCancellationAsync(cancellationToken).ConfigureAwait(false);
+            }
+            finally
+            {
+                Exited -= handler;
+            }
+        }
+
         /// <devdoc>
         /// <para>
         /// Instructs the <see cref='System.Diagnostics.Process'/> component to start

src/libraries/System.Diagnostics.Process/tests/ProcessTestBase.cs

@@ -60,6 +60,18 @@ protected Process CreateProcess(Func<int> method = null)
             return p;
         }
 
+        protected Process CreateProcess(Func<Task<int>> method)
+        {
+            Process p = null;
+            using (RemoteInvokeHandle handle = RemoteExecutor.Invoke(method, new RemoteInvokeOptions { Start = false }))
+            {
+                p = handle.Process;
+                handle.Process = null;
+            }
+            AddProcessForDispose(p);
+            return p;
+        }
+
         protected Process CreateProcess(Func<string, int> method, string arg, bool autoDispose = true)
         {
             Process p = null;