clash-verge-rev/src-tauri/src/core/tray/speed_task.rs

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<Mutex<Option<JoinHandle<()>>>>,
    speed_connection_id: Arc<Mutex<Option<WsConnectionId>>>,
}

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<Mutex<Option<WsConnectionId>>>,
        connection_id: Option<WsConnectionId>,
    ) {
        *speed_connection_id.lock() = connection_id;
    }

    fn take_speed_connection_id(speed_connection_id: &Arc<Mutex<Option<WsConnectionId>>>) -> Option<WsConnectionId> {
        speed_connection_id.lock().take()
    }

    async fn disconnect_speed_connection(speed_connection_id: &Arc<Mutex<Option<WsConnectionId>>>) {
        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}");
            }
        }
    }
}