Fixes for PowerToys Awake

src/common/interop/interop.cpp

@@ -186,5 +186,9 @@ public
         static String ^ ShowColorPickerSharedEvent() {
             return gcnew String(CommonSharedConstants::SHOW_COLOR_PICKER_SHARED_EVENT);
         }
+
+        static String ^ AwakeExitEvent() {
+            return gcnew String(CommonSharedConstants::AWAKE_EXIT_EVENT);
+        }
     };
 }

src/common/interop/shared_constants.h

@@ -28,6 +28,9 @@ namespace CommonSharedConstants
 
     const wchar_t FANCY_ZONES_EDITOR_TOGGLE_EVENT[] = L"Local\\FancyZones-ToggleEditorEvent-1e174338-06a3-472b-874d-073b21c62f14";
 
+    // Path to the event used by Awake
+    const wchar_t AWAKE_EXIT_EVENT[] = L"Local\\PowerToysAwakeExitEvent-c0d5e305-35fc-4fb5-83ec-f6070cfaf7fe";
+
     // Max DWORD for key code to disable keys.
     const DWORD VK_DISABLED = 0x100;
 }

src/modules/awake/Awake/Core/APIHelper.cs

@@ -21,6 +21,18 @@ public enum EXECUTION_STATE : uint
         ES_SYSTEM_REQUIRED = 0x00000001,
     }
 
+    // See: https://docs.microsoft.com/windows/console/handlerroutine
+    public enum ControlType
+    {
+        CTRL_C_EVENT = 0,
+        CTRL_BREAK_EVENT = 1,
+        CTRL_CLOSE_EVENT = 2,
+        CTRL_LOGOFF_EVENT = 5,
+        CTRL_SHUTDOWN_EVENT = 6,
+    }
+
+    public delegate bool ConsoleEventHandler(ControlType ctrlType);
+
     /// <summary>
     /// Helper class that allows talking to Win32 APIs without having to rely on PInvoke in other parts
     /// of the codebase.
@@ -37,6 +49,10 @@ public class APIHelper
         private static CancellationToken _threadToken;
 
         private static Task? _runnerThread;
+        private static System.Timers.Timer _timedLoopTimer;
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        private static extern bool SetConsoleCtrlHandler(ConsoleEventHandler handler, bool add);
 
         [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
         private static extern EXECUTION_STATE SetThreadExecutionState(EXECUTION_STATE esFlags);
@@ -63,10 +79,16 @@ private static extern IntPtr CreateFile(
 
         static APIHelper()
         {
+            _timedLoopTimer = new System.Timers.Timer();
             _log = LogManager.GetCurrentClassLogger();
             _tokenSource = new CancellationTokenSource();
         }
 
+        public static void SetConsoleControlHandler(ConsoleEventHandler handler, bool addHandler)
+        {
+            SetConsoleCtrlHandler(handler, addHandler);
+        }
+
         public static void AllocateConsole()
         {
             _log.Debug("Bootstrapping the console allocation routine.");
@@ -139,7 +161,7 @@ public static void SetNoKeepAwake()
             }
             catch (OperationCanceledException)
             {
-                _log.Info("Confirmed background thread cancellation when setting passive keep awake.");
+                _log.Info("Confirmed background thread cancellation when disabling explicit keep awake.");
             }
         }
 
@@ -156,7 +178,7 @@ public static void SetTimedKeepAwake(uint seconds, Action<bool> callback, Action
             }
             catch (OperationCanceledException)
             {
-                _log.Info("Confirmed background thread cancellation when setting indefinite keep awake.");
+                _log.Info("Confirmed background thread cancellation when setting timed keep awake.");
             }
 
             _tokenSource = new CancellationTokenSource();
@@ -223,14 +245,25 @@ private static bool RunTimedLoop(uint seconds, bool keepDisplayOn = true)
                 if (success)
                 {
                     _log.Info($"Initiated temporary keep awake in background thread: {GetCurrentThreadId()}. Screen on: {keepDisplayOn}");
-                    var startTime = DateTime.UtcNow;
-                    while (DateTime.UtcNow - startTime < TimeSpan.FromSeconds(Math.Abs(seconds)))
+
+                    _timedLoopTimer = new System.Timers.Timer(seconds * 1000);
+                    _timedLoopTimer.Elapsed += (s, e) =>
+                    {
+                        _tokenSource.Cancel();
+
+                        _timedLoopTimer.Stop();
+                    };
+
+                    _timedLoopTimer.Disposed += (s, e) =>
                     {
-                        if (_threadToken.IsCancellationRequested)
-                        {
-                            _threadToken.ThrowIfCancellationRequested();
-                        }
-                    }
+                        _log.Info("Old timer disposed.");
+                    };
+
+                    _timedLoopTimer.Start();
+
+                    WaitHandle.WaitAny(new[] { _threadToken.WaitHandle });
+                    _timedLoopTimer.Stop();
+                    _timedLoopTimer.Dispose();
 
                     return success;
                 }

