Drain crashed testhost stderr without blocking a thread-pool thread

src/Microsoft.TestPlatform.PlatformAbstractions/common/System/ProcessHelper.cs

@@ -6,10 +6,9 @@
 using System.Diagnostics;
 using System.Globalization;
 using System.IO;
+using System.Runtime.CompilerServices;
 using System.Threading;
-#if !NET
 using System.Threading.Tasks;
-#endif
 
 using Microsoft.VisualStudio.TestPlatform.PlatformAbstractions.Interfaces;
 
@@ -23,6 +22,23 @@ public partial class ProcessHelper : IProcessHelper
     private static readonly string Arm = "arm";
     private readonly Process _currentProcess = Process.GetCurrentProcess();
 
+    // Bounded time (ms) we wait for a crashed process's redirected stderr to reach EOF before reading it,
+    // so a late-delivered crash callstack (e.g. "Stack overflow.") is not dropped. See the Exited handler.
+    private const int CrashErrorDrainTimeout = 5000;
+
+    // Bounded time (ms) we wait for stderr to drain after a clean exit, kept short so neither the common
+    // case nor the rare grandchild-keeps-the-pipe-open case adds latency. See the Exited handler.
+    private const int CleanExitErrorDrainTimeout = 500;
+
+    // Processes we deliberately killed (e.g. when aborting or cleaning up a run). Their abnormal exit code is
+    // expected and is not a crash, so the exit handler must not spend the long stderr-drain budget on them -
+    // that would make aborting a run from an IDE slow whenever a grandchild process (e.g. a browser driver)
+    // keeps the stderr pipe open. ConditionalWeakTable holds only weak references to the processes, so entries
+    // disappear when a process is collected and nothing has to be removed explicitly.
+    private readonly ConditionalWeakTable<Process, object> _deliberatelyTerminatedProcesses = new();
+    private readonly object _deliberatelyTerminatedProcessesLock = new();
+    private static readonly object DeliberateTerminationMarker = new();
+
 #if !NET
     private readonly IEnvironment _environment;
 #endif
@@ -127,20 +143,22 @@ void InitializeAndStart()
                 process.OutputDataReceived += (sender, args) => outputCallBack(sender as Process, args.Data);
             }
 
-            // Set once the redirected stderr stream reaches EOF (signaled by a null Data event,
+            // Completed once the redirected stderr stream reaches EOF (signaled by a null Data event,
             // which is raised after all stderr lines have been handed to errorCallback). This is
             // the only reliable signal that the asynchronously-collected error output is complete:
             // neither WaitForExit(timeout) nor WaitForExitAsync guarantees the ErrorDataReceived
