refactor timer scheduler

src-tauri/src/core/timer.rs

@@ -1,7 +1,7 @@
-use crate::{config::Config, feat, process::AsyncHandler, singleton, utils::resolve::is_resolve_done};
+use crate::{config::Config, core::handle, feat, process::AsyncHandler, singleton, utils::resolve::is_resolve_done};
 use anyhow::Result;
 use clash_verge_logging::{Type, logging, logging_error};
-use parking_lot::RwLock;
+use parking_lot::{Mutex, RwLock};
 use smartstring::alias::String;
 use std::{
     collections::HashMap,
@@ -12,14 +12,37 @@ use std::{
     time::Duration,
 };
 use tokio::{
-    sync::watch,
+    sync::mpsc,
     time::{sleep, timeout},
 };
 use tokio_stream::StreamExt as _;
-use tokio_util::time::DelayQueue;
+use tokio_util::time::{DelayQueue, delay_queue::Key};
+
+enum TimerCommand {
+    Apply(HashMap<String, u64>),
+    RunNow(String),
+    TaskFinished(String),
+}
+
+struct TaskState {
+    key: Option<Key>,
+    interval_minutes: u64,
+    running: bool,
+}
+
+impl TaskState {
+    fn new(key: Key, interval_minutes: u64) -> Self {
+        Self {
+            key: Some(key),
+            interval_minutes,
+            running: false,
+        }
+    }
+}
 
 pub struct Timer {
-    pub notify_tx: watch::Sender<HashMap<String, u64>>,
+    command_tx: mpsc::UnboundedSender<TimerCommand>,
+    command_rx: Mutex<Option<mpsc::UnboundedReceiver<TimerCommand>>>,
     pub timer_map: Arc<RwLock<HashMap<String, u64>>>,
     pub initialized: AtomicBool,
 }
@@ -28,9 +51,10 @@ singleton!(Timer, TIMER_INSTANCE);
 
 impl Timer {
     fn new() -> Self {
-        let (notify_tx, _) = watch::channel(HashMap::new());
+        let (command_tx, command_rx) = mpsc::unbounded_channel();
         Self {
-            notify_tx,
+            command_tx,
+            command_rx: Mutex::new(Some(command_rx)),
             timer_map: Arc::new(RwLock::new(HashMap::new())),
             initialized: AtomicBool::new(false),
         }
@@ -46,6 +70,13 @@ impl Timer {
             return Ok(());
         }
 
+        if let Some(command_rx) = self.command_rx.lock().take() {
+            let command_tx = self.command_tx.clone();
+            AsyncHandler::spawn(move || async move {
+                Self::run_scheduler(command_rx, command_tx).await;
+            });
+        }
+
         if let Err(e) = self.refresh().await {
             self.initialized.store(false, Ordering::SeqCst);
             logging_error!(Type::Timer, "Failed to initialize timer: {}", e);
@@ -66,53 +97,6 @@ impl Timer {
             }
         }
 
-        let mut notify_rx = self.notify_tx.subscribe();
-
-        AsyncHandler::spawn(move || async move {
-            loop {
-                let current_tasks = notify_rx.borrow_and_update().clone();
-
-                if current_tasks.is_empty() {
-                    if notify_rx.changed().await.is_err() {
-                        break;
-                    }
-                    continue;
-                }
-
-                let mut queue = DelayQueue::new();
-                for (uid, interval_minutes) in current_tasks {
-                    let delay = Duration::from_secs(interval_minutes.saturating_mul(60));
-                    queue.insert((uid, interval_minutes), delay);
-                }
-
-                loop {
-                    tokio::select! {
-                        Some(expired) = queue.next() => {
-                            let (uid, interval_minutes) = expired.into_inner();
-                            logging!(info, Type::Timer, "开始执行定时任务: uid={}", uid);
-
-                            let uid_clone = uid.clone();
-                            AsyncHandler::spawn(move || async move {
-                                Self::wait_until_resolve_done(Duration::from_millis(5000)).await;
-                                Self::async_task(&uid_clone).await;
-                            });
-
-                            let next_delay = Duration::from_secs(interval_minutes.saturating_mul(60));
-                            queue.insert((uid, interval_minutes), next_delay);
-                        }
-
-                        res = notify_rx.changed() => {
-                            if res.is_err() {
-                                return;
-                            }
-                            logging!(debug, Type::Timer, "刷新定时任务列表...");
-                            break;
-                        }
-                    }
-                }
-            }
-        });
-
         let cur_timestamp = chrono::Local::now().timestamp();
         if let Some(items) = Config::profiles().await.latest_arc().get_items() {
             for item in items.iter() {
@@ -126,11 +110,7 @@ impl Timer {
                     && cur_timestamp - (updated as i64) >= (interval as i64) * 60
                 {
                     logging!(info, Type::Timer, "立即执行到期落后任务: uid={}", uid);
-                    let uid_clone = uid.clone();
-                    AsyncHandler::spawn(move || async move {
-                        Self::wait_until_resolve_done(Duration::from_millis(5000)).await;
-                        Self::async_task(&uid_clone).await;
-                    });
+                    let _ = self.command_tx.send(TimerCommand::RunNow(uid.clone()));
                 }
             }
         }
@@ -142,25 +122,23 @@ impl Timer {
     pub async fn refresh(&self) -> Result<()> {
         let new_map = self.gen_map().await;
 
-        {
-            let mut cache = self.timer_map.write();
-            if *cache == new_map {
-                logging!(debug, Type::Timer, "No timer changes needed");
-                return Ok(());
-            }
-
-            logging!(
-                info,
-                Type::Timer,
-                "Refreshing timer tasks map, count: {}",
-                new_map.len()
-            );
-            *cache = new_map.clone();
-            drop(cache);
-
-            let _ = self.notify_tx.send(new_map);
+        let mut cache = self.timer_map.write();
+        if *cache == new_map {
+            logging!(debug, Type::Timer, "No timer changes needed");
+            return Ok(());
         }
 
+        logging!(
+            info,
+            Type::Timer,
+            "Refreshing timer tasks map, count: {}",
+            new_map.len()
+        );
+        *cache = new_map.clone();
+        drop(cache);
+
+        let _ = self.command_tx.send(TimerCommand::Apply(new_map));
+
         Ok(())
     }
 
@@ -182,6 +160,187 @@ impl Timer {
         new_map
     }
 
+    async fn run_scheduler(
+        mut command_rx: mpsc::UnboundedReceiver<TimerCommand>,
+        command_tx: mpsc::UnboundedSender<TimerCommand>,
+    ) {
+        let mut queue = DelayQueue::new();
+        let mut tasks = HashMap::new();
+
+        loop {
+            tokio::select! {
+                Some(expired) = queue.next() => {
+                    let uid = expired.into_inner();
+                    Self::run_expired_task(&mut tasks, uid, command_tx.clone());
+                }
+
+                command = command_rx.recv() => {
+                    match command {
+                        Some(TimerCommand::Apply(new_map)) => {
+                            Self::apply_timer_map(&mut queue, &mut tasks, new_map);
+                        }
+                        Some(TimerCommand::RunNow(uid)) => {
+                            Self::run_task_now(&mut queue, &mut tasks, uid, command_tx.clone());
+                        }
+                        Some(TimerCommand::TaskFinished(uid)) => {
+                            Self::finish_task(&mut queue, &mut tasks, uid);
+                        }
+                        None => break,
+                    }
+                }
+            }
+        }
+    }
+
+    fn apply_timer_map(
+        queue: &mut DelayQueue<String>,
+        tasks: &mut HashMap<String, TaskState>,
+        new_map: HashMap<String, u64>,
+    ) {
+        tasks.retain(|uid, state| {
+            if new_map.contains_key(uid) {
+                return true;
+            }
+
+            if let Some(key) = state.key.take() {
+                queue.remove(&key);
+            }
+            logging!(debug, Type::Timer, "Removed timer task for uid={}", uid);
+            false
+        });
+
+        for (uid, interval_minutes) in new_map {
+            let Some(state) = tasks.get_mut(&uid) else {
+                Self::insert_task(queue, tasks, uid, interval_minutes);
+                continue;
+            };
+
+            if state.interval_minutes == interval_minutes {
+                continue;
+            }
+
+            Self::update_task_interval(queue, &uid, state, interval_minutes);
+        }
+    }
+
+    fn insert_task(
+        queue: &mut DelayQueue<String>,
+        tasks: &mut HashMap<String, TaskState>,
+        uid: String,
+        interval_minutes: u64,
+    ) {
+        let key = Self::schedule_task(queue, &uid, interval_minutes);
+        logging!(
+            debug,
+            Type::Timer,
+            "Added timer task: uid={}, interval={}min",
+            uid,
+            interval_minutes
+        );
+        tasks.insert(uid, TaskState::new(key, interval_minutes));
+    }
+
+    fn update_task_interval(queue: &mut DelayQueue<String>, uid: &str, state: &mut TaskState, interval_minutes: u64) {
+        state.interval_minutes = interval_minutes;
+
+        if let Some(key) = state.key.as_ref() {
+            queue.reset(key, Self::interval_duration(interval_minutes));
+        } else if !state.running {
+            state.key = Some(Self::schedule_task(queue, uid, interval_minutes));
+        }
+
+        logging!(
+            debug,
+            Type::Timer,
+            "Updated timer task interval: uid={}, interval={}min",
+            uid,
+            interval_minutes
+        );
+    }
+
+    fn run_expired_task(
+        tasks: &mut HashMap<String, TaskState>,
+        uid: String,
+        command_tx: mpsc::UnboundedSender<TimerCommand>,
+    ) {
+        let Some(state) = tasks.get_mut(&uid) else {
+            return;
+        };
+
+        state.key = None;
+        if !Self::mark_task_running(state, &uid, false) {
+            return;
+        }
+
+        Self::spawn_update_task(uid, command_tx);
+    }
+
+    fn run_task_now(
+        queue: &mut DelayQueue<String>,
+        tasks: &mut HashMap<String, TaskState>,
+        uid: String,
+        command_tx: mpsc::UnboundedSender<TimerCommand>,
+    ) {
+        let Some(state) = tasks.get_mut(&uid) else {
+            return;
+        };
+
+        if !Self::mark_task_running(state, &uid, true) {
+            return;
+        }
+
+        if let Some(key) = state.key.take() {
+            queue.remove(&key);
+        }
+        Self::spawn_update_task(uid, command_tx);
+    }
+
+    fn mark_task_running(state: &mut TaskState, uid: &str, immediate: bool) -> bool {
+        if !state.running {
+            state.running = true;
+            return true;
+        }
+
+        if immediate {
+            logging!(
+                debug,
+                Type::Timer,
+                "Timer task already running, skip immediate uid={}",
+                uid
+            );
+        } else {
+            logging!(debug, Type::Timer, "Timer task already running, skip uid={}", uid);
+        }
+        false
+    }
+
+    fn finish_task(queue: &mut DelayQueue<String>, tasks: &mut HashMap<String, TaskState>, uid: String) {
+        let Some(state) = tasks.get_mut(&uid) else {
+            return;
+        };
+
+        state.running = false;
+        let key = Self::schedule_task(queue, &uid, state.interval_minutes);
+        state.key = Some(key);
+    }
+
+    fn spawn_update_task(uid: String, command_tx: mpsc::UnboundedSender<TimerCommand>) {
+        logging!(info, Type::Timer, "开始执行定时任务: uid={}", uid);
+        AsyncHandler::spawn(move || async move {
+            Self::wait_until_resolve_done(Duration::from_millis(5000)).await;
+            Self::async_task(&uid).await;
+            let _ = command_tx.send(TimerCommand::TaskFinished(uid));
+        });
+    }
+
+    fn interval_duration(interval_minutes: u64) -> Duration {
+        Duration::from_secs(interval_minutes.saturating_mul(60))
+    }
+
+    fn schedule_task(queue: &mut DelayQueue<String>, uid: &str, interval_minutes: u64) -> Key {
+        queue.insert(String::from(uid), Self::interval_duration(interval_minutes))
+    }
+
     pub async fn get_next_update_time(&self, uid: &str) -> Option<i64> {
         logging!(debug, Type::Timer, "获取下次更新时间, uid={}", uid);
 
@@ -230,6 +389,7 @@ impl Timer {
                     uid,
                     task_start.elapsed().as_millis()
                 );
+                handle::Handle::notify_profile_update_completed(uid);
             }
             Ok(Err(e)) => logging_error!(Type::Timer, "Failed to update profile uid {}: {}", uid, e),
             Err(_) => logging_error!(Type::Timer, "Timer task timed out for uid: {}", uid),