Using CancellationToken in dotnet-trace (#367)

diff --git a/src/Tools/dotnet-trace/CommandLine/Commands/CollectCommand.cs b/src/Tools/dotnet-trace/CommandLine/Commands/CollectCommand.cs
index be8cfd7fdb51caa8904350e6267355953b83927f..dbd329d7e2b2bb2f542b0263b4ac8fd2f5a0efa4 100644 (file)
--- a/src/Tools/dotnet-trace/CommandLine/Commands/CollectCommand.cs
+++ b/src/Tools/dotnet-trace/CommandLine/Commands/CollectCommand.cs
@@ -6,6 +6,7 @@ using Microsoft.Diagnostics.Tools.RuntimeClient;
 using System;
 using System.Collections.Generic;
 using System.CommandLine;
+using System.CommandLine.Binding;
 using System.CommandLine.Rendering;
 using System.Diagnostics;
 using System.IO;
@@ -17,9 +18,12 @@ namespace Microsoft.Diagnostics.Tools.Trace
 {
     internal static class CollectCommandHandler
     {
+        delegate Task<int> CollectDelegate(CancellationToken ct, IConsole console, int processId, FileInfo output, uint buffersize, string providers, string profile, TraceFileFormat format);
+
         /// <summary>
         /// Collects a diagnostic trace from a currently running process.
         /// </summary>
+        /// <param name="ct">The cancellation token</param>
         /// <param name="console"></param>
         /// <param name="processId">The process to collect the trace from.</param>
         /// <param name="output">The output path for the collected trace data.</param>
@@ -28,7 +32,7 @@ namespace Microsoft.Diagnostics.Tools.Trace
         /// <param name="profile">A named pre-defined set of provider configurations that allows common tracing scenarios to be specified succinctly.</param>
         /// <param name="format">The desired format of the created trace file.</param>
         /// <returns></returns>
-        private static async Task<int> Collect(IConsole console, int processId, FileInfo output, uint buffersize, string providers, string profile, TraceFileFormat format)
+        private static async Task<int> Collect(CancellationToken ct, IConsole console, int processId, FileInfo output, uint buffersize, string providers, string profile, TraceFileFormat format)
         {
             try
             {
@@ -96,6 +100,8 @@ namespace Microsoft.Diagnostics.Tools.Trace
                 var failed = false;
                 var terminated = false;
 
+                ct.Register(() => shouldExit.Set());
+
                 ulong sessionId = 0;
                 using (Stream stream = EventPipeClient.CollectTracing(processId, configuration, out sessionId))
                 using (VirtualTerminalMode vTermMode = VirtualTerminalMode.TryEnable())
@@ -149,10 +155,6 @@ namespace Microsoft.Diagnostics.Tools.Trace
                     collectingTask.Start();
 
                     Console.Out.WriteLine("Press <Enter> or <Ctrl+C> to exit...");
-                    Console.CancelKeyPress += (sender, args) => {
-                        args.Cancel = true;
-                        shouldExit.Set();
-                    };
 
                     do {
                         while (!Console.KeyAvailable && !shouldExit.WaitOne(250)) { }
@@ -238,7 +240,7 @@ namespace Microsoft.Diagnostics.Tools.Trace
                     ProfileOption(),
                     CommonOptions.FormatOption(),
                 },
-                handler: System.CommandLine.Invocation.CommandHandler.Create<IConsole, int, FileInfo, uint, string, string, TraceFileFormat>(Collect));
+                handler: HandlerDescriptor.FromDelegate((CollectDelegate)Collect).GetCommandHandler());
 
         private static uint DefaultCircularBufferSizeInMB => 256;