Gracefully shutdown all threads when exiting

Author: Jack251970

Flow.Launcher/App.xaml.cs

@@ -25,12 +25,29 @@
 
 namespace Flow.Launcher
 {
-    public partial class App : IDisposable, ISingleInstanceApp
+    public partial class App : IAsyncDisposable, ISingleInstanceApp
     {
+        #region Public Properties
+
         public static IPublicAPI API { get; private set; }
+        public static CancellationTokenSource NativeThreadCTS { get; private set; }
+
+        #endregion
+
+        #region Private Fields
+
         private static bool _disposed;
+        private MainWindow _mainWindow;
+        private readonly MainViewModel _mainVM;
         private readonly Settings _settings;
 
+        // To prevent two disposals running at the same time.
+        private static readonly object _disposingLock = new();
+
+        #endregion
+
+        #region Constructor
+
         public App()
         {
             // Initialize settings
@@ -78,34 +95,47 @@ public App()
             {
                 API = Ioc.Default.GetRequiredService<IPublicAPI>();
                 _settings.Initialize();
+                _mainVM = Ioc.Default.GetRequiredService<MainViewModel>();
             }
             catch (Exception e)
             {
                 ShowErrorMsgBoxAndFailFast("Cannot initialize api and settings, please open new issue in Flow.Launcher", e);
                 return;
             }
-        }
 
-        private static void ShowErrorMsgBoxAndFailFast(string message, Exception e)
-        {
-            // Firstly show users the message
-            MessageBox.Show(e.ToString(), message, MessageBoxButton.OK, MessageBoxImage.Error);
+            // Local function
+            static void ShowErrorMsgBoxAndFailFast(string message, Exception e)
+            {
+                // Firstly show users the message
+                MessageBox.Show(e.ToString(), message, MessageBoxButton.OK, MessageBoxImage.Error);
 
-            // Flow cannot construct its App instance, so ensure Flow crashes w/ the exception info.
-            Environment.FailFast(message, e);
+                // Flow cannot construct its App instance, so ensure Flow crashes w/ the exception info.
+                Environment.FailFast(message, e);
+            }
         }
 
+        #endregion
+
+        #region Main
+
         [STAThread]
         public static void Main()
         {
+            NativeThreadCTS = new CancellationTokenSource();
+
             if (SingleInstance<App>.InitializeAsFirstInstance())
             {
-                using var application = new App();
+                var application = new App();
                 application.InitializeComponent();
                 application.Run();
+                application.DisposeAsync().AsTask().GetAwaiter().GetResult();
             }
         }
 
+        #endregion
+
+        #region App Events
+
 #pragma warning disable VSTHRD100 // Avoid async void methods
 
         private async void OnStartup(object sender, StartupEventArgs e)
@@ -136,11 +166,11 @@ await Stopwatch.NormalAsync("|App.OnStartup|Startup cost", async () =>
                 await PluginManager.InitializePluginsAsync();
                 await imageLoadertask;
 
-                var window = new MainWindow();
+                _mainWindow = new MainWindow();
 
                 Log.Info($"|App.OnStartup|Dependencies Info:{ErrorReporting.DependenciesInfo()}");
 
-                Current.MainWindow = window;
+                Current.MainWindow = _mainWindow;
                 Current.MainWindow.Title = Constant.FlowLauncher;
 
                 HotKeyMapper.Initialize();
@@ -157,8 +187,7 @@ await Stopwatch.NormalAsync("|App.OnStartup|Startup cost", async () =>
                 AutoUpdates();
 
                 API.SaveAppAllSettings();
-                Log.Info(
-                    "|App.OnStartup|End Flow Launcher startup ----------------------------------------------------  ");
+                Log.Info("|App.OnStartup|End Flow Launcher startup ----------------------------------------------------");
             });
         }
 
@@ -191,7 +220,6 @@ private void AutoStartup()
             }
         }
 
-        //[Conditional("RELEASE")]
         private void AutoUpdates()
         {
             _ = Task.Run(async () =>
@@ -209,11 +237,30 @@ private void AutoUpdates()
             });
         }
 
+        #endregion
+
+        #region Register Events
+
         private void RegisterExitEvents()
         {
-            AppDomain.CurrentDomain.ProcessExit += (s, e) => Dispose();
-            Current.Exit += (s, e) => Dispose();
-            Current.SessionEnding += (s, e) => Dispose();
+            AppDomain.CurrentDomain.ProcessExit += (s, e) =>
+            {
+                Log.Info("|App.RegisterExitEvents|Process Exit");
+                _ = DisposeAsync();
+            };
+
+            Current.Exit += (s, e) =>
+            {
+                NativeThreadCTS.Cancel();
+                Log.Info("|App.RegisterExitEvents|Application Exit");
+                _ = DisposeAsync();
+            };
+
+            Current.SessionEnding += (s, e) =>
+            {
+                Log.Info("|App.RegisterExitEvents|Session Ending");
+                _ = DisposeAsync();
+            };
         }
 
         /// <summary>
@@ -234,20 +281,62 @@ private static void RegisterAppDomainExceptions()
             AppDomain.CurrentDomain.UnhandledException += ErrorReporting.UnhandledExceptionHandle;
         }
 
