From 9a3598513b994a26d823e70edb9e19c6a46b5eb9 Mon Sep 17 00:00:00 2001 From: Slinetrac Date: Mon, 27 Oct 2025 13:53:03 +0800 Subject: [PATCH] feat(profile-switch): unify post-switch cleanup handling - workflow.rs (25-427) returns `SwitchWorkflowResult` (success + CleanupHandle) or `SwitchWorkflowError`. All failure/timeout paths stash post-switch work into a single CleanupHandle. Cleanup helpers (`notify_profile_switch_finished` and `close_connections_after_switch`) run inside that task for proper lifetime handling. - driver.rs (29-439) propagates CleanupHandle through `SwitchJobOutcome`, spawns a bridge to wait for completion, and blocks `start_next_job` until done. Direct driver-side panics now schedule failure cleanup via the shared helper. --- src-tauri/src/cmd/profile_switch/driver.rs | 81 +++++++++++++++---- src-tauri/src/cmd/profile_switch/workflow.rs | 85 +++++++++++++------- 2 files changed, 121 insertions(+), 45 deletions(-) diff --git a/src-tauri/src/cmd/profile_switch/driver.rs b/src-tauri/src/cmd/profile_switch/driver.rs index 9cb3f9fe..45028dae 100644 --- a/src-tauri/src/cmd/profile_switch/driver.rs +++ b/src-tauri/src/cmd/profile_switch/driver.rs @@ -51,8 +51,14 @@ enum SwitchDriverMessage { #[derive(Debug)] enum SwitchJobOutcome { - Completed { success: bool }, - Panicked { info: SwitchPanicInfo }, + Completed { + success: bool, + cleanup: workflow::CleanupHandle, + }, + Panicked { + info: SwitchPanicInfo, + cleanup: workflow::CleanupHandle, + }, } pub(super) async fn switch_profile( @@ -253,7 +259,7 @@ fn handle_completion( manager: &'static SwitchManager, ) { match &outcome { - SwitchJobOutcome::Completed { success } => { + SwitchJobOutcome::Completed { success, .. } => { logging!( info, Type::Cmd, @@ -262,7 +268,7 @@ fn handle_completion( success ); } - SwitchJobOutcome::Panicked { info } => { + SwitchJobOutcome::Panicked { info, .. } => { logging!( error, Type::Cmd, @@ -286,8 +292,17 @@ fn handle_completion( state.latest_tokens.remove(request.profile_id()); } - // Schedule cleanup tracking removal once background task finishes. - track_cleanup(state, driver_tx.clone(), request.profile_id().clone()); + let cleanup = match outcome { + SwitchJobOutcome::Completed { cleanup, .. } => cleanup, + SwitchJobOutcome::Panicked { cleanup, .. } => cleanup, + }; + + track_cleanup( + state, + driver_tx.clone(), + request.profile_id().clone(), + cleanup, + ); start_next_job(state, driver_tx, manager); } @@ -338,13 +353,25 @@ fn start_switch_job( tokio::select! { res = workflow_fut.as_mut() => { break match res { - Ok(Ok(success)) => SwitchJobOutcome::Completed { success }, - Ok(Err(info)) => SwitchJobOutcome::Panicked { info }, - Err(payload) => SwitchJobOutcome::Panicked { - info: SwitchPanicInfo::driver_task( - workflow::describe_panic_payload(payload.as_ref()), - ), + Ok(Ok(result)) => SwitchJobOutcome::Completed { + success: result.success, + cleanup: result.cleanup, }, + Ok(Err(error)) => SwitchJobOutcome::Panicked { + info: error.info, + cleanup: error.cleanup, + }, + Err(payload) => { + let info = SwitchPanicInfo::driver_task( + workflow::describe_panic_payload(payload.as_ref()), + ); + let cleanup = workflow::schedule_post_switch_failure( + profile_label.clone(), + completion_request.notify(), + completion_request.task_id(), + ); + SwitchJobOutcome::Panicked { info, cleanup } + } }; } _ = watchdog_interval.tick() => { @@ -391,19 +418,39 @@ fn track_cleanup( state: &mut SwitchDriverState, driver_tx: mpsc::Sender, profile: SmartString, + cleanup: workflow::CleanupHandle, ) { - if state.cleanup_profiles.contains_key(&profile) { - return; + if let Some(existing) = state.cleanup_profiles.remove(&profile) { + existing.abort(); } let profile_clone = profile.clone(); + let driver_clone = driver_tx.clone(); let handle = tokio::spawn(async move { - time::sleep(Duration::from_millis(10)).await; - let _ = driver_tx + let profile_label = profile_clone.clone(); + if let Err(err) = cleanup.await { + logging!( + warn, + Type::Cmd, + "Cleanup task for profile {} failed: {}", + profile_label.as_str(), + err + ); + } + if let Err(err) = driver_clone .send(SwitchDriverMessage::CleanupDone { profile: profile_clone, }) - .await; + .await + { + logging!( + error, + Type::Cmd, + "Failed to push cleanup completion for profile {}: {}", + profile_label.as_str(), + err + ); + } }); state.cleanup_profiles.insert(profile, handle); } diff --git a/src-tauri/src/cmd/profile_switch/workflow.rs b/src-tauri/src/cmd/profile_switch/workflow.rs index 75b3bfdb..10d41110 100644 --- a/src-tauri/src/cmd/profile_switch/workflow.rs +++ b/src-tauri/src/cmd/profile_switch/workflow.rs @@ -22,10 +22,22 @@ mod state_machine; use state_machine::{CONFIG_APPLY_TIMEOUT, SAVE_PROFILES_TIMEOUT, SwitchStateMachine}; pub(super) use state_machine::{SwitchPanicInfo, SwitchStage}; +pub(super) type CleanupHandle = tauri::async_runtime::JoinHandle<()>; + +pub(super) struct SwitchWorkflowResult { + pub success: bool, + pub cleanup: CleanupHandle, +} + +pub(super) struct SwitchWorkflowError { + pub info: SwitchPanicInfo, + pub cleanup: CleanupHandle, +} + pub(super) async fn run_switch_job( manager: &'static SwitchManager, request: SwitchRequest, -) -> Result { +) -> Result { if request.cancel_token().is_cancelled() { logging!( info, @@ -33,12 +45,15 @@ pub(super) async fn run_switch_job( "Switch task {} cancelled before validation", request.task_id() ); - schedule_post_switch_failure( + let cleanup = schedule_post_switch_failure( request.profile_id().clone(), request.notify(), request.task_id(), ); - return Ok(false); + return Ok(SwitchWorkflowResult { + success: false, + cleanup, + }); } let profile_id = request.profile_id().clone(); @@ -55,8 +70,11 @@ pub(super) async fn run_switch_job( err ); handle::Handle::notice_message("config_validate::error", err.clone()); - schedule_post_switch_failure(profile_id.clone(), notify, task_id); - return Ok(false); + let cleanup = schedule_post_switch_failure(profile_id.clone(), notify, task_id); + return Ok(SwitchWorkflowResult { + success: false, + cleanup, + }); } logging!( @@ -101,8 +119,11 @@ pub(super) async fn run_switch_job( "config_validate::error", format!("profile switch timed out: {}", profile_id), ); - schedule_post_switch_failure(profile_id.clone(), notify, task_id); - Ok(false) + let cleanup = schedule_post_switch_failure(profile_id.clone(), notify, task_id); + Ok(SwitchWorkflowResult { + success: false, + cleanup, + }) } Ok(Err(panic_payload)) => { let panic_message = describe_panic_payload(panic_payload.as_ref()); @@ -118,14 +139,18 @@ pub(super) async fn run_switch_job( "config_validate::panic", format!("profile switch panic: {}", profile_id), ); - schedule_post_switch_failure(profile_id.clone(), notify, task_id); - Err(SwitchPanicInfo::workflow_root(panic_message)) + let cleanup = schedule_post_switch_failure(profile_id.clone(), notify, task_id); + Err(SwitchWorkflowError { + info: SwitchPanicInfo::workflow_root(panic_message), + cleanup, + }) } Ok(Ok(machine_result)) => match machine_result { Ok(cmd_result) => match cmd_result { Ok(success) => { - schedule_post_switch_success(profile_id.clone(), success, notify, task_id); - Ok(success) + let cleanup = + schedule_post_switch_success(profile_id.clone(), success, notify, task_id); + Ok(SwitchWorkflowResult { success, cleanup }) } Err(err) => { logging!( @@ -136,8 +161,11 @@ pub(super) async fn run_switch_job( err ); handle::Handle::notice_message("config_validate::error", err.clone()); - schedule_post_switch_failure(profile_id.clone(), notify, task_id); - Ok(false) + let cleanup = schedule_post_switch_failure(profile_id.clone(), notify, task_id); + Ok(SwitchWorkflowResult { + success: false, + cleanup, + }) } }, Err(panic_info) => { @@ -154,8 +182,11 @@ pub(super) async fn run_switch_job( "config_validate::panic", format!("profile switch panic: {}", profile_id), ); - schedule_post_switch_failure(profile_id.clone(), notify, task_id); - Err(panic_info) + let cleanup = schedule_post_switch_failure(profile_id.clone(), notify, task_id); + Err(SwitchWorkflowError { + info: panic_info, + cleanup, + }) } }, } @@ -336,7 +367,7 @@ pub(super) async fn restore_previous_profile(previous: Option) -> C Ok(()) } -async fn close_connections_after_switch(profile_id: &SmartString) { +async fn close_connections_after_switch(profile_id: SmartString) { match time::timeout(SWITCH_CLEANUP_TIMEOUT, async { handle::Handle::mihomo().await.close_all_connections().await }) @@ -364,18 +395,12 @@ async fn close_connections_after_switch(profile_id: &SmartString) { } } -fn schedule_close_connections(profile_id: SmartString) { - AsyncHandler::spawn(|| async move { - close_connections_after_switch(&profile_id).await; - }); -} - fn schedule_post_switch_success( profile_id: SmartString, success: bool, notify: bool, task_id: u64, -) { +) -> CleanupHandle { AsyncHandler::spawn(move || async move { handle::Handle::notify_profile_switch_finished( profile_id.clone(), @@ -386,15 +411,19 @@ fn schedule_post_switch_success( if notify && success { handle::Handle::notice_message("info", "Profile Switched"); } - schedule_close_connections(profile_id); - }); + close_connections_after_switch(profile_id).await; + }) } -fn schedule_post_switch_failure(profile_id: SmartString, notify: bool, task_id: u64) { +pub(super) fn schedule_post_switch_failure( + profile_id: SmartString, + notify: bool, + task_id: u64, +) -> CleanupHandle { AsyncHandler::spawn(move || async move { handle::Handle::notify_profile_switch_finished(profile_id.clone(), false, notify, task_id); - schedule_close_connections(profile_id); - }); + close_connections_after_switch(profile_id).await; + }) } pub(super) fn describe_panic_payload(payload: &(dyn Any + Send)) -> String {