Fixes for PowerToys Awake

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);
@@ -45,6 +61,10 @@ public class APIHelper
         [return: MarshalAs(UnmanagedType.Bool)]
         private static extern bool AllocConsole();
 
+        [DllImport("kernel32.dll", SetLastError = true)]
+        [return: MarshalAs(UnmanagedType.Bool)]
+        private static extern bool FreeConsole();
+
         [DllImport("kernel32.dll", SetLastError = true)]
         private static extern bool SetStdHandle(int nStdHandle, IntPtr hHandle);
 
@@ -63,10 +83,21 @@ 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 DeallocateConsole()
+        {
+            FreeConsole();
+        }
+
         public static void AllocateConsole()
         {
             _log.Debug("Bootstrapping the console allocation routine.");
@@ -139,7 +170,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 +187,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 +254,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) =>
                     {
-                        if (_threadToken.IsCancellationRequested)
-                        {
-                            _threadToken.ThrowIfCancellationRequested();
-                        }
-                    }
+                        _tokenSource.Cancel();
+
+                        _timedLoopTimer.Stop();
+                    };
+
+                    _timedLoopTimer.Disposed += (s, e) =>
+                    {
+                        _log.Info("Old timer disposed.");
+                    };
+
+                    _timedLoopTimer.Start();
+
+                    WaitHandle.WaitAny(new[] { _threadToken.WaitHandle });
+                    _timedLoopTimer.Stop();
+                    _timedLoopTimer.Dispose();
 
                     return success;
                 }

src/modules/awake/Awake/Core/Constants.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 Constants
+    {
+        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(Constants.AppName),
+                IndefiniteKeepAwakeCallback(Constants.AppName),
+                TimedKeepAwakeCallback(Constants.AppName),
+                KeepDisplayOnCallback(Constants.AppName),
                 ExitCallback());
         }
 
         private static Action ExitCallback()
         {
             return () =>
             {
-                Environment.Exit(0);
+                Environment.Exit(Environment.ExitCode);
             };
         }