From e8ff6c785a49a7fc8609dfdeffb104438ed739e5 Mon Sep 17 00:00:00 2001
From: huzibaca <masterorigin91@gmail.com>
Date: Thu, 2 Jan 2025 15:27:28 +0800
Subject: [PATCH] fix: modify the external control access key, the tray rate
 display is abnormal

---
 src-tauri/src/core/core.rs     |  7 +++++-
 src-tauri/src/core/tray/mod.rs | 46 ++++++++++++++++++++++++++--------
 src-tauri/src/utils/resolve.rs |  4 ---
 3 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/src-tauri/src/core/core.rs b/src-tauri/src/core/core.rs
index d9bfe3fe..4b0d4162 100644
--- a/src-tauri/src/core/core.rs
+++ b/src-tauri/src/core/core.rs
@@ -1,3 +1,4 @@
+use super::tray::Tray;
 use crate::config::*;
 use crate::core::{clash_api, handle, service};
 use crate::log_err;
@@ -93,8 +94,12 @@ impl CoreManager {
         if service::check_service().await.is_ok() {
             log::info!(target: "app", "try to run core in service mode");
             service::run_core_by_service(&config_path).await?;
-            *running = true;
         }
+        // 流量订阅
+        #[cfg(target_os = "macos")]
+        log_err!(Tray::global().subscribe_traffic().await);
+
+        *running = true;
 
         Ok(())
     }
diff --git a/src-tauri/src/core/tray/mod.rs b/src-tauri/src/core/tray/mod.rs
index 62face10..794e893b 100644
--- a/src-tauri/src/core/tray/mod.rs
+++ b/src-tauri/src/core/tray/mod.rs
@@ -35,6 +35,7 @@ use super::handle;
 pub struct Tray {
     pub speed_rate: Arc<Mutex<Option<SpeedRate>>>,
     shutdown_tx: Arc<RwLock<Option<broadcast::Sender<()>>>>,
+    is_subscribed: Arc<RwLock<bool>>,
 }
 
 #[cfg(not(target_os = "macos"))]
@@ -48,6 +49,7 @@ impl Tray {
         return TRAY.get_or_init(|| Tray {
             speed_rate: Arc::new(Mutex::new(None)),
             shutdown_tx: Arc::new(RwLock::new(None)),
+            is_subscribed: Arc::new(RwLock::new(false)),
         });
 
         #[cfg(not(target_os = "macos"))]
@@ -286,27 +288,48 @@ impl Tray {
     #[cfg(target_os = "macos")]
     pub async fn subscribe_traffic(&self) -> Result<()> {
         log::info!(target: "app", "subscribe traffic");
+
+        // 如果已经订阅，先取消订阅
+        if *self.is_subscribed.read() {
+            self.unsubscribe_traffic();
+        }
+
         let (shutdown_tx, shutdown_rx) = broadcast::channel(1);
         *self.shutdown_tx.write() = Some(shutdown_tx);
+        *self.is_subscribed.write() = true;
 
         let speed_rate = Arc::clone(&self.speed_rate);
+        let is_subscribed = Arc::clone(&self.is_subscribed);
 
         tauri::async_runtime::spawn(async move {
             let mut shutdown = shutdown_rx;
-            if let Ok(mut stream) = Traffic::get_traffic_stream().await {
-                loop {
-                    tokio::select! {
-                        Some(traffic) = stream.next() => {
-                            if let Ok(traffic) = traffic {
-                                let guard = speed_rate.lock();
-                                if let Some(sr) = guard.as_ref() {
-                                    if let Some(rate) = sr.update_and_check_changed(traffic.up, traffic.down) {
-                                        let _ = Tray::global().update_icon(Some(rate));
+
+            'outer: loop {
+                match Traffic::get_traffic_stream().await {
+                    Ok(mut stream) => loop {
+                        tokio::select! {
+                            Some(traffic) = stream.next() => {
+                                if let Ok(traffic) = traffic {
+                                    let guard = speed_rate.lock();
+                                    if let Some(sr) = guard.as_ref() {
+                                        if let Some(rate) = sr.update_and_check_changed(traffic.up, traffic.down) {
+                                            let _ = Tray::global().update_icon(Some(rate));
+                                        }
                                     }
                                 }
                             }
+                            _ = shutdown.recv() => break 'outer,
+                        }
+                    },
+                    Err(e) => {
+                        log::error!(target: "app", "Failed to get traffic stream: {}", e);
+                        // 如果获取流失败，等待一段时间后重试
+                        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+
+                        // 检查是否应该继续重试
+                        if !*is_subscribed.read() {
+                            break;
                         }
-                        _ = shutdown.recv() => break,
                     }
                 }
             }
@@ -319,8 +342,9 @@ impl Tray {
     #[cfg(target_os = "macos")]
     pub fn unsubscribe_traffic(&self) {
         log::info!(target: "app", "unsubscribe traffic");
+        *self.is_subscribed.write() = false;
         if let Some(tx) = self.shutdown_tx.write().take() {
-            drop(tx); // 发送端被丢弃时会自动发送关闭信号
+            drop(tx);
         }
     }
 }
diff --git a/src-tauri/src/utils/resolve.rs b/src-tauri/src/utils/resolve.rs
index a2110ae4..8f3c42d8 100644
--- a/src-tauri/src/utils/resolve.rs
+++ b/src-tauri/src/utils/resolve.rs
@@ -106,10 +106,6 @@ pub async fn resolve_setup(app: &mut App) {
     log_err!(tray::Tray::global().update_part());
     log_err!(hotkey::Hotkey::global().init());
     log_err!(timer::Timer::global().init());
-
-    // 流量订阅
-    #[cfg(target_os = "macos")]
-    log_err!(tray::Tray::global().subscribe_traffic().await);
 }
 
 /// reset system proxy