[PTRun] Gracefully shutdown all threads when exiting (#20450)

microsoft · Sep 13, 2022 · 9d8bfea · 9d8bfea
1 parent 8872cf6
commit 9d8bfea
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 26 deletions.
diff --git a/src/modules/launcher/PowerLauncher/App.xaml.cs b/src/modules/launcher/PowerLauncher/App.xaml.cs
@@ -31,7 +31,10 @@ public partial class App : IDisposable, ISingleInstanceApp
     {
         public static PublicAPIInstance API { get; private set; }
 
+        public static CancellationTokenSource NativeThreadCTS { get; private set; }
+
         private static bool _disposed;
+
         private PowerToysRunSettings _settings;
         private MainViewModel _mainVM;
         private MainWindow _mainWindow;
@@ -46,6 +49,8 @@ public partial class App : IDisposable, ISingleInstanceApp
         [STAThread]
         public static void Main()
         {
+            NativeThreadCTS = new CancellationTokenSource();
+
             Log.Info($"Starting PowerToys Run with PID={Environment.ProcessId}", typeof(App));
             int powerToysPid = GetPowerToysPId();
             if (powerToysPid != 0)
@@ -67,15 +72,13 @@ public static void Main()
             using (var application = new App())
             {
                 application.InitializeComponent();
-                new Thread(() =>
+                NativeEventWaiter.WaitForEventLoop(
+                    Constants.RunExitEvent(),
+                    () =>
                 {
-                    var eventHandle = new EventWaitHandle(false, EventResetMode.AutoReset, Constants.RunExitEvent());
-                    if (eventHandle.WaitOne())
-                    {
-                        Log.Warn("RunExitEvent was signaled. Exiting PowerToys", typeof(App));
-                        ExitPowerToys(application);
-                    }
-                }).Start();
+                    Log.Warn("RunExitEvent was signaled. Exiting PowerToys", typeof(App));
+                    ExitPowerToys(application);
+                }, NativeThreadCTS.Token);
 
                 if (powerToysPid != 0)
                 {
@@ -118,8 +121,8 @@ private void OnStartup(object sender, StartupEventArgs e)
                 StringMatcher.Instance = _stringMatcher;
                 _stringMatcher.UserSettingSearchPrecision = _settings.QuerySearchPrecision;
 
-                _mainVM = new MainViewModel(_settings);
-                _mainWindow = new MainWindow(_settings, _mainVM);
+                _mainVM = new MainViewModel(_settings, NativeThreadCTS.Token);
+                _mainWindow = new MainWindow(_settings, _mainVM, NativeThreadCTS.Token);
                 API = new PublicAPIInstance(_settingsVM, _mainVM, _themeManager);
                 _settingsReader = new SettingsReader(_settings, _themeManager);
                 _settingsReader.ReadSettings();
@@ -154,14 +157,7 @@ private static void ExitPowerToys(App app)
         {
             SingleInstance<App>.SingleInstanceMutex.Close();
 
-            try
-            {
-                app.Dispose();
-            }
-            finally
-            {
-                Environment.Exit(0);
-            }
+            app.Dispatcher.Invoke(() => app.Shutdown());
         }
 
         private static int GetPowerToysPId()
@@ -194,6 +190,7 @@ private void RegisterExitEvents()
 
             Current.Exit += (s, e) =>
             {
+                NativeThreadCTS.Cancel();
                 Log.Info("Application.Current.Exit", GetType());
                 Dispose();
             };

diff --git a/src/modules/launcher/PowerLauncher/Helper/NativeEventWaiter.cs b/src/modules/launcher/PowerLauncher/Helper/NativeEventWaiter.cs
@@ -12,18 +12,22 @@ namespace PowerLauncher.Helper
 {
     public static class NativeEventWaiter
     {
-        public static void WaitForEventLoop(string eventName, Action callback)
+        public static void WaitForEventLoop(string eventName, Action callback, CancellationToken cancel)
         {
             new Thread(() =>
             {
                 var eventHandle = new EventWaitHandle(false, EventResetMode.AutoReset, eventName);
                 while (true)
                 {
-                    if (eventHandle.WaitOne())
+                    if (WaitHandle.WaitAny(new WaitHandle[] { cancel.WaitHandle, eventHandle }) == 1)
                     {
                         Log.Info($"Successfully waited for {eventName}", MethodBase.GetCurrentMethod().DeclaringType);
                         Application.Current.Dispatcher.Invoke(callback);
                     }
+                    else
+                    {
+                        return;
+                    }
                 }
             }).Start();
         }

diff --git a/src/modules/launcher/PowerLauncher/MainWindow.xaml.cs b/src/modules/launcher/PowerLauncher/MainWindow.xaml.cs
@@ -20,6 +20,7 @@
 using PowerLauncher.Telemetry.Events;
 using PowerLauncher.ViewModel;
 using Wox.Infrastructure.UserSettings;
+using CancellationToken = System.Threading.CancellationToken;
 using KeyEventArgs = System.Windows.Input.KeyEventArgs;
 using Log = Wox.Plugin.Logger.Log;
 using Screen = System.Windows.Forms.Screen;
@@ -30,6 +31,7 @@ public partial class MainWindow : IDisposable
     {
         private readonly PowerToysRunSettings _settings;
         private readonly MainViewModel _viewModel;
+        private readonly CancellationToken _nativeWaiterCancelToken;
         private bool _isTextSetProgrammatically;
         private bool _deletePressed;
         private HwndSource _hwndSource;
@@ -38,18 +40,19 @@ public partial class MainWindow : IDisposable
         private bool _disposedValue;
         private IDisposable _reactiveSubscription;
 
-        public MainWindow(PowerToysRunSettings settings, MainViewModel mainVM)
+        public MainWindow(PowerToysRunSettings settings, MainViewModel mainVM, CancellationToken nativeWaiterCancelToken)
             : this()
         {
             DataContext = mainVM;
             _viewModel = mainVM;
+            _nativeWaiterCancelToken = nativeWaiterCancelToken;
             _settings = settings;
 
             InitializeComponent();
 
             _firstDeleteTimer.Elapsed += CheckForFirstDelete;
             _firstDeleteTimer.Interval = 1000;
-            NativeEventWaiter.WaitForEventLoop(Constants.RunSendSettingsTelemetryEvent(), SendSettingsTelemetry);
+            NativeEventWaiter.WaitForEventLoop(Constants.RunSendSettingsTelemetryEvent(), SendSettingsTelemetry, _nativeWaiterCancelToken);
         }
 
         private void SendSettingsTelemetry()

diff --git a/src/modules/launcher/PowerLauncher/ViewModel/MainViewModel.cs b/src/modules/launcher/PowerLauncher/ViewModel/MainViewModel.cs
@@ -48,6 +48,7 @@ public class MainViewModel : BaseModel, ISavable, IDisposable
         private CancellationTokenSource _updateSource;
 
         private CancellationToken _updateToken;
+        private CancellationToken _nativeWaiterCancelToken;
         private bool _saved;
         private ushort _hotkeyHandle;
 
@@ -59,15 +60,15 @@ public class MainViewModel : BaseModel, ISavable, IDisposable
 
         internal HotkeyManager HotkeyManager { get; private set; }
 
-        public MainViewModel(PowerToysRunSettings settings)
+        public MainViewModel(PowerToysRunSettings settings, CancellationToken nativeThreadCancelToken)
         {
             _saved = false;
             _queryTextBeforeLeaveResults = string.Empty;
             _currentQuery = _emptyQuery;
             _disposed = false;
 
             _settings = settings ?? throw new ArgumentNullException(nameof(settings));
-
+            _nativeWaiterCancelToken = nativeThreadCancelToken;
             _historyItemsStorage = new WoxJsonStorage<QueryHistory>();
             _userSelectedRecordStorage = new WoxJsonStorage<UserSelectedRecord>();
             _history = _historyItemsStorage.Load();
@@ -92,12 +93,12 @@ public void RegisterHotkey(IntPtr hwnd)
             Log.Info("RegisterHotkey()", GetType());
 
             // Allow OOBE to call PowerToys Run.
-            NativeEventWaiter.WaitForEventLoop(Constants.PowerLauncherSharedEvent(), OnHotkey);
+            NativeEventWaiter.WaitForEventLoop(Constants.PowerLauncherSharedEvent(), OnHotkey, _nativeWaiterCancelToken);
 
             if (_settings.StartedFromPowerToysRunner)
             {
                 // Allow runner to call PowerToys Run from the centralized keyboard hook.
-                NativeEventWaiter.WaitForEventLoop(Constants.PowerLauncherCentralizedHookSharedEvent(), OnCentralizedKeyboardHookHotKey);
+                NativeEventWaiter.WaitForEventLoop(Constants.PowerLauncherCentralizedHookSharedEvent(), OnCentralizedKeyboardHookHotKey, _nativeWaiterCancelToken);
             }
 
             _settings.PropertyChanged += (s, e) =>