Handle SIGINTs gracefully while waiting for client connection (#2045)

* Fix the exception from dotnet-trace

* Fix the same issue in dotnet-counters

diff --git a/src/Tools/Common/ReversedServerHelpers/ReversedServerHelpers.cs b/src/Tools/Common/ReversedServerHelpers/ReversedServerHelpers.cs
index f1a83e3757fffb3ead0ed044665cd230323a1b0e..2e5890dc6963f53770f9399d78ded724e15dd0fd 100644 (file)
--- 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 5bedd0a84a65f4473418ea34813ce194323a577c..2753ee100104e5d8c89207ddc66ddcd260a141a0 100644 (file)
--- 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 8c959a359f81b06d18284997d477d67674584568..9c29fd454e8bef65a81528201e29a7d613c7ba5a 100644 (file)
--- 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)
                     {