Panic cleanup

omni-led-lib/src/events/event_loop.rs

@@ -11,7 +11,7 @@ impl EventLoop {
         Self {}
     }
 
-    pub async fn run<F: FnMut(Vec<Event>)>(
+    pub fn run<F: FnMut(Vec<Event>)>(
         &self,
         interval: Duration,
         running: &AtomicBool,
@@ -29,7 +29,7 @@ impl EventLoop {
             let update_duration = end - begin;
             trace!("Update took {:?}", update_duration);
 
-            tokio::time::sleep(interval.saturating_sub(update_duration)).await;
+            std::thread::sleep(interval.saturating_sub(update_duration));
         }
     }
 }

omni-led-lib/src/server/server.rs

@@ -6,6 +6,8 @@ use serde::Serialize;
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use tokio::net::TcpListener;
+use tokio::runtime::Runtime;
+use tokio::sync::oneshot;
 use tokio_stream::StreamExt;
 use tonic::transport::Server;
 use tonic::{Code, Request, Response, Status, Streaming};
@@ -22,21 +24,24 @@ pub struct PluginServer {
 }
 
 impl PluginServer {
-    pub async fn load(lua: &Lua) {
+    pub fn load(lua: &Lua, rt: &Runtime) -> oneshot::Sender<()> {
         const LOCALHOST: &str = "127.0.0.1";
 
         let settings = UserDataRef::<Settings>::load(lua);
 
         let port: u16 = settings.get().server_port;
-        let listener = TcpListener::bind(format!("{LOCALHOST}:{port}"))
-            .await
+
+        let listener = rt
+            .block_on(TcpListener::bind(format!("{LOCALHOST}:{port}")))
             .unwrap();
+
         let address = listener.local_addr().unwrap();
         let bound_port = address.port();
 
         let log_level_filter = settings.get().log_level.into();
 
-        tokio::task::spawn(
+        let (tx, rx) = oneshot::channel();
+        rt.spawn(
             Server::builder()
                 .add_service(
                     omni_led_api::types::plugin_server::PluginServer::new(Self::new(
@@ -45,7 +50,13 @@ impl PluginServer {
                     .max_decoding_message_size(64 * 1024 * 1024)
                     .max_encoding_message_size(64 * 1024 * 1024),
                 )
-                .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener)),
+                .serve_with_incoming_shutdown(
+                    tokio_stream::wrappers::TcpListenerStream::new(listener),
+                    async {
+                        _ = rx.await;
+                        debug!("Server shutting down")
+                    },
+                ),
         );
 
         let address = format!("{LOCALHOST}:{bound_port}");
@@ -71,6 +82,8 @@ impl PluginServer {
         .unwrap();
 
         lua.globals().set("SERVER", info).unwrap();
+
+        tx
     }
 
     fn new(log_level_filter: log::LevelFilter) -> Self {

omni-led-lib/src/ui/mod.rs

@@ -1,6 +1,6 @@
+pub mod event;
 pub mod handler;
 pub mod window;
 
-mod event;
 mod icon_image;
 mod tray_icon;

omni-led/src/main.rs

@@ -17,17 +17,20 @@ use omni_led_lib::{
     script_handler::script_handler::ScriptHandler,
     server::server::PluginServer,
     settings::settings::Settings,
-    ui::handler::HandlerBuilder,
+    ui::event::Event,
+    ui::handler::{HandlerBuilder, PROXY},
 };
 use std::sync;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::time::Instant;
+use std::time::{Duration, Instant};
 
 mod logging;
 
 static RUNNING: AtomicBool = AtomicBool::new(true);
 
 fn main() {
+    set_panic_hook();
+
     let (constants_tx, constants_rx) = sync::mpsc::channel();
     let (ready_tx, ready_rx) = sync::mpsc::channel();
 
@@ -36,65 +39,77 @@ fn main() {
             .enable_all()
             .build()
             .unwrap();
-        rt.block_on(async {
-            let init_begin = Instant::now();
-
-            let lua = Lua::new();
-
-            load_internal_functions(&lua);
-            Constants::load(&lua);
-
-            constants_tx
-                .send(UserDataRef::<Constants>::load(&lua).get().clone())
-                .unwrap();
-
-            let _ = ready_rx.recv().unwrap();
-
-            let log_handle = logging::init(&lua);
-            Log::load(&lua, log_handle);
-
-            let applications_config = read_config(&lua, ConfigType::Applications).unwrap();
-            let devices_config = read_config(&lua, ConfigType::Devices).unwrap();
-            let scripts_config = read_config(&lua, ConfigType::Scripts).unwrap();
-            let settings_config = read_config(&lua, ConfigType::Settings).unwrap();
 
-            Settings::load(&lua, settings_config);
-            PluginServer::load(&lua).await;
-            Events::load(&lua);
-            Shortcuts::load(&lua);
-            Devices::load(&lua, devices_config);
-            ScriptHandler::load(&lua, scripts_config);
-            AppLoader::load(&lua, applications_config);
+        let init_begin = Instant::now();
 
-            let keyboard_handle = std::thread::spawn(|| process_events(&RUNNING));
+        let lua = Lua::new();
 
-            let init_end = Instant::now();
-            debug!("Initialized in {:?}", init_end - init_begin);
+        load_internal_functions(&lua);
+        Constants::load(&lua);
 
-            let settings = UserDataRef::<Settings>::load(&lua);
-            let interval = settings.get().update_interval;
-            let event_loop = EventLoop::new();
-            event_loop
-                .run(interval, &RUNNING, |events| {
-                    let dispatcher = UserDataRef::<Events>::load(&lua);
-                    for event in events {
-                        dispatcher.get().dispatch(&lua, event).unwrap();
-                    }
-
-                    let mut script_handler = UserDataRef::<ScriptHandler>::load(&lua);
-                    script_handler.get_mut().update(&lua, interval).unwrap();
-                })
-                .await;
+        constants_tx
+            .send(UserDataRef::<Constants>::load(&lua).get().clone())
+            .unwrap();
 
-            keyboard_handle.join().unwrap();
+        let _ = ready_rx.recv().unwrap();
+
+        let log_handle = logging::init(&lua);
+        Log::load(&lua, log_handle);
+
+        let applications_config = read_config(&lua, ConfigType::Applications).unwrap();
+        let devices_config = read_config(&lua, ConfigType::Devices).unwrap();
+        let scripts_config = read_config(&lua, ConfigType::Scripts).unwrap();
+        let settings_config = read_config(&lua, ConfigType::Settings).unwrap();
+
+        Settings::load(&lua, settings_config);
+        let server_shutdown_tx = PluginServer::load(&lua, &rt);
+        Events::load(&lua);
+        Shortcuts::load(&lua);
+        Devices::load(&lua, devices_config);
+        ScriptHandler::load(&lua, scripts_config);
+        AppLoader::load(&lua, applications_config);
+
+        let init_end = Instant::now();
+        debug!("Initialized in {:?}", init_end - init_begin);
+
+        let settings = UserDataRef::<Settings>::load(&lua);
+        let interval = settings.get().update_interval;
+        let event_loop = EventLoop::new();
+        event_loop.run(interval, &RUNNING, |events| {
+            let dispatcher = UserDataRef::<Events>::load(&lua);
+            for event in events {
+                dispatcher.get().dispatch(&lua, event).unwrap();
+            }
+
+            let mut script_handler = UserDataRef::<ScriptHandler>::load(&lua);
+            script_handler.get_mut().update(&lua, interval).unwrap();
         });
+
+        server_shutdown_tx.send(()).unwrap();
+        rt.shutdown_timeout(Duration::from_secs(1));
     });
 
+    let keyboard_thread = std::thread::spawn(|| process_events(&RUNNING));
+
     HandlerBuilder::new()
         .with_constants(constants_rx.recv().unwrap())
         .with_on_init(move || ready_tx.send(true).unwrap())
         .run();
 
     RUNNING.store(false, Ordering::Relaxed);
     _ = scripting_thread.join();
+    _ = keyboard_thread.join().unwrap();
+}
+
+fn set_panic_hook() {
+    // Exit all loops on panic to avoid deadlocks on cleanup
+
+    let hook = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |info| {
+        if let Some(proxy) = PROXY.get() {
+            proxy.send(Event::Quit)
+        }
+        RUNNING.store(false, Ordering::Relaxed);
+        hook(info);
+    }));
 }