From 8653451cbe5aff5012fdfde2b8ce41209705bd9d Mon Sep 17 00:00:00 2001
From: Sung Yoon Whang <suwhang@microsoft.com>
Date: Wed, 3 Mar 2021 17:24:58 -0800
Subject: [PATCH] Handle SIGINTs gracefully while waiting for client connection
 (#2045)

* Fix the exception from dotnet-trace

* Fix the same issue in dotnet-counters
---
 .../ReversedServerHelpers.cs                   | 15 +++++++++++++--
 src/Tools/dotnet-counters/CounterMonitor.cs    | 18 ++++++++++++++++--
 .../CommandLine/Commands/CollectCommand.cs     |  9 +++++++--
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/Tools/Common/ReversedServerHelpers/ReversedServerHelpers.cs b/src/Tools/Common/ReversedServerHelpers/ReversedServerHelpers.cs
index f1a83e375..2e5890dc6 100644
--- a/src/Tools/Common/ReversedServerHelpers/ReversedServerHelpers.cs
+++ b/src/Tools/Common/ReversedServerHelpers/ReversedServerHelpers.cs
@@ -231,8 +231,19 @@ namespace Microsoft.Internal.Common.Utils
                 Console.WriteLine($"Waiting for connection on {fullPort}");
                 Console.WriteLine($"Start an application with the following environment variable: DOTNET_DiagnosticPorts={fullPort}");
 
-                IpcEndpointInfo endpointInfo = await server.AcceptAsync(ct);
-                return new DiagnosticsClientHolder(new DiagnosticsClient(endpointInfo.Endpoint), endpointInfo, fullPort, server);
+                try
+                {
+                    IpcEndpointInfo endpointInfo = await server.AcceptAsync(ct);
+                    return new DiagnosticsClientHolder(new DiagnosticsClient(endpointInfo.Endpoint), endpointInfo, fullPort, server);
+                }
+                catch (TaskCanceledException)
+                {
+                    if (!ct.IsCancellationRequested)
+                    {
+                        throw;
+                    }
+                    return null;
+                }
             }
             else
             {
diff --git a/src/Tools/dotnet-counters/CounterMonitor.cs b/src/Tools/dotnet-counters/CounterMonitor.cs
index 5bedd0a84..2753ee100 100644
--- a/src/Tools/dotnet-counters/CounterMonitor.cs
+++ b/src/Tools/dotnet-counters/CounterMonitor.cs
@@ -31,6 +31,7 @@ namespace Microsoft.Diagnostics.Tools.Counters
         private CounterFilter filter;
         private string _output;
         private bool pauseCmdSet;
+        private ManualResetEvent shouldExit;
         private bool shouldResumeRuntime;
         private DiagnosticsClient _diagnosticsClient;
         private EventPipeSession _session;
@@ -96,10 +97,16 @@ namespace Microsoft.Diagnostics.Tools.Counters
             {
                 return 0;
             }
+            shouldExit = new ManualResetEvent(false);
+            _ct.Register(() => shouldExit.Set());
 
             DiagnosticsClientBuilder builder = new DiagnosticsClientBuilder("dotnet-counters", 10);
             using (DiagnosticsClientHolder holder = await builder.Build(ct, _processId, diagnosticPort, showChildIO: false, printLaunchCommand: false))
             {
+                if (holder == null)
+                {
+                    return 1;
+                }
                 try
                 {
                     InitializeCounterList(counters, counter_list);
@@ -134,9 +141,18 @@ namespace Microsoft.Diagnostics.Tools.Counters
             {
                 return 0;
             }
+
+            shouldExit = new ManualResetEvent(false);
+            _ct.Register(() => shouldExit.Set());
+
             DiagnosticsClientBuilder builder = new DiagnosticsClientBuilder("dotnet-counters", 10);
             using (DiagnosticsClientHolder holder = await builder.Build(ct, _processId, diagnosticPort, showChildIO: false, printLaunchCommand: false))
             {
+                if (holder == null)
+                {
+                    return 1;
+                }
+
                 try
                 {
                     InitializeCounterList(counters, counter_list);
@@ -321,8 +337,6 @@ namespace Microsoft.Diagnostics.Tools.Counters
 
             _renderer.Initialize();
 
-            ManualResetEvent shouldExit = new ManualResetEvent(false);
-            _ct.Register(() => shouldExit.Set());
             Task monitorTask = new Task(() => {
                 try
                 {
diff --git a/src/Tools/dotnet-trace/CommandLine/Commands/CollectCommand.cs b/src/Tools/dotnet-trace/CommandLine/Commands/CollectCommand.cs
index 8c959a359..9c29fd454 100644
--- a/src/Tools/dotnet-trace/CommandLine/Commands/CollectCommand.cs
+++ b/src/Tools/dotnet-trace/CommandLine/Commands/CollectCommand.cs
@@ -151,9 +151,16 @@ namespace Microsoft.Diagnostics.Tools.Trace
                 Process process;
                 DiagnosticsClientBuilder builder = new DiagnosticsClientBuilder("dotnet-trace", 10);
                 bool shouldResumeRuntime = ProcessLauncher.Launcher.HasChildProc || !string.IsNullOrEmpty(diagnosticPort);
+                var shouldExit = new ManualResetEvent(false);
+                ct.Register(() => shouldExit.Set());
 
                 using (DiagnosticsClientHolder holder = await builder.Build(ct, processId, diagnosticPort, showChildIO: showchildio, printLaunchCommand: true))
                 {
+                    // if builder returned null, it means we received ctrl+C while waiting for clients to connect. Exit gracefully.
+                    if (holder == null)
+                    {
+                        return await Task.FromResult(ret);
+                    }
                     diagnosticsClient = holder.Client;
                     if (shouldResumeRuntime)
                     {
@@ -185,12 +192,10 @@ namespace Microsoft.Diagnostics.Tools.Trace
                         }
                     }
 
-                    var shouldExit = new ManualResetEvent(false);
                     var shouldStopAfterDuration = duration != default(TimeSpan);
                     var rundownRequested = false;
                     System.Timers.Timer durationTimer = null;
 
-                    ct.Register(() => shouldExit.Set());
 
                     using (VirtualTerminalMode vTermMode = printStatusOverTime ? VirtualTerminalMode.TryEnable() : null)
                     {
-- 
2.34.1