Changes for Dev

src/main.rs

@@ -44,14 +44,14 @@ async fn main() {
         .and_then(|s| s.parse().ok())
         .unwrap_or(9187);
 
-    tokio::spawn(async move {
+    let omikron_handle = tokio::spawn(async move {
         match omikron_connection::start(omikron_port).await {
             Err(e) => log_err!(0, PrintType::General, "{:?}", e),
             _ => {}
         }
     });
 
-    tokio::spawn(async move {
+    let tauri_handle = tokio::spawn(async move {
         match tauri::start(9189).await {
             Err(e) => log_err!(0, PrintType::General, "{:?}", e),
             _ => {}
@@ -80,7 +80,17 @@ async fn main() {
         .and_then(|s| s.parse().ok())
         .unwrap_or(9188);
 
-    let _ = server::server::start(api_port).await;
+    tokio::select! {
+        result = server::server::start(api_port) => {
+            if let Err(e) = result {
+                log_err!(0, PrintType::General, "Server error: {:?}", e);
+            }
+        }
+        _ = tokio::signal::ctrl_c() => {
+            log!("Shutting down on signal...");
+        }
+    }
 
-    tokio::signal::ctrl_c().await.unwrap();
+    omikron_handle.abort();
+    tauri_handle.abort();
 }

src/transport/omikron_connection.rs

@@ -103,7 +103,15 @@ pub struct OmikronConnection {
     pub_key: RwLock<Option<Vec<u8>>>,
     pub ping: RwLock<i64>,
     waiting_tasks: DashMap<u32, WaitingTask>,
-    cleanup_handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
+    cleanup_handle: std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
+}
+
+impl Drop for OmikronConnection {
+    fn drop(&mut self) {
+        if let Some(handle) = self.cleanup_handle.lock().unwrap().take() {
+            handle.abort();
+        }
+    }
 }
 
 impl OmikronConnection {
@@ -120,7 +128,7 @@ impl OmikronConnection {
             pub_key: RwLock::new(None),
             ping: RwLock::new(-1),
             waiting_tasks: DashMap::new(),
-            cleanup_handle: Mutex::new(None),
+            cleanup_handle: std::sync::Mutex::new(None),
         });
 
         conn
@@ -148,7 +156,7 @@ impl OmikronConnection {
                     .retain(|_, v| v.inserted_at.elapsed() < MAX_WAITING_AGE);
             }
         });
-        *self.cleanup_handle.lock().await = Some(cleanup_handle);
+        *self.cleanup_handle.lock().unwrap() = Some(cleanup_handle);
 
         while let Ok(cv) = receiver.receive().await {
             if let Err(e) = self.clone().process_message(cv).await {
@@ -1114,8 +1122,7 @@ impl OmikronConnection {
             }
         }
 
-        // Cancel cleanup task
-        if let Some(handle) = self.cleanup_handle.lock().await.take() {
+        if let Some(handle) = self.cleanup_handle.lock().unwrap().take() {
             handle.abort();
         }
     }