use crate::core::handle; use crate::process::AsyncHandler; use crate::utils::{connections_stream, tray_speed}; use crate::{Type, logging}; use parking_lot::Mutex; use std::sync::Arc; use std::time::Duration; use tauri::async_runtime::JoinHandle; use tauri_plugin_mihomo::models::WsConnectionId; const TRAY_SPEED_RETRY_DELAY: Duration = Duration::from_secs(1); const TRAY_SPEED_STALE_TIMEOUT: Duration = Duration::from_secs(5); #[derive(Clone)] pub struct TraySpeedController { speed_task: Arc>>>, speed_connection_id: Arc>>, } impl Default for TraySpeedController { fn default() -> Self { Self { speed_task: Arc::new(Mutex::new(None)), speed_connection_id: Arc::new(Mutex::new(None)), } } } impl TraySpeedController { pub fn new() -> Self { Self::default() } pub fn update_task(&self, enable_tray_speed: bool) { if enable_tray_speed { self.start_task(); } else { self.stop_task(); } } fn start_task(&self) { if handle::Handle::global().is_exiting() { return; } if !Self::has_main_tray() { logging!(warn, Type::Tray, "托盘不可用,跳过启动托盘速率任务"); return; } let mut guard = self.speed_task.lock(); if guard.as_ref().is_some_and(|task| !task.inner().is_finished()) { return; } let speed_connection_id = Arc::clone(&self.speed_connection_id); let task = AsyncHandler::spawn(move || async move { loop { if handle::Handle::global().is_exiting() { break; } if !Self::has_main_tray() { logging!(warn, Type::Tray, "托盘已不可用,停止托盘速率任务"); break; } let stream_connect_result = connections_stream::connect_traffic_stream().await; let mut speed_stream = match stream_connect_result { Ok(stream) => stream, Err(err) => { logging!(debug, Type::Tray, "托盘速率流连接失败,稍后重试: {err}"); Self::apply_tray_speed(0, 0); tokio::time::sleep(TRAY_SPEED_RETRY_DELAY).await; continue; } }; Self::set_speed_connection_id(&speed_connection_id, Some(speed_stream.connection_id)); loop { let next_state = speed_stream .next_event(TRAY_SPEED_STALE_TIMEOUT, || handle::Handle::global().is_exiting()) .await; match next_state { connections_stream::StreamConsumeState::Event(speed_event) => { Self::apply_tray_speed(speed_event.up, speed_event.down); } connections_stream::StreamConsumeState::Stale => { logging!(debug, Type::Tray, "托盘速率流长时间未收到有效数据,触发重连"); Self::apply_tray_speed(0, 0); break; } connections_stream::StreamConsumeState::Closed | connections_stream::StreamConsumeState::ExitRequested => { break; } } } Self::disconnect_speed_connection(&speed_connection_id).await; if handle::Handle::global().is_exiting() || !Self::has_main_tray() { break; } // `Stale` already reset the display; this covers remote close. Self::apply_tray_speed(0, 0); tokio::time::sleep(TRAY_SPEED_RETRY_DELAY).await; } Self::set_speed_connection_id(&speed_connection_id, None); }); *guard = Some(task); } fn stop_task(&self) { let task = self.speed_task.lock().take(); let speed_connection_id = Arc::clone(&self.speed_connection_id); AsyncHandler::spawn(move || async move { // Await abort before disconnect so the task cannot consume and lose the connection ID. if let Some(task) = task { task.abort(); let _ = task.await; } Self::disconnect_speed_connection(&speed_connection_id).await; }); let app_handle = handle::Handle::app_handle(); if let Some(tray) = app_handle.tray_by_id(super::TRAY_ID) { let result = tray.with_inner_tray_icon(|inner| { if let Some(status_item) = inner.ns_status_item() { tray_speed::clear_speed_attributed_title(&status_item); } }); if let Err(err) = result { logging!(warn, Type::Tray, "清除富文本速率失败: {err}"); } } } fn has_main_tray() -> bool { handle::Handle::app_handle().tray_by_id(super::TRAY_ID).is_some() } fn set_speed_connection_id( speed_connection_id: &Arc>>, connection_id: Option, ) { *speed_connection_id.lock() = connection_id; } fn take_speed_connection_id(speed_connection_id: &Arc>>) -> Option { speed_connection_id.lock().take() } async fn disconnect_speed_connection(speed_connection_id: &Arc>>) { if let Some(connection_id) = Self::take_speed_connection_id(speed_connection_id) { connections_stream::disconnect_connection(connection_id).await; } } fn apply_tray_speed(up: u64, down: u64) { let app_handle = handle::Handle::app_handle(); if let Some(tray) = app_handle.tray_by_id(super::TRAY_ID) { let result = tray.with_inner_tray_icon(move |inner| { if let Some(status_item) = inner.ns_status_item() { tray_speed::set_speed_attributed_title(&status_item, up, down); } }); if let Err(err) = result { logging!(warn, Type::Tray, "设置富文本速率失败: {err}"); } } } }