-        public void Dispose()
+        #endregion
+
+        #region IAsyncDisposable
+
+        protected virtual async ValueTask DisposeAsync(bool disposing)
         {
-            // if sessionending is called, exit proverbially be called when log off / shutdown
-            // but if sessionending is not called, exit won't be called when log off / shutdown
-            if (!_disposed)
+            // Prevent two disposes at the same time.
+            lock (_disposingLock)
             {
-                API.SaveAppAllSettings();
+                if (!disposing)
+                {
+                    return;
+                }
+
+                if (_disposed)
+                {
+                    return;
+                }
+
                 _disposed = true;
             }
+
+            await Stopwatch.NormalAsync("|App.Dispose|Dispose cost", async () =>
+            {
+                Log.Info("|App.Dispose|Begin Flow Launcher dispose ----------------------------------------------------");
+
+                if (disposing)
+                {
+                    API?.SaveAppAllSettings();
+                    await PluginManager.DisposePluginsAsync();
+
+                    // Dispose needs to be called on the main Windows thread, since some resources owned by the thread need to be disposed.
+                    await _mainWindow?.Dispatcher.InvokeAsync(DisposeAsync);
+                    _mainVM?.Dispose();
+                }
+
+                Log.Info("|App.Dispose|End Flow Launcher dispose ----------------------------------------------------");
+            });
         }
 
+        public async ValueTask DisposeAsync()
+        {
+            // Do not change this code. Put cleanup code in 'DisposeAsync(bool disposing)' method
+            await DisposeAsync(disposing: true);
+            GC.SuppressFinalize(this);
+        }
+
+        #endregion
+
+        #region ISingleInstanceApp
+
         public void OnSecondAppStarted()
         {
             Ioc.Default.GetRequiredService<MainViewModel>().Show();
         }
+
+        #endregion
     }
 }

Flow.Launcher/MainWindow.xaml

@@ -17,6 +17,7 @@
     AllowDrop="True"
     AllowsTransparency="True"
     Background="Transparent"
+    Closed="OnClosed"
     Closing="OnClosing"
     Deactivated="OnDeactivated"
     Icon="Images/app.png"

Flow.Launcher/MainWindow.xaml.cs

@@ -27,7 +27,7 @@
 
 namespace Flow.Launcher
 {
-    public partial class MainWindow
+    public partial class MainWindow : IDisposable
     {
         #region Private Fields
 
@@ -50,13 +50,17 @@ public partial class MainWindow
         private SoundPlayer animationSoundWPF;
 
         // Window WndProc
+        private HwndSource _hwndSource;
         private int _initialWidth;
         private int _initialHeight;
 
         // Window Animation
         private const double DefaultRightMargin = 66; //* this value from base.xaml
         private bool _animating;
-        private bool _isClockPanelAnimating = false; // 애니메이션 실행 중인지 여부
+        private bool _isClockPanelAnimating = false;
+
+        // IDisposable
+        private bool _disposedValue = false;
 
         #endregion
 
@@ -85,8 +89,8 @@ public MainWindow()
         private void OnSourceInitialized(object sender, EventArgs e)
         {
             var handle = Win32Helper.GetWindowHandle(this, true);
-            var win = HwndSource.FromHwnd(handle);
-            win.AddHook(WndProc);
+            _hwndSource = HwndSource.FromHwnd(handle);
+            _hwndSource.AddHook(WndProc);
             Win32Helper.HideFromAltTab(this);
             Win32Helper.DisableControlBox(this);
         }
@@ -227,20 +231,31 @@ private async void OnLoaded(object sender, RoutedEventArgs _)
                 .AddValueChanged(History, (s, e) => UpdateClockPanelVisibility());
         }
 
-        private async void OnClosing(object sender, CancelEventArgs e)
+        private void OnClosing(object sender, CancelEventArgs e)
         {
+            _viewModel.Save();
             _notifyIcon.Visible = false;
-            App.API.SaveAppAllSettings();
-            e.Cancel = true;
-            await PluginManager.DisposePluginsAsync();
             Notification.Uninstall();
-            Environment.Exit(0);
+        }
+
+        private void OnClosed(object sender, EventArgs e)
+        {
+            try
+            {
+                _hwndSource.RemoveHook(WndProc);
+            }
+            catch (Exception)
+            {
+                // Ignored
+            }
+
+            _hwndSource = null;
         }
 
         private void OnLocationChanged(object sender, EventArgs e)
         {
-            if (_animating)
-                return;
+            if (_animating) return;
+
             if (_settings.SearchWindowScreen == SearchWindowScreens.RememberLastLaunchLocation)
             {
                 _settings.WindowLeft = Left;
@@ -990,5 +1005,29 @@ private void QueryTextBox_OnPreviewDragOver(object sender, DragEventArgs e)
         }
 
         #endregion
+
+        #region IDisposable
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!_disposedValue)
+            {
+                if (disposing)
+                {
+                    _hwndSource?.Dispose();
+                }
+
+                _disposedValue = true;
+            }
+        }
+
+        public void Dispose()
+        {
+            // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
+            Dispose(disposing: true);
+            GC.SuppressFinalize(this);
+        }
+
+        #endregion
     }
 }

Flow.Launcher/ViewModel/MainViewModel.cs

@@ -27,7 +27,7 @@
 
 namespace Flow.Launcher.ViewModel
 {
-    public partial class MainViewModel : BaseModel, ISavable
+    public partial class MainViewModel : BaseModel, ISavable, IDisposable
     {
         #region Private Fields
 
@@ -1542,5 +1542,29 @@ public void UpdateResultView(ICollection<ResultsForUpdate> resultsForUpdates)
         }
 
         #endregion
+
+        #region IDisposable
+
+        private bool _disposed = false;
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!_disposed)
+            {
+                if (disposing)
+                {
+                    _updateSource?.Dispose();
+                    _disposed = true;
+                }
+            }
+        }
+
+        public void Dispose()
+        {
+            Dispose(disposing: true);
+            GC.SuppressFinalize(this);
+        }
+
+        #endregion
     }
 }