src/modules/awake/Awake/Core/InternalConstants.cs

@@ -0,0 +1,12 @@
+// Copyright (c) Microsoft Corporation
+// The Microsoft Corporation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+namespace Awake.Core
+{
+    internal static class InternalConstants
+    {
+        internal const string AppName = "Awake";
+        internal const string FullAppName = "PowerToys " + AppName;
+    }
+}

src/modules/awake/Awake/Core/TrayHelper.cs

@@ -6,8 +6,10 @@
 using System.Drawing;
 using System.IO;
 using System.Text.Json;
+using System.Threading.Tasks;
 using System.Windows.Forms;
 using Microsoft.PowerToys.Settings.UI.Library;
+using NLog;
 
 #pragma warning disable CS8602 // Dereference of a possibly null reference.
 #pragma warning disable CS8603 // Possible null reference return.
@@ -16,32 +18,43 @@ namespace Awake.Core
 {
     internal static class TrayHelper
     {
-        private static NotifyIcon? trayIcon;
+        private static readonly Logger _log;
 
-        private static NotifyIcon TrayIcon { get => trayIcon; set => trayIcon = value; }
+        private static NotifyIcon? _trayIcon;
 
-        private static SettingsUtils? moduleSettings;
+        private static NotifyIcon TrayIcon { get => _trayIcon; set => _trayIcon = value; }
 
-        private static SettingsUtils ModuleSettings { get => moduleSettings; set => moduleSettings = value; }
+        private static SettingsUtils? _moduleSettings;
+
+        private static SettingsUtils ModuleSettings { get => _moduleSettings; set => _moduleSettings = value; }
 
         static TrayHelper()
         {
+            _log = LogManager.GetCurrentClassLogger();
             TrayIcon = new NotifyIcon();
             ModuleSettings = new SettingsUtils();
         }
 
         public static void InitializeTray(string text, Icon icon, ContextMenuStrip? contextMenu = null)
         {
-            System.Threading.Tasks.Task.Factory.StartNew(
+            Task.Factory.StartNew(
                 (tray) =>
-            {
-                ((NotifyIcon?)tray).Text = text;
-                ((NotifyIcon?)tray).Icon = icon;
-                ((NotifyIcon?)tray).ContextMenuStrip = contextMenu;
-                ((NotifyIcon?)tray).Visible = true;
+                {
+                    ((NotifyIcon?)tray).Text = text;
+                    ((NotifyIcon?)tray).Icon = icon;
+                    ((NotifyIcon?)tray).ContextMenuStrip = contextMenu;
+                    ((NotifyIcon?)tray).Visible = true;
+
+                    _log.Info("Setting up the tray.");
+                    Application.Run();
+                    _log.Info("Tray setup complete.");
+                }, TrayIcon);
+        }
 
-                Application.Run();
-            }, TrayIcon);
+        public static void ClearTray()
+        {
+            TrayIcon.Icon = null;
+            TrayIcon.Dispose();
         }
 
         internal static void SetTray(string text, AwakeSettings settings)
@@ -50,18 +63,18 @@ internal static void SetTray(string text, AwakeSettings settings)
                 text,
                 settings.Properties.KeepDisplayOn,
                 settings.Properties.Mode,
-                PassiveKeepAwakeCallback(text),
-                IndefiniteKeepAwakeCallback(text),
-                TimedKeepAwakeCallback(text),
-                KeepDisplayOnCallback(text),
+                PassiveKeepAwakeCallback(InternalConstants.AppName),
+                IndefiniteKeepAwakeCallback(InternalConstants.AppName),
+                TimedKeepAwakeCallback(InternalConstants.AppName),
+                KeepDisplayOnCallback(InternalConstants.AppName),
                 ExitCallback());
         }
 
         private static Action ExitCallback()
         {
             return () =>
             {
-                Environment.Exit(0);
+                Environment.Exit(Environment.ExitCode);
             };
         }
 