-            // callbacks have run. The exit handler below waits (bounded) on this before reading.
-            ManualResetEventSlim? errorStreamClosed = null;
+            // callbacks have run. The exit handler below awaits (bounded) on this before reading.
+            TaskCompletionSource<bool>? errorStreamClosed = null;
             if (errorCallback != null)
             {
-                errorStreamClosed = new ManualResetEventSlim(initialState: false);
+                // RunContinuationsAsynchronously so completing this from the ErrorDataReceived callback does not
+                // inline the exit handler's continuation onto the stderr-reader thread.
+                errorStreamClosed = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
                 process.ErrorDataReceived += (sender, args) =>
                 {
                     if (args.Data is null)
                     {
-                        errorStreamClosed.Set();
+                        errorStreamClosed.TrySetResult(true);
                     }
 
                     errorCallback(sender as Process, args.Data);
@@ -151,8 +169,8 @@ void InitializeAndStart()
             {
                 process.Exited += async (sender, args) =>
                 {
-                    const int timeout = 500;
-                    var stopwatch = Stopwatch.StartNew();
+                    // Bounded time we give the process to fully exit after we are notified of its exit.
+                    const int processExitTimeout = 500;
 
                     if (sender is Process p)
                     {
@@ -172,7 +190,9 @@ void InitializeAndStart()
                             //
                             // For older frameworks, the solution is more tricky but it seems we can get the expected
                             // behavior using the parameterless 'WaitForExit()' combined with an awaited Task.Run call.
-                            var cts = new CancellationTokenSource(timeout);
+                            // 'using' so the timer the timeout allocates is released as soon as we are done waiting,
+                            // instead of leaking one per process exit when many test hosts are spawned.
+                            using var cts = new CancellationTokenSource(processExitTimeout);
 #if NET
                             await p.WaitForExitAsync(cts.Token);
 #else
@@ -191,7 +211,7 @@ void InitializeAndStart()
                             // the testhost to become a zombie process in the first place.
                             if (_environment.OperatingSystem is PlatformOperatingSystem.Windows)
                             {
-                                p.WaitForExit(timeout);
+                                p.WaitForExit(processExitTimeout);
                             }
                             else
                             {
@@ -201,6 +221,12 @@ void InitializeAndStart()
                                     {
                                         if (!p.HasExited)
                                         {
+                                            // We are force-killing a process that overran the exit budget (e.g. a
+                                            // grandchild keeps it hanging). Record it as a deliberate termination -
+                                            // exactly like TerminateProcess does - BEFORE killing, so the stderr
+                                            // drain below uses the short abort budget instead of treating our own
+                                            // kill as a crash and waiting the generous 5s budget unnecessarily.
+                                            MarkDeliberatelyTerminated(p);
                                             p.Kill();
                                         }
                                     }
@@ -222,12 +248,28 @@ void InitializeAndStart()
                             // the exit) or InvalidOperationException.
                         }
 
-                        // The process has exited. Within the SAME bounded budget used above, wait for the
-                        // redirected stderr to reach EOF so that asynchronously-collected error output
-                        // (e.g. a testhost crash callstack such as "Stack overflow.") is complete before the
-                        // exit callback consumes it. WaitForExit(timeout)/WaitForExitAsync do not guarantee
-                        // the ErrorDataReceived callbacks have run.
-                        WaitForErrorStreamToDrain(errorStreamClosed, timeout, stopwatch.ElapsedMilliseconds);
+                        // The process has exited. Asynchronously wait (bounded) for the redirected stderr to reach
+                        // EOF so that asynchronously-collected error output (e.g. a testhost crash callstack such as
+                        // "Stack overflow.") is complete before the exit callback consumes it.
+                        // WaitForExit(timeout)/WaitForExitAsync do not guarantee the ErrorDataReceived callbacks
+                        // have run.
+                        //
+                        // We await rather than block here on purpose: the crash callstack can be delivered to
+                        // ErrorDataReceived noticeably late under load (e.g. thread-pool starvation while many test
+                        // hosts run in parallel on CI), and blocking a thread-pool thread for the whole drain budget
+                        // would compete with the very ErrorDataReceived callback we are waiting for and could starve
+                        // it out. Dropping that output both produces a misleading error message and makes
+                        // RunTestsShouldThrowOnStackOverflowException flaky.
+                        //
+                        // This drain budget is intentionally separate from (and far more generous than) the
+                        // process-exit budget above, and the generous part is only spent when the process crashed -
+                        // i.e. it exited abnormally on its own. A clean exit, or a process we deliberately killed
+                        // (e.g. aborting a run from an IDE), gets only a short grace period so we never add latency
+                        // to those cases - in particular we must not hang for seconds on abort when a grandchild
+                        // process keeps the stderr pipe open and EOF never arrives. In every case the wait returns
+                        // as soon as EOF is observed, so a process that exits and drains promptly pays almost nothing.
+                        var errorDrainTimeout = GetErrorDrainTimeout(DidProcessExitCleanly(p), WasDeliberatelyTerminated(p));
+                        await WaitForErrorStreamToDrainAsync(errorStreamClosed, errorDrainTimeout).ConfigureAwait(false);
                     }
 
                     // If exit callback has code that access Process object, ensure that the exceptions handling should be done properly.
@@ -252,24 +294,78 @@ void InitializeAndStart()
     }
 
     /// <summary>
-    /// Waits, bounded by the time remaining in <paramref name="budgetMilliseconds"/>, for the redirected
-    /// standard error stream to reach EOF (signaled via <paramref name="errorStreamClosed"/>). This ensures
-    /// all <see cref="Process.ErrorDataReceived"/> callbacks have completed - and therefore the captured
-    /// error output is complete - before it is consumed by the exit callback. It returns immediately when
-    /// there is no redirected error stream, when it has already drained, or when the budget is already
-    /// exhausted (e.g. a grandchild process keeps the pipe open), so the caller can never hang.
+    /// Asynchronously waits, bounded by <paramref name="timeoutMilliseconds"/>, for the redirected standard
+    /// error stream to reach EOF (signaled by completing <paramref name="errorStreamClosed"/>). This ensures all
+    /// <see cref="Process.ErrorDataReceived"/> callbacks have completed - and therefore the captured error
+    /// output is complete - before it is consumed by the exit callback. It returns immediately when there is
+    /// no redirected error stream or when the timeout is not positive, and is otherwise bounded by the timeout
+    /// (e.g. a grandchild process keeps the pipe open), so the caller can never hang. It deliberately does not
+    /// block the calling thread while waiting, so it does not consume a thread-pool thread that the pending
+    /// <see cref="Process.ErrorDataReceived"/> callback may itself need in order to deliver EOF under
+    /// thread-pool starvation.
     /// </summary>
-    internal static void WaitForErrorStreamToDrain(ManualResetEventSlim? errorStreamClosed, int budgetMilliseconds, long elapsedMilliseconds)
+    internal static async Task WaitForErrorStreamToDrainAsync(TaskCompletionSource<bool>? errorStreamClosed, int timeoutMilliseconds)
     {
-        if (errorStreamClosed is null)
+        if (errorStreamClosed is null || timeoutMilliseconds <= 0)
         {
             return;
         }
 
-        var remainingMilliseconds = budgetMilliseconds - (int)elapsedMilliseconds;
-        if (remainingMilliseconds > 0)
+        using var timeoutCancellation = new CancellationTokenSource();
+        var delayTask = Task.Delay(timeoutMilliseconds, timeoutCancellation.Token);
+        var completedTask = await Task.WhenAny(errorStreamClosed.Task, delayTask).ConfigureAwait(false);
+
+        // Stop the timer as soon as the stream drains so we don't leave it pending for the whole timeout.
+        if (completedTask != delayTask)
+        {
+            timeoutCancellation.Cancel();
+        }
+    }
+
+    /// <summary>
+    /// Returns <see langword="true"/> when the process has exited with a zero exit code. A non-zero exit code
+    /// (a crash) - or an exit code that cannot be retrieved - is treated as not-clean so the caller waits the
+    /// longer stderr drain budget and does not truncate potentially important crash output.
+    /// </summary>
+    private static bool DidProcessExitCleanly(Process process)
+    {
+        try
+        {
+            return process.HasExited && process.ExitCode == 0;
+        }
+        catch
+        {
+            // If the exit code is not retrievable (e.g. the process handle is gone), assume a crash so we
+            // give the redirected stderr the longer budget to drain.
+            return false;
+        }
+    }
+
+    /// <summary>
+    /// Picks the bounded time we are willing to wait for the redirected stderr to reach EOF. A genuine crash
+    /// (an abnormal exit we did not cause) gets the generous budget so a late-delivered crash callstack is
+    /// captured; a clean exit, or a process we deliberately killed (an abort/cleanup), gets only the short
+    /// budget so we never add latency to those cases.
+    /// </summary>
+    internal static int GetErrorDrainTimeout(bool exitedCleanly, bool deliberatelyTerminated)
+        => exitedCleanly || deliberatelyTerminated ? CleanExitErrorDrainTimeout : CrashErrorDrainTimeout;
+
+    private void MarkDeliberatelyTerminated(Process process)
+    {
+        lock (_deliberatelyTerminatedProcessesLock)
+        {
+            if (!_deliberatelyTerminatedProcesses.TryGetValue(process, out _))
+            {
+                _deliberatelyTerminatedProcesses.Add(process, DeliberateTerminationMarker);
+            }
+        }
+    }
+
+    private bool WasDeliberatelyTerminated(Process process)
+    {
+        lock (_deliberatelyTerminatedProcessesLock)
         {
-            errorStreamClosed.Wait(remainingMilliseconds);
+            return _deliberatelyTerminatedProcesses.TryGetValue(process, out _);
         }
     }
 
@@ -346,6 +442,11 @@ public void TerminateProcess(object? process)
         {
             if (process is Process proc && !proc.HasExited)
             {
+                // Killing a still-running process on purpose (abort/cleanup): record it so the exit handler
+                // treats the resulting abnormal exit as an abort (short stderr drain), not a crash (long
+                // stderr drain). A process that exits on its own is never recorded here, so a genuine crash
+                // still gets the generous budget.
+                MarkDeliberatelyTerminated(proc);
                 proc.Kill();
             }
         }