stop the ETW session on Ctrl+C + restore console colors ;), fixes dot…

…net#729 (dotnet#761)

* stop the ETW session on Ctrl+C + restore console colors ;), fixes dotnet#729

* post code review fixes, dotnet#729

src/BenchmarkDotNet.Diagnostics.Windows/EtwDiagnoser.cs

@@ -40,6 +40,10 @@ protected void Start(DiagnoserActionParameters parameters)
 
             Session = CreateSession(parameters.Benchmark);
 
+            Console.CancelKeyPress += OnConsoleCancelKeyPress;
+
+            NativeWindowsConsoleHelper.OnExit += OnConsoleCancelKeyPress;
+
             EnableProvider();
 
             AttachToEvents(Session, parameters.Benchmark);
@@ -76,6 +80,9 @@ protected void Stop()
             WaitForDelayedEvents();
 
             Session.Dispose();
+
+            Console.CancelKeyPress -= OnConsoleCancelKeyPress;
+            NativeWindowsConsoleHelper.OnExit -= OnConsoleCancelKeyPress;
         }
 
         private void Clear()
@@ -84,6 +91,8 @@ private void Clear()
             StatsPerProcess.Clear();
         }
 
+        private void OnConsoleCancelKeyPress(object sender, ConsoleCancelEventArgs e) => Session?.Dispose();
+
         private static string GetSessionName(string prefix, Benchmark benchmark, ParameterInstances parameters = null)
         {
             if (parameters != null && parameters.Items.Count > 0)

src/BenchmarkDotNet.Diagnostics.Windows/NativeWindowsConsoleHelper.cs

@@ -0,0 +1,64 @@
+using System;
+using System.Runtime.InteropServices;
+
+namespace BenchmarkDotNet.Diagnostics.Windows
+{
+    // we need this class because when the Console window gets closed (by simply pressing x)
+    // the managed event is not raised for CTRL_CLOSE_EVENT, more https://stackoverflow.com/questions/474679/capture-console-exit-c-sharp
+    internal static class NativeWindowsConsoleHelper
+    {
+        private enum CtrlType
+        {
+            CTRL_C_EVENT = 0,
+            CTRL_BREAK_EVENT = 1,
+            CTRL_CLOSE_EVENT = 2,
+            CTRL_LOGOFF_EVENT = 5,
+            CTRL_SHUTDOWN_EVENT = 6
+        }
+
+        private delegate bool ConsoleCtrlHandler(CtrlType sig);
+
+        [DllImport("Kernel32")]
+        private static extern bool SetConsoleCtrlHandler(ConsoleCtrlHandler handler, bool add);
+
+        private static event EventHandler<ConsoleCancelEventArgs> onExit;
+
+        public static event EventHandler<ConsoleCancelEventArgs> OnExit
+        {
+            add
+            {
+                if (Portability.RuntimeInformation.IsWindows())
+                {
+                    SetConsoleCtrlHandler(NativeHandler, true);
+
+                    onExit += value;
+                }
+            }
+            remove
+            {
+                if (Portability.RuntimeInformation.IsWindows())
+                {
+                    SetConsoleCtrlHandler(NativeHandler, false);
+
+                    onExit -= value;
+                }
+            }
+        }
+
+        private static bool NativeHandler(CtrlType sig)
+        {
+            switch (sig)
+            {
+                case CtrlType.CTRL_C_EVENT:
+                case CtrlType.CTRL_BREAK_EVENT:
+                case CtrlType.CTRL_LOGOFF_EVENT:
+                case CtrlType.CTRL_SHUTDOWN_EVENT:
+                case CtrlType.CTRL_CLOSE_EVENT:
+                    onExit?.Invoke(null, null);
+                    return false; // if we return true from here, the event is marked as handled and the managed handler is not fired
+                default:
+                    return false;
+            }
+        }
+    }
+}

src/BenchmarkDotNet/Loggers/ConsoleLogger.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using BenchmarkDotNet.Toolchains;
 
 namespace BenchmarkDotNet.Loggers
 {
@@ -25,16 +26,17 @@ public ConsoleLogger(Dictionary<LogKind, ConsoleColor> colorScheme = null)
 
         private void Write(LogKind logKind, Action<string> write, string text)
         {
-            var colorBefore = Console.ForegroundColor;
+            ConsoleHandler.EnsureInitialized(this);
 
             try
             {
-                Console.ForegroundColor = GetColor(logKind);
+                ConsoleHandler.SetForegroundColor(GetColor(logKind));
+
                 write(text);
             }
             finally
             {
-                Console.ForegroundColor = colorBefore;
+                ConsoleHandler.RestoreForegroundColor();
             }
         }
 

src/BenchmarkDotNet/Running/TypeParser.cs

@@ -17,9 +17,13 @@ internal class TypeParser
         private static readonly Dictionary<string, string> Configuration = CreateConfiguration();
         private static readonly char[] TrimChars = { ' ' };
 
+        private static bool consoleCancelKeyPressed = false;
+
         private readonly Type[] allTypes;
         private readonly ILogger logger;
 
+        static TypeParser() => Console.CancelKeyPress += (_, __) => consoleCancelKeyPressed = true; 
+
         internal TypeParser(Type[] types, ILogger logger)
         {
             this.logger = logger;
@@ -42,9 +46,13 @@ public TypeWithMethods(Type type, MethodInfo[] methods = null)
 
         internal string[] ReadArgumentList(string[] args)
         {
-            while (args.Length == 0)
+            while (args.Length == 0 && !consoleCancelKeyPressed)
             {
                 PrintAvailable();
+
+                if (consoleCancelKeyPressed)
+                    break;
+
                 var benchmarkCaptionExample = allTypes.Length == 0 ? "Intro_00" : allTypes.First().GetDisplayName();
                 logger.WriteLineHelp(
                     $"You should select the target benchmark. Please, print a number of a benchmark (e.g. '0') or a benchmark caption (e.g. '{benchmarkCaptionExample}'):");
@@ -237,10 +245,14 @@ private void PrintAvailable()
             logger.WriteLineHelp($"Available Benchmark{(allTypes.Length > 1 ? "s" : "")}:");
 
             int numberWidth = allTypes.Length.ToString().Length;
-            for (int i = 0; i < allTypes.Length; i++)
+            for (int i = 0; i < allTypes.Length && !consoleCancelKeyPressed; i++)
                 logger.WriteLineHelp(string.Format(CultureInfo.InvariantCulture, "  #{0} {1}", i.ToString().PadRight(numberWidth), allTypes[i].GetDisplayName()));
-            logger.WriteLine();
-            logger.WriteLine();
+
+            if (!consoleCancelKeyPressed)
+            {
+                logger.WriteLine();
+                logger.WriteLine();
+            }
         }
 
         private static Dictionary<string, string> CreateConfiguration()

src/BenchmarkDotNet/Toolchains/ConsoleHandler.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Diagnostics;
+using System.Drawing;
 using System.Linq;
 using System.Threading;
 using BenchmarkDotNet.Loggers;
@@ -15,6 +16,7 @@ internal class ConsoleHandler
 
         private Process process;
         private ILogger logger;
+        private ConsoleColor? colorBefore;
 
         public ConsoleHandler(ILogger logger)
         {
@@ -44,6 +46,8 @@ public void ClearProcess()
         // This method gives us a chance to make a "best-effort" to clean anything up after Ctrl-C is type in the Console
         private void HandlerCallback(object sender, ConsoleCancelEventArgs e)
         {
+            Console.ResetColor();
+
             if (e.SpecialKey != ConsoleSpecialKey.ControlC && e.SpecialKey != ConsoleSpecialKey.ControlBreak)
                 return;
 
@@ -110,5 +114,20 @@ private bool HasProcessDied(Process process)
             }
             return true;
         }
+
+        public static void SetForegroundColor(ConsoleColor color)
+        {
+            Instance.colorBefore = Console.ForegroundColor;
+            Console.ForegroundColor = color;
+        }
+
+        public static void RestoreForegroundColor()
+        {
+            if (Instance.colorBefore.HasValue)
+            {
+                Console.ForegroundColor = Instance.colorBefore.Value;
+                Instance.colorBefore = null;
+            }
+        }
     }
 }