From 3d8ed3905e9ef029d01e3484330a23d939c2698c Mon Sep 17 00:00:00 2001
From: Den Delimarsky <1389609+dend@users.noreply.github.com>
Date: Tue, 11 May 2021 19:10:15 -0700
Subject: [PATCH] Proper threading support for all parts of the keep-awake
 process

---
 .../espresso/Espresso/Core/APIHelper.cs       | 60 ++++++++++++++-----
 src/modules/espresso/Espresso/Program.cs      | 40 +++----------
 2 files changed, 54 insertions(+), 46 deletions(-)

diff --git a/src/modules/espresso/Espresso/Core/APIHelper.cs b/src/modules/espresso/Espresso/Core/APIHelper.cs
index ac7b67e064..b56025286e 100644
--- a/src/modules/espresso/Espresso/Core/APIHelper.cs
+++ b/src/modules/espresso/Espresso/Core/APIHelper.cs
@@ -78,21 +78,14 @@ namespace Espresso.Shell.Core
             return SetAwakeState(EXECUTION_STATE.ES_CONTINUOUS);
         }
 
-        /// <summary>
-        /// Sets up the machine to be awake indefinitely.
-        /// </summary>
-        /// <param name="keepDisplayOn">Determines whether the display should be kept on while the machine is awake.</param>
-        /// <returns>Status of the attempt. True if successful, false if not.</returns>
-        public static bool SetIndefiniteKeepAwake(bool keepDisplayOn = true)
+        public static void SetIndefiniteKeepAwake(Action<bool> callback, Action failureCallback, bool keepDisplayOn = true)
         {
-            if (keepDisplayOn)
-            {
-                return SetAwakeState(EXECUTION_STATE.ES_SYSTEM_REQUIRED | EXECUTION_STATE.ES_DISPLAY_REQUIRED | EXECUTION_STATE.ES_CONTINUOUS);
-            }
-            else
-            {
-                return SetAwakeState(EXECUTION_STATE.ES_SYSTEM_REQUIRED | EXECUTION_STATE.ES_CONTINUOUS);
-            }
+            _tokenSource = new CancellationTokenSource();
+            _threadToken = _tokenSource.Token;
+
+            Task.Run(() => RunIndefiniteLoop(keepDisplayOn), _threadToken)
+                .ContinueWith((result) => callback(result.Result), TaskContinuationOptions.OnlyOnRanToCompletion)
+                .ContinueWith((result) => failureCallback, TaskContinuationOptions.NotOnRanToCompletion);
         }
 
         public static void SetTimedKeepAwake(long seconds, Action<bool> callback, Action failureCallback, bool keepDisplayOn = true)
@@ -105,6 +98,45 @@ namespace Espresso.Shell.Core
                 .ContinueWith((result) => failureCallback, TaskContinuationOptions.NotOnRanToCompletion);
         }
 
+        private static bool RunIndefiniteLoop(bool keepDisplayOn = true)
+        {
+            bool success;
+            if (keepDisplayOn)
+            {
+                success = SetAwakeState(EXECUTION_STATE.ES_SYSTEM_REQUIRED | EXECUTION_STATE.ES_DISPLAY_REQUIRED | EXECUTION_STATE.ES_CONTINUOUS);
+            }
+            else
+            {
+                success = SetAwakeState(EXECUTION_STATE.ES_SYSTEM_REQUIRED | EXECUTION_STATE.ES_CONTINUOUS);
+            }
+
+            try
+            {
+                if (success)
+                {
+                    _log.Info("Initiated indefinite keep awake in background thread.");
+                    while (true)
+                    {
+                        if (_threadToken.IsCancellationRequested)
+                        {
+                            _threadToken.ThrowIfCancellationRequested();
+                        }
+                    }
+                }
+                else
+                {
+                    _log.Info("Could not successfully set up indefinite keep awake.");
+                    return success;
+                }
+            }
+            catch (OperationCanceledException ex)
+            {
+                // Task was clearly cancelled.
+                _log.Info($"Background thread termination. Message: {ex.Message}");
+                return success;
+            }
+        }
+
         private static bool RunTimedLoop(long seconds, bool keepDisplayOn = true)
         {
             bool success = false;
diff --git a/src/modules/espresso/Espresso/Program.cs b/src/modules/espresso/Espresso/Program.cs
index 025d163d88..2e3f06cb39 100644
--- a/src/modules/espresso/Espresso/Program.cs
+++ b/src/modules/espresso/Espresso/Program.cs
@@ -186,38 +186,29 @@ namespace Espresso.Shell
                 }
             }
 
+            var exitSignal = new ManualResetEvent(false);
             if (pid != 0)
             {
                 RunnerHelper.WaitForPowerToysRunner(pid, () =>
                 {
+                    exitSignal.Set();
                     Environment.Exit(0);
                 });
             }
 
-            new ManualResetEvent(false).WaitOne();
+            exitSignal.WaitOne();
         }
 
         private static void SetupIndefiniteKeepAwake(bool displayOn)
         {
             // Indefinite keep awake.
-            bool success = APIHelper.SetIndefiniteKeepAwake(displayOn);
-            if (success)
-            {
-                _log.Info($"Currently in indefinite keep awake. Display always on: {displayOn}");
-            }
-            else
-            {
-                var errorMessage = "Could not set up the state to be indefinite keep awake.";
-                _log.Info(errorMessage);
-                _log.Debug(errorMessage);
-            }
+            APIHelper.SetIndefiniteKeepAwake(LogCompletedKeepAwakeThread, LogUnexpectedOrCancelledKeepAwakeThreadCompletion, displayOn);
         }
 
         private static void HandleEspressoConfigChange(FileSystemEventArgs fileEvent)
         {
             _log.Info("Detected a settings file change. Updating configuration...");
             _log.Info("Resetting keep-awake to normal state due to settings change.");
-            ResetNormalPowerState();
             ProcessSettings();
         }
 
@@ -277,34 +268,19 @@ namespace Espresso.Shell
         {
             _log.Info($"Timed keep-awake. Expected runtime: {time} seconds with display on setting set to {displayOn}.");
 
-            APIHelper.SetTimedKeepAwake(time, LogTimedKeepAwakeCompletion, LogUnexpectedOrCancelledKeepAwakeCompletion, displayOn);
+            APIHelper.SetTimedKeepAwake(time, LogCompletedKeepAwakeThread, LogUnexpectedOrCancelledKeepAwakeThreadCompletion, displayOn);
         }
 
-        private static void LogUnexpectedOrCancelledKeepAwakeCompletion()
+        private static void LogUnexpectedOrCancelledKeepAwakeThreadCompletion()
         {
             var errorMessage = "The keep-awake thread was terminated early.";
             _log.Info(errorMessage);
             _log.Debug(errorMessage);
         }
 
-        private static void LogTimedKeepAwakeCompletion(bool result)
+        private static void LogCompletedKeepAwakeThread(bool result)
         {
-            _log.Info($"Completed timed keep-awake successfully: {result}");
-        }
-
-        private static void ResetNormalPowerState()
-        {
-            bool success = APIHelper.SetNormalKeepAwake();
-            if (success)
-            {
-                _log.Info("Returned to normal keep-awake state.");
-            }
-            else
-            {
-                var errorMessage = "Could not return to normal keep-awake state.";
-                _log.Info(errorMessage);
-                _log.Debug(errorMessage);
-            }
+            _log.Info($"Exited keep-awake thread successfully: {result}");
         }
     }
 }