@@ -153,28 +166,21 @@ private static Action IndefiniteKeepAwakeCallback(string moduleName)
 
         public static void SetTray(string text, bool keepDisplayOn, AwakeMode mode, Action passiveKeepAwakeCallback, Action indefiniteKeepAwakeCallback, Action<uint, uint> timedKeepAwakeCallback, Action keepDisplayOnCallback, Action exitCallback)
         {
-            var contextMenuStrip = new ContextMenuStrip();
+            ContextMenuStrip? contextMenuStrip = new ContextMenuStrip();
 
             // Main toolstrip.
-            var operationContextMenu = new ToolStripMenuItem
+            ToolStripMenuItem? operationContextMenu = new ToolStripMenuItem
             {
                 Text = "Mode",
             };
 
             // No keep-awake menu item.
-            var passiveMenuItem = new ToolStripMenuItem
+            ToolStripMenuItem? passiveMenuItem = new ToolStripMenuItem
             {
                 Text = "Off (Passive)",
             };
 
-            if (mode == AwakeMode.PASSIVE)
-            {
-                passiveMenuItem.Checked = true;
-            }
-            else
-            {
-                passiveMenuItem.Checked = false;
-            }
+            passiveMenuItem.Checked = mode == AwakeMode.PASSIVE;
 
             passiveMenuItem.Click += (e, s) =>
             {
@@ -183,38 +189,25 @@ public static void SetTray(string text, bool keepDisplayOn, AwakeMode mode, Acti
             };
 
             // Indefinite keep-awake menu item.
-            var indefiniteMenuItem = new ToolStripMenuItem
+            ToolStripMenuItem? indefiniteMenuItem = new ToolStripMenuItem
             {
                 Text = "Keep awake indefinitely",
             };
 
-            if (mode == AwakeMode.INDEFINITE)
-            {
-                indefiniteMenuItem.Checked = true;
-            }
-            else
-            {
-                indefiniteMenuItem.Checked = false;
-            }
+            indefiniteMenuItem.Checked = mode == AwakeMode.INDEFINITE;
 
             indefiniteMenuItem.Click += (e, s) =>
             {
                 // User opted to set the mode to indefinite, so we need to write new settings.
                 indefiniteKeepAwakeCallback();
             };
 
-            var displayOnMenuItem = new ToolStripMenuItem
+            ToolStripMenuItem? displayOnMenuItem = new ToolStripMenuItem
             {
                 Text = "Keep screen on",
             };
-            if (keepDisplayOn)
-            {
-                displayOnMenuItem.Checked = true;
-            }
-            else
-            {
-                displayOnMenuItem.Checked = false;
-            }
+
+            displayOnMenuItem.Checked = keepDisplayOn;
 
             displayOnMenuItem.Click += (e, s) =>
             {
@@ -223,54 +216,52 @@ public static void SetTray(string text, bool keepDisplayOn, AwakeMode mode, Acti
             };
 
             // Timed keep-awake menu item
-            var timedMenuItem = new ToolStripMenuItem
+            ToolStripMenuItem? timedMenuItem = new ToolStripMenuItem
             {
                 Text = "Keep awake temporarily",
             };
-            if (mode == AwakeMode.TIMED)
-            {
-                timedMenuItem.Checked = true;
-            }
-            else
-            {
-                timedMenuItem.Checked = false;
-            }
 
-            var halfHourMenuItem = new ToolStripMenuItem
+            timedMenuItem.Checked = mode == AwakeMode.TIMED;
+
+            ToolStripMenuItem? halfHourMenuItem = new ToolStripMenuItem
             {
                 Text = "30 minutes",
             };
+
             halfHourMenuItem.Click += (e, s) =>
             {
                 // User is setting the keep-awake to 30 minutes.
                 timedKeepAwakeCallback(0, 30);
             };
 
-            var oneHourMenuItem = new ToolStripMenuItem
+            ToolStripMenuItem? oneHourMenuItem = new ToolStripMenuItem
             {
                 Text = "1 hour",
             };
+
             oneHourMenuItem.Click += (e, s) =>
             {
                 // User is setting the keep-awake to 1 hour.
                 timedKeepAwakeCallback(1, 0);
             };
 
-            var twoHoursMenuItem = new ToolStripMenuItem
+            ToolStripMenuItem? twoHoursMenuItem = new ToolStripMenuItem
             {
                 Text = "2 hours",
             };
+
             twoHoursMenuItem.Click += (e, s) =>
             {
                 // User is setting the keep-awake to 2 hours.
                 timedKeepAwakeCallback(2, 0);
             };
 
             // Exit menu item.
-            var exitContextMenu = new ToolStripMenuItem
+            ToolStripMenuItem? exitContextMenu = new ToolStripMenuItem
             {
                 Text = "Exit",
             };
+
             exitContextMenu.Click += (e, s) =>
             {
                 // User is setting the keep-awake to 2 hours.