From fdeeacd2df7ef3c3a5ec38a5e5490600260f776b Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Fri, 17 Jul 2026 09:38:42 +0000 Subject: [PATCH 01/10] fix(shutdown): stop wallet backend on GUI close and clear stale shutdown-channel log Await every initialized wallet backend within the bounded GUI shutdown flow, including contexts completed by an in-flight network switch, and consume terminal shutdown receivers exactly once. The platform-wallet coordinator thread join race remains tracked upstream in dashpay/platform#3954 and is not fixed here. Co-Authored-By: OpenAI Codex GPT-5 --- src/app.rs | 120 +++++++++++++++++++++++++++++----------- src/backend_task/mod.rs | 7 ++- 2 files changed, 94 insertions(+), 33 deletions(-) diff --git a/src/app.rs b/src/app.rs index b76cf44ce..f5aede5c4 100644 --- a/src/app.rs +++ b/src/app.rs @@ -682,8 +682,8 @@ pub struct AppState { spv_block: SpvBlockReconciler, /// Data-migration banner + cold-start `FinishUnwire` dispatch reconciler. migration: MigrationReconciler, - /// Async shutdown receiver. `Some` while a graceful shutdown is in progress; - /// the viewport is closed once the receiver resolves. + /// Async shutdown receiver. `Some` until a graceful shutdown reaches a + /// terminal state; the viewport is closed once the receiver resolves. shutdown_receiver: Option>, /// Timestamp when the async shutdown was initiated, used as a hard deadline /// to force-close the viewport if the shutdown task stalls. @@ -1827,6 +1827,60 @@ impl AppState { ); } } + + fn collect_created_context(contexts: &mut Vec>, task_result: TaskResult) { + if let TaskResult::Success { result, .. } = task_result + && let BackendTaskSuccessResult::NetworkContextCreated { context, .. } = *result + { + contexts.push(context); + } + } + + fn start_async_shutdown(&mut self) -> tokio::sync::oneshot::Receiver<()> { + self.subtasks.cancellation_token.cancel(); + let mut contexts = self.network_contexts.values().cloned().collect::>(); + let subtasks = Arc::clone(&self.subtasks); + let (_empty_sender, empty_receiver) = tokiompsc::channel(1); + let mut task_result_receiver = + std::mem::replace(&mut self.task_result_receiver, empty_receiver); + let (tx, rx) = tokio::sync::oneshot::channel(); + + tokio::spawn(async move { + let mut task_shutdown = subtasks.shutdown_async(); + let task_shutdown_completed = loop { + tokio::select! { + result = &mut task_shutdown => break result.is_ok(), + task_result = task_result_receiver.recv() => { + let Some(task_result) = task_result else { + break task_shutdown.await.is_ok(); + }; + Self::collect_created_context(&mut contexts, task_result); + } + } + }; + + while let Ok(task_result) = task_result_receiver.try_recv() { + Self::collect_created_context(&mut contexts, task_result); + } + + let wallet_backends = contexts + .iter() + .filter_map(|context| context.wallet_backend().ok()) + .collect::>(); + futures::future::join_all( + wallet_backends + .iter() + .map(|wallet_backend| wallet_backend.shutdown()), + ) + .await; + + if task_shutdown_completed { + let _ = tx.send(()); + } + }); + + rx + } } impl App for AppState { @@ -1837,37 +1891,41 @@ impl App for AppState { // When the user closes the window we cancel the native close, show a banner, // and start an async shutdown. Once all tasks have finished (or timed out) // we issue Close ourselves. - if let Some(rx) = &mut self.shutdown_receiver { + if self.shutdown_started.is_some() { // Shutdown already in progress — check if it's done. - let should_close = match rx.try_recv() { - Ok(()) => { - tracing::debug!("Async shutdown finished, closing viewport"); - true - } - Err(tokio::sync::oneshot::error::TryRecvError::Closed) => { - // Sender dropped without sending — shutdown task likely panicked. - tracing::warn!("Shutdown channel closed unexpectedly (possible panic)"); - true - } - Err(tokio::sync::oneshot::error::TryRecvError::Empty) => { - // Still waiting — check hard deadline to prevent infinite loop. - if let Some(started) = self.shutdown_started { - let grace = crate::utils::tasks::SHUTDOWN_TIMEOUT - + std::time::Duration::from_secs(5); - if started.elapsed() > grace { - tracing::warn!( - "Shutdown hard deadline exceeded, force-closing viewport" - ); - true + let should_close = match self.shutdown_receiver.as_mut() { + Some(rx) => match rx.try_recv() { + Ok(()) => { + tracing::debug!("Async shutdown finished, closing viewport"); + true + } + Err(tokio::sync::oneshot::error::TryRecvError::Closed) => { + // Sender dropped without sending — shutdown task likely panicked. + tracing::warn!("Shutdown channel closed unexpectedly (possible panic)"); + true + } + Err(tokio::sync::oneshot::error::TryRecvError::Empty) => { + // Still waiting — check hard deadline to prevent infinite loop. + if let Some(started) = self.shutdown_started { + let grace = crate::utils::tasks::SHUTDOWN_TIMEOUT + + std::time::Duration::from_secs(5); + if started.elapsed() > grace { + tracing::warn!( + "Shutdown hard deadline exceeded, force-closing viewport" + ); + true + } else { + false + } } else { false } - } else { - false } - } + }, + None => true, }; if should_close { + self.shutdown_receiver.take(); ctx.send_viewport_cmd(egui::ViewportCommand::Close); } else { ctx.request_repaint(); @@ -1887,7 +1945,7 @@ impl App for AppState { MessageType::Warning, ); tracing::debug!("Close requested, starting async shutdown"); - self.shutdown_receiver = Some(self.subtasks.shutdown_async()); + self.shutdown_receiver = Some(self.start_async_shutdown()); self.shutdown_started = Some(std::time::Instant::now()); ctx.request_repaint(); return; @@ -2431,11 +2489,9 @@ impl App for AppState { // observers (TouchBar, etc.) while views are still alive. crate::platform::order_out_all_windows(); - // If shutdown_receiver is Some, the async shutdown was already initiated - // in update(). Skip the blocking fallback to avoid double-shutdown. - // The blocking path only runs when the window was force-closed without - // going through update() (e.g., OS-level kill, alt-F4 on some platforms). - if self.shutdown_receiver.is_some() { + // The start marker persists after the completion receiver is consumed. + // Only force-closes that bypass update() need the blocking fallback. + if self.shutdown_started.is_some() { tracing::debug!("on_exit: async shutdown was initiated, skipping blocking fallback"); return; } diff --git a/src/backend_task/mod.rs b/src/backend_task/mod.rs index e2565b284..2389a81d7 100644 --- a/src/backend_task/mod.rs +++ b/src/backend_task/mod.rs @@ -899,7 +899,7 @@ impl AppContext { // context, which fast-failed with `WalletBackendNotYetWired` and // reported `spv_started=false`. Wiring first removes that race so // `spv_started` reflects whether sync actually began. - let spv_started = if start_spv { + let spv_started = if start_spv && !self.subtasks.cancellation_token.is_cancelled() { match new_ctx .ensure_wallet_backend_and_start_spv(sender.clone()) .await @@ -916,6 +916,11 @@ impl AppContext { } else { false }; + if self.subtasks.cancellation_token.is_cancelled() + && let Ok(backend) = new_ctx.wallet_backend() + { + backend.shutdown().await; + } Ok(BackendTaskSuccessResult::NetworkContextCreated { network, context: new_ctx, From 80e27aa124e733b415918de86ca5efa5ba43bcb9 Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Fri, 17 Jul 2026 15:46:22 +0000 Subject: [PATCH 02/10] fix(shutdown): make wallet teardown race-safe Close task registration before draining managed work, include the final MCP context, clear remembered secrets, and bound both graceful and fallback wallet backend shutdown paths. Co-Authored-By: OpenAI Codex GPT-5 --- src/app.rs | 238 +++++++++++++++++++++++++++++++++++++++------ src/utils/tasks.rs | 126 ++++++++++++++---------- 2 files changed, 284 insertions(+), 80 deletions(-) diff --git a/src/app.rs b/src/app.rs index f5aede5c4..1e31b965d 100644 --- a/src/app.rs +++ b/src/app.rs @@ -81,6 +81,22 @@ pub const MIGRATION_IDENTITIES_ACK_ACTION_ID: &str = "migration:ack:unreadable_i pub const MIGRATION_UNREADABLE_ACK_ACTION_ID: &str = "migration:ack:unreadable_identities_and_votes"; +const WALLET_BACKEND_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30); +const SHUTDOWN_DEADLINE_MARGIN: Duration = Duration::from_secs(5); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ShutdownOutcome { + Complete, + TaskManagerFailed, + WalletBackendTimedOut, +} + +fn shutdown_hard_deadline() -> Duration { + crate::utils::tasks::SHUTDOWN_TIMEOUT + + WALLET_BACKEND_SHUTDOWN_TIMEOUT + + SHUTDOWN_DEADLINE_MARGIN +} + fn migration_allows_scheduled_vote_sweep(state: &MigrationState) -> bool { matches!( state, @@ -684,10 +700,12 @@ pub struct AppState { migration: MigrationReconciler, /// Async shutdown receiver. `Some` until a graceful shutdown reaches a /// terminal state; the viewport is closed once the receiver resolves. - shutdown_receiver: Option>, + shutdown_receiver: Option>, /// Timestamp when the async shutdown was initiated, used as a hard deadline /// to force-close the viewport if the shutdown task stalls. shutdown_started: Option, + /// True once every managed shutdown stage reached a terminal outcome. + shutdown_finished: bool, /// Platform-level accessibility (AccessKit) activation reconciler. accessibility: AccessibilityActivator, /// Shared MCP context -- follows network switches via `ArcSwap`. @@ -1347,6 +1365,7 @@ impl AppState { migration: MigrationReconciler::new(), shutdown_receiver: None, shutdown_started: None, + shutdown_finished: false, accessibility: AccessibilityActivator::new(accessibility_enforced), #[cfg(feature = "mcp")] mcp_app_context, @@ -1828,25 +1847,77 @@ impl AppState { } } + fn push_unique_context(contexts: &mut Vec>, context: Arc) { + if !contexts + .iter() + .any(|existing| Arc::ptr_eq(existing, &context)) + { + contexts.push(context); + } + } + fn collect_created_context(contexts: &mut Vec>, task_result: TaskResult) { if let TaskResult::Success { result, .. } = task_result && let BackendTaskSuccessResult::NetworkContextCreated { context, .. } = *result { - contexts.push(context); + Self::push_unique_context(contexts, context); } } - fn start_async_shutdown(&mut self) -> tokio::sync::oneshot::Receiver<()> { - self.subtasks.cancellation_token.cancel(); - let mut contexts = self.network_contexts.values().cloned().collect::>(); - let subtasks = Arc::clone(&self.subtasks); + async fn shutdown_wallet_backends(contexts: Vec>) -> ShutdownOutcome { + let mut wallet_backends = Vec::new(); + for context in contexts { + if let Ok(backend) = context.wallet_backend() + && !wallet_backends + .iter() + .any(|existing| Arc::ptr_eq(existing, &backend)) + { + wallet_backends.push(backend); + } + } + + for backend in &wallet_backends { + backend.forget_all_secrets(); + } + + let shutdowns = futures::future::join_all( + wallet_backends + .iter() + .map(|wallet_backend| wallet_backend.shutdown()), + ); + if tokio::time::timeout(WALLET_BACKEND_SHUTDOWN_TIMEOUT, shutdowns) + .await + .is_err() + { + tracing::warn!( + timeout_secs = WALLET_BACKEND_SHUTDOWN_TIMEOUT.as_secs(), + "Wallet backend shutdown timed out; closing with degraded teardown" + ); + ShutdownOutcome::WalletBackendTimedOut + } else { + ShutdownOutcome::Complete + } + } + + fn initial_shutdown_contexts(&self) -> Vec> { + let mut contexts = Vec::new(); + for context in self.network_contexts.values() { + Self::push_unique_context(&mut contexts, Arc::clone(context)); + } + contexts + } + + fn start_async_shutdown(&mut self) -> tokio::sync::oneshot::Receiver { + let mut contexts = self.initial_shutdown_contexts(); + let mut task_shutdown = self.subtasks.shutdown_async(); let (_empty_sender, empty_receiver) = tokiompsc::channel(1); let mut task_result_receiver = std::mem::replace(&mut self.task_result_receiver, empty_receiver); + #[cfg(feature = "mcp")] + let mcp_app_context = self.mcp_app_context.clone(); let (tx, rx) = tokio::sync::oneshot::channel(); tokio::spawn(async move { - let mut task_shutdown = subtasks.shutdown_async(); let task_shutdown_completed = loop { tokio::select! { result = &mut task_shutdown => break result.is_ok(), @@ -1863,24 +1934,70 @@ impl AppState { Self::collect_created_context(&mut contexts, task_result); } - let wallet_backends = contexts - .iter() - .filter_map(|context| context.wallet_backend().ok()) - .collect::>(); - futures::future::join_all( - wallet_backends - .iter() - .map(|wallet_backend| wallet_backend.shutdown()), - ) - .await; - - if task_shutdown_completed { - let _ = tx.send(()); + #[cfg(feature = "mcp")] + if let Some(mcp_app_context) = mcp_app_context { + Self::push_unique_context(&mut contexts, mcp_app_context.load_full()); } + + let wallet_outcome = Self::shutdown_wallet_backends(contexts).await; + let outcome = if task_shutdown_completed { + wallet_outcome + } else { + ShutdownOutcome::TaskManagerFailed + }; + let _ = tx.send(outcome); }); rx } + + fn run_blocking_shutdown_fallback(&mut self) { + const POLL_INTERVAL: Duration = Duration::from_millis(10); + + let mut contexts = self.initial_shutdown_contexts(); + let mut task_shutdown = self.subtasks.shutdown_async(); + let task_shutdown_completed = loop { + while let Ok(task_result) = self.task_result_receiver.try_recv() { + Self::collect_created_context(&mut contexts, task_result); + } + + match task_shutdown.try_recv() { + Ok(()) => break true, + Err(tokio::sync::oneshot::error::TryRecvError::Closed) => break false, + Err(tokio::sync::oneshot::error::TryRecvError::Empty) => { + std::thread::sleep(POLL_INTERVAL); + } + } + }; + + while let Ok(task_result) = self.task_result_receiver.try_recv() { + Self::collect_created_context(&mut contexts, task_result); + } + + #[cfg(feature = "mcp")] + if let Some(mcp_app_context) = &self.mcp_app_context { + Self::push_unique_context(&mut contexts, mcp_app_context.load_full()); + } + + let (tx, rx) = std::sync::mpsc::sync_channel(1); + tokio::spawn(async move { + let outcome = Self::shutdown_wallet_backends(contexts).await; + let _ = tx.send(outcome); + }); + + match rx.recv_timeout(WALLET_BACKEND_SHUTDOWN_TIMEOUT + SHUTDOWN_DEADLINE_MARGIN) { + Ok(ShutdownOutcome::Complete) if task_shutdown_completed => {} + Ok(outcome) => tracing::warn!( + ?outcome, + task_shutdown_completed, + "Blocking shutdown fallback completed with degraded teardown" + ), + Err(error) => tracing::warn!( + ?error, + "Blocking wallet backend shutdown did not report a terminal outcome" + ), + } + } } impl App for AppState { @@ -1895,8 +2012,19 @@ impl App for AppState { // Shutdown already in progress — check if it's done. let should_close = match self.shutdown_receiver.as_mut() { Some(rx) => match rx.try_recv() { - Ok(()) => { - tracing::debug!("Async shutdown finished, closing viewport"); + Ok(outcome) => { + self.shutdown_finished = true; + match outcome { + ShutdownOutcome::Complete => { + tracing::debug!("Async shutdown finished, closing viewport"); + } + ShutdownOutcome::TaskManagerFailed => tracing::warn!( + "Task shutdown failed; closing with degraded teardown" + ), + ShutdownOutcome::WalletBackendTimedOut => tracing::warn!( + "Wallet backend shutdown exceeded its deadline; closing with degraded teardown" + ), + } true } Err(tokio::sync::oneshot::error::TryRecvError::Closed) => { @@ -1907,8 +2035,7 @@ impl App for AppState { Err(tokio::sync::oneshot::error::TryRecvError::Empty) => { // Still waiting — check hard deadline to prevent infinite loop. if let Some(started) = self.shutdown_started { - let grace = crate::utils::tasks::SHUTDOWN_TIMEOUT - + std::time::Duration::from_secs(5); + let grace = shutdown_hard_deadline(); if started.elapsed() > grace { tracing::warn!( "Shutdown hard deadline exceeded, force-closing viewport" @@ -2489,16 +2616,12 @@ impl App for AppState { // observers (TouchBar, etc.) while views are still alive. crate::platform::order_out_all_windows(); - // The start marker persists after the completion receiver is consumed. - // Only force-closes that bypass update() need the blocking fallback. - if self.shutdown_started.is_some() { - tracing::debug!("on_exit: async shutdown was initiated, skipping blocking fallback"); + if self.shutdown_finished { + tracing::debug!("on_exit: async shutdown finished, skipping blocking fallback"); return; } tracing::debug!("on_exit: fallback blocking shutdown"); - if let Err(e) = self.subtasks.shutdown() { - tracing::error!("Error during task shutdown: {}", e); - } + self.run_blocking_shutdown_fallback(); tracing::debug!("App shutdown complete"); } } @@ -2895,3 +3018,56 @@ mod spv_overlay_tests { } } } + +#[cfg(test)] +mod shutdown_tests { + use super::*; + + #[test] + fn collect_created_context_only_keeps_network_context_results() { + let temp_dir = tempfile::tempdir().expect("temp dir"); + let context = crate::context::test_support::test_app_context(temp_dir.path()); + let mut contexts = Vec::new(); + + AppState::collect_created_context( + &mut contexts, + TaskResult::unattributed_success(BackendTaskSuccessResult::NetworkContextCreated { + network: Network::Testnet, + context: Arc::clone(&context), + spv_started: false, + }), + ); + assert_eq!(contexts.len(), 1); + assert!(Arc::ptr_eq(&contexts[0], &context)); + + AppState::collect_created_context( + &mut contexts, + TaskResult::unattributed_success(BackendTaskSuccessResult::NetworkContextCreated { + network: Network::Testnet, + context: Arc::clone(&context), + spv_started: false, + }), + ); + assert_eq!(contexts.len(), 1); + + AppState::collect_created_context( + &mut contexts, + TaskResult::unattributed_success(BackendTaskSuccessResult::None), + ); + assert_eq!(contexts.len(), 1); + + AppState::collect_created_context( + &mut contexts, + TaskResult::unattributed_error(TaskError::NoIdentitiesFound), + ); + assert_eq!(contexts.len(), 1); + } + + #[test] + fn viewport_deadline_covers_task_and_wallet_shutdown_budgets() { + assert!( + shutdown_hard_deadline() + >= crate::utils::tasks::SHUTDOWN_TIMEOUT + WALLET_BACKEND_SHUTDOWN_TIMEOUT + ); + } +} diff --git a/src/utils/tasks.rs b/src/utils/tasks.rs index 14193e83d..c2a8fc449 100644 --- a/src/utils/tasks.rs +++ b/src/utils/tasks.rs @@ -1,4 +1,4 @@ -use std::sync::{Arc, Mutex, atomic::AtomicUsize}; +use std::sync::{Arc, Mutex}; use tokio::time::{Duration, timeout}; use tokio_util::sync::CancellationToken; @@ -8,19 +8,27 @@ pub const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); #[derive(Debug, Clone)] pub struct TaskManager { pub cancellation_token: CancellationToken, // Cancellation token for graceful shutdown - tasks: Arc>>, // Subtasks for graceful shutdown + task_state: Arc>, // Subtasks and their registration barrier active_names: Arc>>, // Names of currently running tasks } +#[derive(Debug)] +struct TaskState { + accepting: bool, + tasks: tokio::task::JoinSet<&'static str>, +} + /// TaskManager tracks spawned subtasks and allows for graceful shutdown of all tasks. impl TaskManager { pub fn new() -> Self { let cancellation_token = CancellationToken::new(); - let subtasks = Arc::new(tokio::sync::Mutex::new(tokio::task::JoinSet::new())); TaskManager { cancellation_token, - tasks: subtasks, + task_state: Arc::new(Mutex::new(TaskState { + accepting: true, + tasks: tokio::task::JoinSet::new(), + })), active_names: Arc::new(Mutex::new(Vec::new())), } } @@ -34,11 +42,19 @@ impl TaskManager { F: std::future::Future + Send + 'static, F::Output: Send + 'static, { - if let Ok(mut names) = self.active_names.lock() { - names.push(name); + let mut state = self.task_state.lock().unwrap_or_else(|e| e.into_inner()); + if !state.accepting { + tracing::debug!(task = name, "Rejected task registration during shutdown"); + return; } - let subtasks = self.tasks.clone(); - tokio::spawn(spawn_subtask(subtasks, name, future)); + state.tasks.spawn(async move { + future.await; + name + }); + self.active_names + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(name); } /// Start an asynchronous graceful shutdown of all subtasks. @@ -47,14 +63,13 @@ impl TaskManager { /// shutdown is complete (or timed out). This does **not** block the calling /// thread, so the UI can keep repainting while tasks wind down. pub fn shutdown_async(&self) -> tokio::sync::oneshot::Receiver<()> { - let cancel = self.cancellation_token.clone(); - let subtasks = self.tasks.clone(); + let tasks = self.begin_shutdown(); let active_names = self.active_names.clone(); let (tx, rx) = tokio::sync::oneshot::channel::<()>(); tokio::task::spawn(async move { - let completed = shutdown_inner(&cancel, &subtasks, &active_names, "async").await; + let completed = shutdown_inner(tasks, &active_names, "async").await; tracing::debug!( "Async shutdown complete, {} subtasks finished cleanly", @@ -74,8 +89,7 @@ impl TaskManager { /// /// This is an equivalent of `Runtime::shutdown_timeout` but for subtasks. pub fn shutdown(&self) -> Result<(), String> { - let cancel = self.cancellation_token.clone(); - let subtasks = self.tasks.clone(); + let tasks = self.begin_shutdown(); let active_names = self.active_names.clone(); // a bit naive synchronization to wait for shutdown @@ -83,7 +97,7 @@ impl TaskManager { // we need to run this task in separate task to avoid cancelling it during shutdown tokio::task::spawn(async move { - let completed = shutdown_inner(&cancel, &subtasks, &active_names, "blocking").await; + let completed = shutdown_inner(tasks, &active_names, "blocking").await; // notify that shutdown is complete if tx.send(completed).is_err() { @@ -107,33 +121,35 @@ impl TaskManager { Ok(()) } + + fn begin_shutdown(&self) -> tokio::task::JoinSet<&'static str> { + let tasks = { + let mut state = self.task_state.lock().unwrap_or_else(|e| e.into_inner()); + state.accepting = false; + std::mem::take(&mut state.tasks) + }; + self.cancellation_token.cancel(); + tasks + } } -/// Shared shutdown logic: cancel all tasks, join with timeout, abort remaining. +/// Join the tasks captured by the registration barrier, aborting on timeout. /// /// Returns the number of tasks that completed cleanly within the timeout. /// The `label` is used in log messages to distinguish async vs blocking callers. async fn shutdown_inner( - cancel: &CancellationToken, - subtasks: &Arc>>, + mut tasks: tokio::task::JoinSet<&'static str>, active_names: &Arc>>, label: &str, ) -> usize { - tracing::trace!("{label}: cancelling all tasks"); - cancel.cancel(); - - let completed = Arc::new(AtomicUsize::new(0)); - - let counter = completed.clone(); - let tasks_list = subtasks.clone(); + let mut completed = 0; let names_for_join = active_names.clone(); - let timed_out = timeout(SHUTDOWN_TIMEOUT, async move { - let mut tasks = tasks_list.lock().await; + let timed_out = timeout(SHUTDOWN_TIMEOUT, async { let total = tasks.len(); tracing::trace!(total, "{label}: joining tasks"); let start = std::time::Instant::now(); while let Some(handle) = tasks.join_next().await { - let i = counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1; + completed += 1; match &handle { Ok(name) => { // Remove one instance of this name from active list @@ -144,14 +160,14 @@ async fn shutdown_inner( } tracing::trace!( task = name, - task_num = i, + task_num = completed, total, elapsed_ms = start.elapsed().as_millis() as u64, "{label}: task joined OK" ); } Err(e) => tracing::trace!( - task_num = i, + task_num = completed, total, elapsed_ms = start.elapsed().as_millis() as u64, error = %e, @@ -163,10 +179,9 @@ async fn shutdown_inner( .await; if timed_out.is_err() { - let done = completed.load(std::sync::atomic::Ordering::Relaxed); let remaining: Vec<&str> = active_names.lock().map(|n| n.clone()).unwrap_or_default(); tracing::trace!( - completed = done, + completed, remaining_count = remaining.len(), remaining_tasks = ?remaining, "{label}: timed out waiting for tasks, aborting remaining" @@ -187,25 +202,9 @@ async fn shutdown_inner( } // Abort all remaining tasks - subtasks.lock().await.shutdown().await; - - completed.load(std::sync::atomic::Ordering::Relaxed) -} + tasks.shutdown().await; -#[inline(always)] -async fn spawn_subtask( - subtasks: Arc>>, - name: &'static str, - future: F, -) where - F: std::future::Future + Send + 'static, - F::Output: Send + 'static, -{ - let mut subtasks_lock = subtasks.lock().await; - subtasks_lock.spawn(async move { - future.await; - name - }); + completed } impl Default for TaskManager { @@ -213,3 +212,32 @@ impl Default for TaskManager { TaskManager::new() } } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, Ordering}; + + #[tokio::test(flavor = "current_thread")] + async fn shutdown_rejects_tasks_submitted_after_the_barrier() { + let manager = TaskManager::new(); + let accepted_task_ran = Arc::new(AtomicBool::new(false)); + let ran = Arc::clone(&accepted_task_ran); + manager.spawn_sync("accepted-task", async move { + ran.store(true, Ordering::Release); + }); + + let shutdown = manager.shutdown_async(); + let late_task_ran = Arc::new(AtomicBool::new(false)); + let ran = Arc::clone(&late_task_ran); + + manager.spawn_sync("late-task", async move { + ran.store(true, Ordering::Release); + }); + + shutdown.await.expect("shutdown completion"); + tokio::task::yield_now().await; + assert!(accepted_task_ran.load(Ordering::Acquire)); + assert!(!late_task_ran.load(Ordering::Acquire)); + } +} From e5ddd8dce6b2ae2e3be2ff0ea67ea8b03cd84ff9 Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Sat, 18 Jul 2026 16:25:58 +0000 Subject: [PATCH 03/10] fix(shutdown): wait for backend-task blocking work before wallet teardown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Backend tasks run inside tokio::spawn_blocking closures wrapped by an abortable async watcher. On shutdown timeout, aborting the watcher only cancelled the cancellable wrapper — the underlying spawn_blocking closure cannot be forcibly stopped and kept running, unobserved, while wallet backend secrets were cleared and storage torn down underneath it. TaskManager now tracks blocking-task completion via a Shared future kept independent of the abortable watcher, so shutdown can genuinely await real completion (bounded by a new blocking-task timeout) instead of an abortable wrapper. Both the async shutdown path and the on_exit force-close fallback route through the same two-phase wait before shutdown_wallet_backends runs. A still-running blocking task past the bound now reports a distinct degraded ShutdownOutcome::BackendTasksTimedOut instead of silently proceeding as if nothing were wrong. shutdown_hard_deadline() extended to cover both phases. Addresses PR #905 review thread PRRT_kwDOM8GK3c6R6sht. Co-Authored-By: OpenAI Codex GPT-5 Co-Authored-By: Claude Sonnet 4.5 --- src/app.rs | 154 ++++++++++++++------------ src/utils/tasks.rs | 270 +++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 335 insertions(+), 89 deletions(-) diff --git a/src/app.rs b/src/app.rs index bd2637a83..849167bca 100644 --- a/src/app.rs +++ b/src/app.rs @@ -38,7 +38,7 @@ use crate::ui::wallets::wallets_screen::WalletsBalancesScreen; use crate::ui::welcome_screen::WelcomeScreen; use crate::ui::{MessageType, RootScreenType, Screen, ScreenLike, ScreenType}; use crate::utils::egui_mpsc::{self, EguiMpscAsync}; -use crate::utils::tasks::TaskManager; +use crate::utils::tasks::{TaskManager, TaskShutdownOutcome}; use crate::wallet_backend::DetScope; use dash_sdk::dpp::dashcore::Network; use dash_sdk::platform::Identifier; @@ -88,11 +88,12 @@ const SHUTDOWN_DEADLINE_MARGIN: Duration = Duration::from_secs(5); enum ShutdownOutcome { Complete, TaskManagerFailed, + BackendTasksTimedOut, WalletBackendTimedOut, } fn shutdown_hard_deadline() -> Duration { - crate::utils::tasks::SHUTDOWN_TIMEOUT + 2 * crate::utils::tasks::SHUTDOWN_TIMEOUT + WALLET_BACKEND_SHUTDOWN_TIMEOUT + SHUTDOWN_DEADLINE_MARGIN } @@ -261,8 +262,13 @@ mod backend_task_join_tests { let sender = SenderAsync::new(tx, egui::Context::default()); let join_handle = tokio::task::spawn_blocking(|| panic!("backend task panic")); - forward_backend_task_join_error(join_handle, sender, None, BackendTaskContext::Unknown) - .await; + forward_backend_task_join_error( + join_handle.await, + sender, + None, + BackendTaskContext::Unknown, + ) + .await; let result = tokio::time::timeout(Duration::from_secs(1), rx.recv()) .await @@ -289,7 +295,7 @@ mod backend_task_join_tests { let join_handle = tokio::task::spawn_blocking(|| panic!("backend task panic")); forward_backend_task_join_error( - join_handle, + join_handle.await, sender, Some(request_id), BackendTaskContext::Unknown, @@ -576,12 +582,12 @@ impl TaskResult { } async fn forward_backend_task_join_error( - join_handle: tokio::task::JoinHandle<()>, + join_result: Result<(), tokio::task::JoinError>, sender: egui_mpsc::SenderAsync, request_id: Option, context: BackendTaskContext, ) { - if let Err(source) = join_handle.await { + if let Err(source) = join_result { let stopped = TaskError::BackendTaskFailed { source: source.into(), }; @@ -1484,25 +1490,27 @@ impl AppState { let watcher_context = context.clone(); let app_context = self.current_app_context().clone(); let handle = tokio::runtime::Handle::current(); - let join_handle = tokio::task::spawn_blocking(move || { - handle.block_on(async move { - let result = app_context.run_backend_task(task, sender.clone()).await; - if let Err(e) = sender - .send(TaskResult::from_backend_task_result(context, result)) - .await - { - tracing::error!("Failed to send task result: {}", e); - } - }); - }); - self.subtasks.spawn_sync( + self.subtasks.spawn_blocking_sync( "backend_task_join_watcher", - forward_backend_task_join_error( - join_handle, - watcher_sender, - request_id, - watcher_context, - ), + move || { + handle.block_on(async move { + let result = app_context.run_backend_task(task, sender.clone()).await; + if let Err(e) = sender + .send(TaskResult::from_backend_task_result(context, result)) + .await + { + tracing::error!("Failed to send task result: {}", e); + } + }); + }, + move |join_result| { + forward_backend_task_join_error( + join_result, + watcher_sender, + request_id, + watcher_context, + ) + }, ); } @@ -1517,39 +1525,41 @@ impl AppState { let app_context = self.current_app_context().clone(); let handle = tokio::runtime::Handle::current(); - let join_handle = tokio::task::spawn_blocking(move || { - handle.block_on(async move { - let results = match mode { - BackendTasksExecutionMode::Sequential => { - app_context - .run_backend_tasks_sequential(tasks, sender.clone()) - .await - } - BackendTasksExecutionMode::Concurrent => { - app_context - .run_backend_tasks_concurrent(tasks, sender.clone()) - .await - } - }; + self.subtasks.spawn_blocking_sync( + "backend_tasks_join_watcher", + move || { + handle.block_on(async move { + let results = match mode { + BackendTasksExecutionMode::Sequential => { + app_context + .run_backend_tasks_sequential(tasks, sender.clone()) + .await + } + BackendTasksExecutionMode::Concurrent => { + app_context + .run_backend_tasks_concurrent(tasks, sender.clone()) + .await + } + }; - for (context, result) in contexts.into_iter().zip(results) { - if let Err(e) = sender - .send(TaskResult::from_backend_task_result(context, result)) - .await - { - tracing::error!("Failed to send task result: {}", e); + for (context, result) in contexts.into_iter().zip(results) { + if let Err(e) = sender + .send(TaskResult::from_backend_task_result(context, result)) + .await + { + tracing::error!("Failed to send task result: {}", e); + } } - } - }); - }); - self.subtasks.spawn_sync( - "backend_tasks_join_watcher", - forward_backend_task_join_error( - join_handle, - watcher_sender, - None, - BackendTaskContext::Unknown, - ), + }); + }, + move |join_result| { + forward_backend_task_join_error( + join_result, + watcher_sender, + None, + BackendTaskContext::Unknown, + ) + }, ); } @@ -1918,12 +1928,12 @@ impl AppState { let (tx, rx) = tokio::sync::oneshot::channel(); tokio::spawn(async move { - let task_shutdown_completed = loop { + let task_shutdown_outcome = loop { tokio::select! { - result = &mut task_shutdown => break result.is_ok(), + result = &mut task_shutdown => break result.ok(), task_result = task_result_receiver.recv() => { let Some(task_result) = task_result else { - break task_shutdown.await.is_ok(); + break task_shutdown.await.ok(); }; Self::collect_created_context(&mut contexts, task_result); } @@ -1940,10 +1950,12 @@ impl AppState { } let wallet_outcome = Self::shutdown_wallet_backends(contexts).await; - let outcome = if task_shutdown_completed { - wallet_outcome - } else { - ShutdownOutcome::TaskManagerFailed + let outcome = match task_shutdown_outcome { + Some(TaskShutdownOutcome::Complete) => wallet_outcome, + Some(TaskShutdownOutcome::BackendTasksTimedOut) => { + ShutdownOutcome::BackendTasksTimedOut + } + None => ShutdownOutcome::TaskManagerFailed, }; let _ = tx.send(outcome); }); @@ -1956,14 +1968,14 @@ impl AppState { let mut contexts = self.initial_shutdown_contexts(); let mut task_shutdown = self.subtasks.shutdown_async(); - let task_shutdown_completed = loop { + let task_shutdown_outcome = loop { while let Ok(task_result) = self.task_result_receiver.try_recv() { Self::collect_created_context(&mut contexts, task_result); } match task_shutdown.try_recv() { - Ok(()) => break true, - Err(tokio::sync::oneshot::error::TryRecvError::Closed) => break false, + Ok(outcome) => break Some(outcome), + Err(tokio::sync::oneshot::error::TryRecvError::Closed) => break None, Err(tokio::sync::oneshot::error::TryRecvError::Empty) => { std::thread::sleep(POLL_INTERVAL); } @@ -1986,10 +1998,11 @@ impl AppState { }); match rx.recv_timeout(WALLET_BACKEND_SHUTDOWN_TIMEOUT + SHUTDOWN_DEADLINE_MARGIN) { - Ok(ShutdownOutcome::Complete) if task_shutdown_completed => {} + Ok(ShutdownOutcome::Complete) + if task_shutdown_outcome == Some(TaskShutdownOutcome::Complete) => {} Ok(outcome) => tracing::warn!( ?outcome, - task_shutdown_completed, + ?task_shutdown_outcome, "Blocking shutdown fallback completed with degraded teardown" ), Err(error) => tracing::warn!( @@ -2021,6 +2034,9 @@ impl App for AppState { ShutdownOutcome::TaskManagerFailed => tracing::warn!( "Task shutdown failed; closing with degraded teardown" ), + ShutdownOutcome::BackendTasksTimedOut => tracing::warn!( + "Backend task blocking work exceeded its deadline; closing with degraded teardown" + ), ShutdownOutcome::WalletBackendTimedOut => tracing::warn!( "Wallet backend shutdown exceeded its deadline; closing with degraded teardown" ), @@ -3069,7 +3085,7 @@ mod shutdown_tests { fn viewport_deadline_covers_task_and_wallet_shutdown_budgets() { assert!( shutdown_hard_deadline() - >= crate::utils::tasks::SHUTDOWN_TIMEOUT + WALLET_BACKEND_SHUTDOWN_TIMEOUT + >= 2 * crate::utils::tasks::SHUTDOWN_TIMEOUT + WALLET_BACKEND_SHUTDOWN_TIMEOUT ); } } diff --git a/src/utils/tasks.rs b/src/utils/tasks.rs index c2a8fc449..b5a2141fb 100644 --- a/src/utils/tasks.rs +++ b/src/utils/tasks.rs @@ -1,3 +1,4 @@ +use futures::FutureExt; use std::sync::{Arc, Mutex}; use tokio::time::{Duration, timeout}; use tokio_util::sync::CancellationToken; @@ -5,6 +6,34 @@ use tokio_util::sync::CancellationToken; /// Timeout duration for graceful shutdown. pub const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); +type BlockingTaskCompletion = futures::future::Shared< + futures::future::BoxFuture<'static, Arc>>>>, +>; + +#[derive(Clone)] +struct TrackedBlockingTask { + name: &'static str, + completion: BlockingTaskCompletion, +} + +impl std::fmt::Debug for TrackedBlockingTask { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_tuple("TrackedBlockingTask") + .field(&self.name) + .finish() + } +} + +/// Terminal state of managed task shutdown. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TaskShutdownOutcome { + /// Ordinary and blocking tasks reached their bounded shutdown points. + Complete, + /// Blocking work was still running when its shutdown wait expired. + BackendTasksTimedOut, +} + #[derive(Debug, Clone)] pub struct TaskManager { pub cancellation_token: CancellationToken, // Cancellation token for graceful shutdown @@ -16,6 +45,13 @@ pub struct TaskManager { struct TaskState { accepting: bool, tasks: tokio::task::JoinSet<&'static str>, + blocking_tasks: Vec, +} + +#[derive(Debug)] +struct ShutdownTasks { + tasks: tokio::task::JoinSet<&'static str>, + blocking_tasks: Vec, } /// TaskManager tracks spawned subtasks and allows for graceful shutdown of all tasks. @@ -28,6 +64,7 @@ impl TaskManager { task_state: Arc::new(Mutex::new(TaskState { accepting: true, tasks: tokio::task::JoinSet::new(), + blocking_tasks: Vec::new(), })), active_names: Arc::new(Mutex::new(Vec::new())), } @@ -57,26 +94,78 @@ impl TaskManager { .push(name); } + /// Spawn blocking work and retain its real task handle through shutdown. + /// + /// The async join observer remains an ordinary abortable subtask, while a + /// separate completion handle lets shutdown await non-cancellable blocking work. + pub fn spawn_blocking_sync(&self, name: &'static str, task: F, on_join: C) + where + F: FnOnce() + Send + 'static, + C: FnOnce(Result<(), tokio::task::JoinError>) -> Fut + Send + 'static, + Fut: std::future::Future + Send + 'static, + { + let mut state = self.task_state.lock().unwrap_or_else(|e| e.into_inner()); + if !state.accepting { + tracing::debug!( + task = name, + "Rejected blocking task registration during shutdown" + ); + return; + } + + let join_handle = tokio::task::spawn_blocking(task); + let completion = async move { Arc::new(Mutex::new(Some(join_handle.await))) } + .boxed() + .shared(); + state.blocking_tasks.push(TrackedBlockingTask { + name, + completion: completion.clone(), + }); + state.tasks.spawn(async move { + let result = completion + .await + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take(); + if let Some(result) = result { + on_join(result).await; + } + name + }); + self.active_names + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(name); + } + /// Start an asynchronous graceful shutdown of all subtasks. /// - /// Cancels all tasks and returns a `oneshot::Receiver` that resolves when - /// shutdown is complete (or timed out). This does **not** block the calling - /// thread, so the UI can keep repainting while tasks wind down. - pub fn shutdown_async(&self) -> tokio::sync::oneshot::Receiver<()> { + /// Cancels ordinary tasks and returns a receiver that resolves after both + /// ordinary and blocking task waits reach a bounded outcome. This does + /// **not** block the calling thread, so the UI can keep repainting. + pub fn shutdown_async(&self) -> tokio::sync::oneshot::Receiver { let tasks = self.begin_shutdown(); let active_names = self.active_names.clone(); - let (tx, rx) = tokio::sync::oneshot::channel::<()>(); + let (tx, rx) = tokio::sync::oneshot::channel::(); tokio::task::spawn(async move { - let completed = shutdown_inner(tasks, &active_names, "async").await; + let (completed, outcome) = shutdown_all_inner( + tasks, + &active_names, + "async", + SHUTDOWN_TIMEOUT, + SHUTDOWN_TIMEOUT, + ) + .await; tracing::debug!( + ?outcome, "Async shutdown complete, {} subtasks finished cleanly", completed ); - let _ = tx.send(()); + let _ = tx.send(outcome); }); rx @@ -84,8 +173,9 @@ impl TaskManager { /// Shutdown all subtasks gracefully (blocking). /// - /// Wait for all subtasks to finish within a specified timeout, and then abort them. - /// Blocks the calling thread. Prefer [`shutdown_async`] when the UI must stay responsive. + /// Ordinary tasks are aborted after [`SHUTDOWN_TIMEOUT`]. Blocking work is + /// then awaited for the same bound and returns an error if still running. + /// Blocks the calling thread. Prefer [`shutdown_async`] for a responsive UI. /// /// This is an equivalent of `Runtime::shutdown_timeout` but for subtasks. pub fn shutdown(&self) -> Result<(), String> { @@ -93,14 +183,21 @@ impl TaskManager { let active_names = self.active_names.clone(); // a bit naive synchronization to wait for shutdown - let (tx, mut rx) = tokio::sync::oneshot::channel::(); + let (tx, mut rx) = tokio::sync::oneshot::channel::<(usize, TaskShutdownOutcome)>(); // we need to run this task in separate task to avoid cancelling it during shutdown tokio::task::spawn(async move { - let completed = shutdown_inner(tasks, &active_names, "blocking").await; + let result = shutdown_all_inner( + tasks, + &active_names, + "blocking", + SHUTDOWN_TIMEOUT, + SHUTDOWN_TIMEOUT, + ) + .await; // notify that shutdown is complete - if tx.send(completed).is_err() { + if tx.send(result).is_err() { tracing::error!("Failed to send shutdown completion signal"); } }); @@ -108,9 +205,11 @@ impl TaskManager { // wait for the shutdown task to finish const WAIT_TIME: Duration = Duration::from_millis(100); let mut completed = 0; - for _ in 0..SHUTDOWN_TIMEOUT.as_millis() / WAIT_TIME.as_millis() { - if let Ok(count) = rx.try_recv() { + let mut outcome = TaskShutdownOutcome::BackendTasksTimedOut; + for _ in 0..(2 * SHUTDOWN_TIMEOUT.as_millis()) / WAIT_TIME.as_millis() { + if let Ok((count, shutdown_outcome)) = rx.try_recv() { completed = count; + outcome = shutdown_outcome; break; } // wait for a short time to avoid busy waiting @@ -119,20 +218,49 @@ impl TaskManager { tracing::debug!("Shutdown complete, {} subtasks finished cleanly", completed); - Ok(()) + match outcome { + TaskShutdownOutcome::Complete => Ok(()), + TaskShutdownOutcome::BackendTasksTimedOut => { + Err("backend task blocking work timed out during shutdown".to_owned()) + } + } } - fn begin_shutdown(&self) -> tokio::task::JoinSet<&'static str> { - let tasks = { + fn begin_shutdown(&self) -> ShutdownTasks { + let (tasks, blocking_tasks) = { let mut state = self.task_state.lock().unwrap_or_else(|e| e.into_inner()); state.accepting = false; - std::mem::take(&mut state.tasks) + ( + std::mem::take(&mut state.tasks), + std::mem::take(&mut state.blocking_tasks), + ) }; self.cancellation_token.cancel(); - tasks + ShutdownTasks { + tasks, + blocking_tasks, + } } } +async fn shutdown_all_inner( + tasks: ShutdownTasks, + active_names: &Arc>>, + label: &str, + task_timeout: Duration, + blocking_task_timeout: Duration, +) -> (usize, TaskShutdownOutcome) { + let completed = shutdown_inner(tasks.tasks, active_names, label, task_timeout).await; + let blocking_tasks_completed = + shutdown_blocking_tasks(tasks.blocking_tasks, label, blocking_task_timeout).await; + let outcome = if blocking_tasks_completed { + TaskShutdownOutcome::Complete + } else { + TaskShutdownOutcome::BackendTasksTimedOut + }; + (completed, outcome) +} + /// Join the tasks captured by the registration barrier, aborting on timeout. /// /// Returns the number of tasks that completed cleanly within the timeout. @@ -141,10 +269,11 @@ async fn shutdown_inner( mut tasks: tokio::task::JoinSet<&'static str>, active_names: &Arc>>, label: &str, + shutdown_timeout: Duration, ) -> usize { let mut completed = 0; let names_for_join = active_names.clone(); - let timed_out = timeout(SHUTDOWN_TIMEOUT, async { + let timed_out = timeout(shutdown_timeout, async { let total = tasks.len(); tracing::trace!(total, "{label}: joining tasks"); let start = std::time::Instant::now(); @@ -207,6 +336,39 @@ async fn shutdown_inner( completed } +async fn shutdown_blocking_tasks( + tasks: Vec, + label: &str, + shutdown_timeout: Duration, +) -> bool { + let total = tasks.len(); + if total == 0 { + return true; + } + + tracing::trace!(total, "{label}: joining backend task blocking work"); + let joined = timeout( + shutdown_timeout, + futures::future::join_all(tasks.into_iter().map(|task| async move { + task.completion.await; + task.name + })), + ) + .await; + + if joined.is_err() { + tracing::warn!( + total, + timeout_secs = shutdown_timeout.as_secs(), + "Backend task blocking work exceeded shutdown wait; continuing with degraded teardown" + ); + false + } else { + tracing::trace!(total, "{label}: backend task blocking work joined"); + true + } +} + impl Default for TaskManager { fn default() -> Self { TaskManager::new() @@ -240,4 +402,72 @@ mod tests { assert!(accepted_task_ran.load(Ordering::Acquire)); assert!(!late_task_ran.load(Ordering::Acquire)); } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn shutdown_waits_for_blocking_work_past_abort_timeout() { + let manager = TaskManager::new(); + let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); + let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + manager.spawn_blocking_sync( + "backend-task", + move || { + started_tx.send(()).expect("report task start"); + release_rx.recv().expect("wait for task release"); + }, + |_| async {}, + ); + started_rx.recv().expect("blocking task started"); + + let teardown_started = Arc::new(AtomicBool::new(false)); + let teardown_flag = Arc::clone(&teardown_started); + let shutdown = manager.shutdown_async(); + let gated_teardown = tokio::spawn(async move { + let outcome = shutdown.await.expect("shutdown completion"); + teardown_flag.store(true, Ordering::Release); + outcome + }); + tokio::time::sleep(SHUTDOWN_TIMEOUT + Duration::from_millis(100)).await; + let teardown_started_while_blocking = teardown_started.load(Ordering::Acquire); + + release_tx.send(()).expect("release blocking task"); + let outcome = tokio::time::timeout(Duration::from_secs(1), gated_teardown) + .await + .expect("shutdown after blocking task release") + .expect("gated teardown task"); + + assert!( + !teardown_started_while_blocking, + "wallet teardown must stay gated while blocking work is running" + ); + assert_eq!(outcome, TaskShutdownOutcome::Complete); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn shutdown_reports_backend_task_timeout_as_degraded() { + let manager = TaskManager::new(); + let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); + let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + manager.spawn_blocking_sync( + "stuck-backend-task", + move || { + started_tx.send(()).expect("report task start"); + release_rx.recv().expect("wait for task release"); + }, + |_| async {}, + ); + started_rx.recv().expect("blocking task started"); + + let outcome = shutdown_all_inner( + manager.begin_shutdown(), + &manager.active_names, + "test", + Duration::from_millis(10), + Duration::from_millis(20), + ) + .await + .1; + release_tx.send(()).expect("release blocking task"); + + assert_eq!(outcome, TaskShutdownOutcome::BackendTasksTimedOut); + } } From 578e74ff5678122616f4c9c728aa36be71d0e2c5 Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Sat, 18 Jul 2026 23:27:05 +0000 Subject: [PATCH 04/10] fix(shutdown): close double-teardown and secret re-caching races found in review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Latch `shutdown_started` at the first close attempt (graceful or force-close) so on_exit()'s guard blocks fallback re-entry on every branch, not only the clean-finish path -- stops a second close signal from racing wallet teardown against a still-running first attempt (SEC-001). - Call forget_all_secrets() again after the bounded wallet-shutdown wait resolves, including the timeout branch, so a backend task that re-caches a secret mid-wait is still cleared before exit (SEC-002). - Register a SwitchNetwork-created wallet backend before its slow network call resolves, and race that call against shutdown's cancellation token, so an in-flight network switch is discoverable to teardown instead of being invisible to it (CODE-003). - Prune completed blocking-task entries, drop dead TaskManager::shutdown(), centralize the shutdown budget constant, and add a CHANGELOG entry (non-blocking cleanup from the same review). Adds regression tests for all three fixes; cargo fmt/clippy/tests all clean (verified independently, not from the fixing agent's own report). Co-Authored-By: Codex Sol 🤖 Co-authored by [Claudius the Magnificent](https://github.com/lklimek/claudius) AI Agent --- CHANGELOG.md | 4 + src/app.rs | 332 ++++++++++++++++++++++++++++++++++------ src/backend_task/mod.rs | 151 ++++++++++++++++-- src/utils/tasks.rs | 175 ++++++++++++--------- 4 files changed, 527 insertions(+), 135 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1eddbec28..2832a1622 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -186,6 +186,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Fixed +- **Closing the app now finishes wallet activity and clears in-memory secrets**: + wallet cleanup still completes if you close the app again while it is already + closing or if a network change is still connecting. + - **Submitted Platform actions are no longer reported as rejected when only confirmation failed**: if a state transition was broadcast but its result could not be confirmed, the app now tells you to check whether it completed diff --git a/src/app.rs b/src/app.rs index 849167bca..78fe114d5 100644 --- a/src/app.rs +++ b/src/app.rs @@ -39,7 +39,7 @@ use crate::ui::welcome_screen::WelcomeScreen; use crate::ui::{MessageType, RootScreenType, Screen, ScreenLike, ScreenType}; use crate::utils::egui_mpsc::{self, EguiMpscAsync}; use crate::utils::tasks::{TaskManager, TaskShutdownOutcome}; -use crate::wallet_backend::DetScope; +use crate::wallet_backend::{DetScope, WalletBackend}; use dash_sdk::dpp::dashcore::Network; use dash_sdk::platform::Identifier; use eframe::{App, egui}; @@ -93,11 +93,26 @@ enum ShutdownOutcome { } fn shutdown_hard_deadline() -> Duration { - 2 * crate::utils::tasks::SHUTDOWN_TIMEOUT + TaskManager::graceful_shutdown_budget() + WALLET_BACKEND_SHUTDOWN_TIMEOUT + SHUTDOWN_DEADLINE_MARGIN } +trait ShutdownWalletBackend: Send + Sync { + fn forget_all_secrets(&self); + fn shutdown(&self) -> futures::future::BoxFuture<'_, ()>; +} + +impl ShutdownWalletBackend for WalletBackend { + fn forget_all_secrets(&self) { + WalletBackend::forget_all_secrets(self); + } + + fn shutdown(&self) -> futures::future::BoxFuture<'_, ()> { + Box::pin(WalletBackend::shutdown(self)) + } +} + fn migration_allows_scheduled_vote_sweep(state: &MigrationState) -> bool { matches!( state, @@ -1867,26 +1882,22 @@ impl AppState { } fn collect_created_context(contexts: &mut Vec>, task_result: TaskResult) { - if let TaskResult::Success { result, .. } = task_result - && let BackendTaskSuccessResult::NetworkContextCreated { context, .. } = *result - { - Self::push_unique_context(contexts, context); - } - } - - async fn shutdown_wallet_backends(contexts: Vec>) -> ShutdownOutcome { - let mut wallet_backends = Vec::new(); - for context in contexts { - if let Ok(backend) = context.wallet_backend() - && !wallet_backends - .iter() - .any(|existing| Arc::ptr_eq(existing, &backend)) - { - wallet_backends.push(backend); + if let TaskResult::Success { result, .. } = task_result { + match *result { + BackendTaskSuccessResult::NetworkContextRegistered { context, .. } + | BackendTaskSuccessResult::NetworkContextCreated { context, .. } => { + Self::push_unique_context(contexts, context); + } + _ => {} } } + } - for backend in &wallet_backends { + async fn shutdown_wallet_backend_instances( + wallet_backends: &[Arc], + shutdown_timeout: Duration, + ) -> ShutdownOutcome { + for backend in wallet_backends { backend.forget_all_secrets(); } @@ -1895,18 +1906,40 @@ impl AppState { .iter() .map(|wallet_backend| wallet_backend.shutdown()), ); - if tokio::time::timeout(WALLET_BACKEND_SHUTDOWN_TIMEOUT, shutdowns) + let outcome = if tokio::time::timeout(shutdown_timeout, shutdowns) .await .is_err() { tracing::warn!( - timeout_secs = WALLET_BACKEND_SHUTDOWN_TIMEOUT.as_secs(), + timeout_secs = shutdown_timeout.as_secs(), "Wallet backend shutdown timed out; closing with degraded teardown" ); ShutdownOutcome::WalletBackendTimedOut } else { ShutdownOutcome::Complete + }; + + for backend in wallet_backends { + backend.forget_all_secrets(); + } + + outcome + } + + async fn shutdown_wallet_backends(contexts: Vec>) -> ShutdownOutcome { + let mut wallet_backends = Vec::new(); + for context in contexts { + if let Ok(backend) = context.wallet_backend() + && !wallet_backends + .iter() + .any(|existing| Arc::ptr_eq(existing, &backend)) + { + wallet_backends.push(backend); + } } + + Self::shutdown_wallet_backend_instances(&wallet_backends, WALLET_BACKEND_SHUTDOWN_TIMEOUT) + .await } fn initial_shutdown_contexts(&self) -> Vec> { @@ -1917,6 +1950,23 @@ impl AppState { contexts } + async fn finish_wallet_shutdown( + mut contexts: Vec>, + mut task_result_receiver: tokiompsc::Receiver, + #[cfg(feature = "mcp")] mcp_app_context: Option>>, + ) -> ShutdownOutcome { + while let Ok(task_result) = task_result_receiver.try_recv() { + Self::collect_created_context(&mut contexts, task_result); + } + + #[cfg(feature = "mcp")] + if let Some(mcp_app_context) = mcp_app_context { + Self::push_unique_context(&mut contexts, mcp_app_context.load_full()); + } + + Self::shutdown_wallet_backends(contexts).await + } + fn start_async_shutdown(&mut self) -> tokio::sync::oneshot::Receiver { let mut contexts = self.initial_shutdown_contexts(); let mut task_shutdown = self.subtasks.shutdown_async(); @@ -1940,16 +1990,13 @@ impl AppState { } }; - while let Ok(task_result) = task_result_receiver.try_recv() { - Self::collect_created_context(&mut contexts, task_result); - } - - #[cfg(feature = "mcp")] - if let Some(mcp_app_context) = mcp_app_context { - Self::push_unique_context(&mut contexts, mcp_app_context.load_full()); - } - - let wallet_outcome = Self::shutdown_wallet_backends(contexts).await; + let wallet_outcome = Self::finish_wallet_shutdown( + contexts, + task_result_receiver, + #[cfg(feature = "mcp")] + mcp_app_context, + ) + .await; let outcome = match task_shutdown_outcome { Some(TaskShutdownOutcome::Complete) => wallet_outcome, Some(TaskShutdownOutcome::BackendTasksTimedOut) => { @@ -1968,8 +2015,13 @@ impl AppState { let mut contexts = self.initial_shutdown_contexts(); let mut task_shutdown = self.subtasks.shutdown_async(); + let (_empty_sender, empty_receiver) = tokiompsc::channel(1); + let mut task_result_receiver = + std::mem::replace(&mut self.task_result_receiver, empty_receiver); + #[cfg(feature = "mcp")] + let mcp_app_context = self.mcp_app_context.clone(); let task_shutdown_outcome = loop { - while let Ok(task_result) = self.task_result_receiver.try_recv() { + while let Ok(task_result) = task_result_receiver.try_recv() { Self::collect_created_context(&mut contexts, task_result); } @@ -1982,18 +2034,15 @@ impl AppState { } }; - while let Ok(task_result) = self.task_result_receiver.try_recv() { - Self::collect_created_context(&mut contexts, task_result); - } - - #[cfg(feature = "mcp")] - if let Some(mcp_app_context) = &self.mcp_app_context { - Self::push_unique_context(&mut contexts, mcp_app_context.load_full()); - } - let (tx, rx) = std::sync::mpsc::sync_channel(1); tokio::spawn(async move { - let outcome = Self::shutdown_wallet_backends(contexts).await; + let outcome = Self::finish_wallet_shutdown( + contexts, + task_result_receiver, + #[cfg(feature = "mcp")] + mcp_app_context, + ) + .await; let _ = tx.send(outcome); }); @@ -2278,6 +2327,10 @@ impl App for AppState { self.network_switch_banner.take_and_clear(); self.finalize_network_switch(network); } + BackendTaskSuccessResult::NetworkContextRegistered { network, context } => { + context.install_secret_prompt(Arc::clone(&self.secret_prompt_host)); + self.network_contexts.entry(network).or_insert(context); + } BackendTaskSuccessResult::PlatformAddressSyncPushed { updates } => { // Coordinator push: populate per-address platform_address_info // for all loaded wallets so the per-address tab stays current @@ -2634,10 +2687,11 @@ impl App for AppState { // observers (TouchBar, etc.) while views are still alive. crate::platform::order_out_all_windows(); - if self.shutdown_finished { - tracing::debug!("on_exit: async shutdown finished, skipping blocking fallback"); + if self.shutdown_started.is_some() || self.shutdown_finished { + tracing::debug!("on_exit: shutdown already attempted, skipping blocking fallback"); return; } + self.shutdown_started = Some(std::time::Instant::now()); tracing::debug!("on_exit: fallback blocking shutdown"); self.run_blocking_shutdown_fallback(); tracing::debug!("App shutdown complete"); @@ -3040,6 +3094,176 @@ mod spv_overlay_tests { #[cfg(test)] mod shutdown_tests { use super::*; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + struct MockShutdownBackend { + forget_calls: AtomicUsize, + secret_cached: AtomicBool, + shutdown_completes: bool, + } + + impl ShutdownWalletBackend for MockShutdownBackend { + fn forget_all_secrets(&self) { + self.forget_calls.fetch_add(1, Ordering::Relaxed); + self.secret_cached.store(false, Ordering::Release); + } + + fn shutdown(&self) -> futures::future::BoxFuture<'_, ()> { + self.secret_cached.store(true, Ordering::Release); + if self.shutdown_completes { + Box::pin(async {}) + } else { + Box::pin(std::future::pending()) + } + } + } + + struct TestDataDir { + prior: Option, + _temp_dir: tempfile::TempDir, + _lock: std::sync::MutexGuard<'static, ()>, + } + + impl TestDataDir { + fn enter() -> Self { + let lock = crate::test_support::DASH_EVO_DATA_DIR_LOCK + .lock() + .unwrap_or_else(|error| error.into_inner()); + let temp_dir = tempfile::tempdir().expect("temporary app data directory"); + let prior = std::env::var("DASH_EVO_DATA_DIR").ok(); + // Safety: the process-global test lock serializes this environment override. + unsafe { std::env::set_var("DASH_EVO_DATA_DIR", temp_dir.path()) }; + Self { + prior, + _temp_dir: temp_dir, + _lock: lock, + } + } + } + + impl Drop for TestDataDir { + fn drop(&mut self) { + // Safety: `_lock` remains held until after the prior value is restored. + unsafe { + match &self.prior { + Some(value) => std::env::set_var("DASH_EVO_DATA_DIR", value), + None => std::env::remove_var("DASH_EVO_DATA_DIR"), + } + } + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn second_exit_after_async_attempt_does_not_repeat_wallet_teardown() { + use crate::wallet_backend::{RememberPolicy, SecretPlaintext, SecretScope}; + use zeroize::Zeroizing; + + let _data_dir = TestDataDir::enter(); + let mut app = AppState::new(egui::Context::default()).expect("test AppState"); + let context = app.current_app_context().clone(); + context + .ensure_wallet_backend(app.task_result_sender.clone()) + .await + .expect("wire test wallet backend"); + let backend = context.wallet_backend().expect("wired wallet backend"); + let secret_access = backend.secret_access(); + let scope = SecretScope::HdSeed { + seed_hash: [0x42; 32], + }; + let secret = Zeroizing::new([0x24; 64]); + secret_access.remember_session( + &scope, + SecretPlaintext::HdSeed(&secret), + RememberPolicy::UntilAppClose, + ); + + let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); + let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + app.subtasks.spawn_blocking_sync( + "slow-shutdown-regression-task", + move || { + started_tx.send(()).expect("report slow task start"); + release_rx.recv().expect("wait for slow task release"); + }, + |_| async {}, + ); + started_rx.recv().expect("slow task started"); + + app.shutdown_started = Some(std::time::Instant::now()); + let mut shutdown = app.start_async_shutdown(); + assert!( + tokio::time::timeout(Duration::from_millis(50), &mut shutdown) + .await + .is_err(), + "wallet teardown waits for in-flight blocking work" + ); + release_tx.send(()).expect("release slow task"); + let outcome = tokio::time::timeout(Duration::from_secs(5), &mut shutdown) + .await + .expect("async shutdown stays within its budget") + .expect("async shutdown task reports an outcome"); + assert_eq!(outcome, ShutdownOutcome::Complete); + assert!(!secret_access.is_session_cached(&scope)); + + secret_access.remember_session( + &scope, + SecretPlaintext::HdSeed(&secret), + RememberPolicy::UntilAppClose, + ); + app.shutdown_finished = false; + app.shutdown_receiver = None; + eframe::App::on_exit(&mut app); + + assert!( + secret_access.is_session_cached(&scope), + "on_exit must not invoke wallet teardown again after any async attempt" + ); + secret_access.forget_all(); + } + + #[tokio::test] + async fn wallet_shutdown_clears_secrets_again_after_teardown_wait() { + let backend = Arc::new(MockShutdownBackend { + forget_calls: AtomicUsize::new(0), + secret_cached: AtomicBool::new(true), + shutdown_completes: true, + }); + + let outcome = AppState::shutdown_wallet_backend_instances( + &[Arc::clone(&backend)], + Duration::from_secs(1), + ) + .await; + + assert_eq!(outcome, ShutdownOutcome::Complete); + assert_eq!(backend.forget_calls.load(Ordering::Relaxed), 2); + assert!( + !backend.secret_cached.load(Ordering::Acquire), + "a secret re-cached during shutdown is cleared by the final pass" + ); + } + + #[tokio::test] + async fn wallet_shutdown_clears_secrets_again_after_teardown_timeout() { + let backend = Arc::new(MockShutdownBackend { + forget_calls: AtomicUsize::new(0), + secret_cached: AtomicBool::new(true), + shutdown_completes: false, + }); + + let outcome = AppState::shutdown_wallet_backend_instances( + &[Arc::clone(&backend)], + Duration::from_millis(10), + ) + .await; + + assert_eq!(outcome, ShutdownOutcome::WalletBackendTimedOut); + assert_eq!(backend.forget_calls.load(Ordering::Relaxed), 2); + assert!( + !backend.secret_cached.load(Ordering::Acquire), + "the final clear still runs when backend shutdown exceeds its timeout" + ); + } #[test] fn collect_created_context_only_keeps_network_context_results() { @@ -3047,6 +3271,16 @@ mod shutdown_tests { let context = crate::context::test_support::test_app_context(temp_dir.path()); let mut contexts = Vec::new(); + AppState::collect_created_context( + &mut contexts, + TaskResult::unattributed_success(BackendTaskSuccessResult::NetworkContextRegistered { + network: Network::Testnet, + context: Arc::clone(&context), + }), + ); + assert_eq!(contexts.len(), 1); + assert!(Arc::ptr_eq(&contexts[0], &context)); + AppState::collect_created_context( &mut contexts, TaskResult::unattributed_success(BackendTaskSuccessResult::NetworkContextCreated { @@ -3082,10 +3316,12 @@ mod shutdown_tests { } #[test] - fn viewport_deadline_covers_task_and_wallet_shutdown_budgets() { - assert!( - shutdown_hard_deadline() - >= 2 * crate::utils::tasks::SHUTDOWN_TIMEOUT + WALLET_BACKEND_SHUTDOWN_TIMEOUT + fn viewport_deadline_is_derived_from_shutdown_phase_budgets() { + assert_eq!( + shutdown_hard_deadline(), + TaskManager::graceful_shutdown_budget() + + WALLET_BACKEND_SHUTDOWN_TIMEOUT + + SHUTDOWN_DEADLINE_MARGIN ); } } diff --git a/src/backend_task/mod.rs b/src/backend_task/mod.rs index 2d4282517..205dcb1c5 100644 --- a/src/backend_task/mod.rs +++ b/src/backend_task/mod.rs @@ -711,6 +711,12 @@ pub enum BackendTaskSuccessResult { CoreClientReinitialized, /// A new network context was created asynchronously during a network switch. + NetworkContextRegistered { + network: Network, + context: Arc, + }, + + /// A new network context finished its asynchronous network switch setup. NetworkContextCreated { network: Network, context: Arc, @@ -996,31 +1002,66 @@ impl AppContext { }) .ok_or(TaskError::NetworkContextCreationFailed { network })?; - // Wire the freshly-built context's wallet backend and then start - // chain sync. The old code called `start_spv()` on an unwired - // context, which fast-failed with `WalletBackendNotYetWired` and - // reported `spv_started=false`. Wiring first removes that race so - // `spv_started` reflects whether sync actually began. - let spv_started = if start_spv && !self.subtasks.cancellation_token.is_cancelled() { - match new_ctx - .ensure_wallet_backend_and_start_spv(sender.clone()) - .await - { - Ok(()) => { - tracing::info!(?network, "SPV started after network switch"); - true + let backend_wired = match new_ctx.ensure_wallet_backend(sender.clone()).await { + Ok(()) => { + if let Err(error) = sender + .send(TaskResult::unattributed_success( + BackendTaskSuccessResult::NetworkContextRegistered { + network, + context: Arc::clone(&new_ctx), + }, + )) + .await + { + tracing::debug!( + ?network, + %error, + "Network switch context registration receiver was unavailable" + ); } - Err(e) => { - tracing::warn!(?network, "SPV start failed after network switch: {e}"); - false + true + } + Err(error) => { + tracing::warn!( + ?network, + %error, + "Wallet backend wiring failed after network switch" + ); + false + } + }; + + let cancellation_token = self.subtasks.cancellation_token.clone(); + let spv_started = if start_spv + && backend_wired + && !cancellation_token.is_cancelled() + { + tokio::select! { + result = new_ctx.ensure_wallet_backend_and_start_spv(sender.clone()) => { + match result { + Ok(()) => { + tracing::info!(?network, "SPV started after network switch"); + true + } + Err(error) => { + tracing::warn!( + ?network, + %error, + "SPV start failed after network switch" + ); + false + } + } } + _ = cancellation_token.cancelled() => false, } } else { false }; - if self.subtasks.cancellation_token.is_cancelled() + if cancellation_token.is_cancelled() && let Ok(backend) = new_ctx.wallet_backend() { + backend.forget_all_secrets(); backend.shutdown().await; } Ok(BackendTaskSuccessResult::NetworkContextCreated { @@ -1168,6 +1209,82 @@ impl AppContext { mod tests { use super::*; + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn switch_network_registers_wired_backend_before_cancellation_teardown() { + use crate::context::test_support::test_app_context; + use crate::wallet_backend::{RememberPolicy, SecretPlaintext, SecretScope}; + use zeroize::Zeroizing; + + let temp_dir = tempfile::tempdir().expect("tempdir"); + let context = test_app_context(temp_dir.path()); + let (tx, mut rx) = tokio::sync::mpsc::channel::(32); + let sender = SenderAsync::new(tx, context.egui_ctx().clone()); + let mut switch = Box::pin(context.run_backend_task( + BackendTask::SwitchNetwork { + network: Network::Mainnet, + start_spv: true, + }, + sender, + )); + + let registered_context = tokio::time::timeout(Duration::from_secs(5), async { + loop { + tokio::select! { + result = rx.recv() => { + let result = result.expect("switch registration result"); + if let TaskResult::Success { result, .. } = result + && let BackendTaskSuccessResult::NetworkContextRegistered { + context, .. + } = *result + { + break context; + } + } + _ = switch.as_mut() => { + panic!("switch completed before early registration"); + } + } + } + }) + .await + .expect("wired switch context is registered before network startup"); + + let backend = registered_context + .wallet_backend() + .expect("registered context has a wired backend"); + let secret_access = backend.secret_access(); + let scope = SecretScope::HdSeed { + seed_hash: [0x5a; 32], + }; + secret_access.remember_session( + &scope, + SecretPlaintext::HdSeed(&Zeroizing::new([0x7b; 64])), + RememberPolicy::UntilAppClose, + ); + assert!(secret_access.is_session_cached(&scope)); + + context.subtasks.cancellation_token.cancel(); + let result = tokio::time::timeout(Duration::from_secs(5), switch.as_mut()) + .await + .expect("cancelled switch task terminates") + .expect("cancelled switch returns its context"); + let BackendTaskSuccessResult::NetworkContextCreated { + context: completed_context, + spv_started, + .. + } = result + else { + panic!("expected completed network context"); + }; + + assert!(Arc::ptr_eq(®istered_context, &completed_context)); + assert!(!spv_started); + assert!( + !secret_access.is_session_cached(&scope), + "cancellation teardown clears secrets on the registered backend" + ); + } + fn dapi_connection_refused_error() -> TaskError { use dash_sdk::Error as SdkError; use dash_sdk::dapi_client::DapiClientError; diff --git a/src/utils/tasks.rs b/src/utils/tasks.rs index b5a2141fb..1d29af129 100644 --- a/src/utils/tasks.rs +++ b/src/utils/tasks.rs @@ -6,10 +6,34 @@ use tokio_util::sync::CancellationToken; /// Timeout duration for graceful shutdown. pub const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); +// `Shared` needs a cloneable output; the mutex-wrapped option lets one observer take `JoinError`. type BlockingTaskCompletion = futures::future::Shared< futures::future::BoxFuture<'static, Arc>>>>, >; +struct ActiveTaskGuard { + active_names: Arc>>, + name: &'static str, +} + +impl ActiveTaskGuard { + fn new(active_names: Arc>>, name: &'static str) -> Self { + Self { active_names, name } + } +} + +impl Drop for ActiveTaskGuard { + fn drop(&mut self) { + let mut names = self + .active_names + .lock() + .unwrap_or_else(|error| error.into_inner()); + if let Some(position) = names.iter().position(|name| *name == self.name) { + names.swap_remove(position); + } + } +} + #[derive(Clone)] struct TrackedBlockingTask { name: &'static str, @@ -84,14 +108,16 @@ impl TaskManager { tracing::debug!(task = name, "Rejected task registration during shutdown"); return; } - state.tasks.spawn(async move { - future.await; - name - }); self.active_names .lock() .unwrap_or_else(|e| e.into_inner()) .push(name); + let active_names = Arc::clone(&self.active_names); + state.tasks.spawn(async move { + let _active_task = ActiveTaskGuard::new(active_names, name); + future.await; + name + }); } /// Spawn blocking work and retain its real task handle through shutdown. @@ -117,11 +143,20 @@ impl TaskManager { let completion = async move { Arc::new(Mutex::new(Some(join_handle.await))) } .boxed() .shared(); + state + .blocking_tasks + .retain(|task| task.completion.peek().is_none()); state.blocking_tasks.push(TrackedBlockingTask { name, completion: completion.clone(), }); + self.active_names + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(name); + let active_names = Arc::clone(&self.active_names); state.tasks.spawn(async move { + let _active_task = ActiveTaskGuard::new(active_names, name); let result = completion .await .lock() @@ -132,10 +167,11 @@ impl TaskManager { } name }); - self.active_names - .lock() - .unwrap_or_else(|e| e.into_inner()) - .push(name); + } + + /// Maximum time used by the ordinary-task and blocking-task shutdown phases. + pub const fn graceful_shutdown_budget() -> Duration { + SHUTDOWN_TIMEOUT.saturating_add(SHUTDOWN_TIMEOUT) } /// Start an asynchronous graceful shutdown of all subtasks. @@ -171,61 +207,6 @@ impl TaskManager { rx } - /// Shutdown all subtasks gracefully (blocking). - /// - /// Ordinary tasks are aborted after [`SHUTDOWN_TIMEOUT`]. Blocking work is - /// then awaited for the same bound and returns an error if still running. - /// Blocks the calling thread. Prefer [`shutdown_async`] for a responsive UI. - /// - /// This is an equivalent of `Runtime::shutdown_timeout` but for subtasks. - pub fn shutdown(&self) -> Result<(), String> { - let tasks = self.begin_shutdown(); - let active_names = self.active_names.clone(); - - // a bit naive synchronization to wait for shutdown - let (tx, mut rx) = tokio::sync::oneshot::channel::<(usize, TaskShutdownOutcome)>(); - - // we need to run this task in separate task to avoid cancelling it during shutdown - tokio::task::spawn(async move { - let result = shutdown_all_inner( - tasks, - &active_names, - "blocking", - SHUTDOWN_TIMEOUT, - SHUTDOWN_TIMEOUT, - ) - .await; - - // notify that shutdown is complete - if tx.send(result).is_err() { - tracing::error!("Failed to send shutdown completion signal"); - } - }); - - // wait for the shutdown task to finish - const WAIT_TIME: Duration = Duration::from_millis(100); - let mut completed = 0; - let mut outcome = TaskShutdownOutcome::BackendTasksTimedOut; - for _ in 0..(2 * SHUTDOWN_TIMEOUT.as_millis()) / WAIT_TIME.as_millis() { - if let Ok((count, shutdown_outcome)) = rx.try_recv() { - completed = count; - outcome = shutdown_outcome; - break; - } - // wait for a short time to avoid busy waiting - std::thread::sleep(WAIT_TIME); - } - - tracing::debug!("Shutdown complete, {} subtasks finished cleanly", completed); - - match outcome { - TaskShutdownOutcome::Complete => Ok(()), - TaskShutdownOutcome::BackendTasksTimedOut => { - Err("backend task blocking work timed out during shutdown".to_owned()) - } - } - } - fn begin_shutdown(&self) -> ShutdownTasks { let (tasks, blocking_tasks) = { let mut state = self.task_state.lock().unwrap_or_else(|e| e.into_inner()); @@ -272,7 +253,6 @@ async fn shutdown_inner( shutdown_timeout: Duration, ) -> usize { let mut completed = 0; - let names_for_join = active_names.clone(); let timed_out = timeout(shutdown_timeout, async { let total = tasks.len(); tracing::trace!(total, "{label}: joining tasks"); @@ -281,12 +261,6 @@ async fn shutdown_inner( completed += 1; match &handle { Ok(name) => { - // Remove one instance of this name from active list - if let Ok(mut names) = names_for_join.lock() - && let Some(pos) = names.iter().position(|n| *n == *name) - { - names.swap_remove(pos); - } tracing::trace!( task = name, task_num = completed, @@ -380,6 +354,67 @@ mod tests { use super::*; use std::sync::atomic::{AtomicBool, Ordering}; + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn completed_blocking_entry_is_pruned_on_next_registration() { + let manager = TaskManager::new(); + let (joined_tx, joined_rx) = tokio::sync::oneshot::channel(); + manager.spawn_blocking_sync( + "completed-backend-task", + || {}, + move |_| async move { + let _ = joined_tx.send(()); + }, + ); + joined_rx.await.expect("completed task observer ran"); + + let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + manager.spawn_blocking_sync( + "pending-backend-task", + move || { + release_rx.recv().expect("wait for task release"); + }, + |_| async {}, + ); + + let tracked = manager + .task_state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .blocking_tasks + .len(); + release_tx.send(()).expect("release pending task"); + + assert_eq!(tracked, 1, "only the pending blocking task stays tracked"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn panicking_on_join_callback_is_removed_from_active_names() { + let manager = TaskManager::new(); + manager.spawn_blocking_sync( + "panicking-on-join", + || {}, + |_| async { panic!("on_join panic for regression coverage") }, + ); + + let _ = shutdown_all_inner( + manager.begin_shutdown(), + &manager.active_names, + "test", + Duration::from_secs(1), + Duration::from_secs(1), + ) + .await; + + assert!( + manager + .active_names + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_empty(), + "completed work is not reported as active when its callback panics" + ); + } + #[tokio::test(flavor = "current_thread")] async fn shutdown_rejects_tasks_submitted_after_the_barrier() { let manager = TaskManager::new(); From a978ff38a1f7dbd939f4131e7c181d7c6b41adc5 Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:10:07 +0000 Subject: [PATCH 05/10] fix: update merged test to new forward_backend_task_join_error signature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The merge of origin/v1.0-dev (904ba0b4) auto-merged src/app.rs cleanly at the text level, but left a semantic break: v1.0-dev's new test `panicking_scheduled_vote_sweep_clears_in_progress_guard` (added in #902) called `forward_backend_task_join_error(join_handle, ...)`, passing the raw `JoinHandle` — the pre-#905 signature. This branch's shutdown-race work changed the function to take an already-awaited `Result<(), JoinError>` (so the join can be observed via `tokio::select!` alongside shutdown), which every other call site in this file was updated for except this one new test that didn't exist when the signature changed. Add the missing `.await`, matching the two sibling tests immediately above and below it. No behavioral change to production code. Co-Authored-By: Claude Opus 4.6 --- src/app.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/app.rs b/src/app.rs index 7f4e5119f..03dee4b9e 100644 --- a/src/app.rs +++ b/src/app.rs @@ -371,7 +371,7 @@ mod backend_task_join_tests { let join_handle = tokio::task::spawn_blocking(|| panic!("scheduled sweep panic")); forward_backend_task_join_error( - join_handle, + join_handle.await, sender, None, BackendTaskContext::ScheduledVoteSweep { network }, From 029060881668078501dc741a87e41dba3f326c1c Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:59:51 +0000 Subject: [PATCH 06/10] fix: correct mislabeled platform PR reference in shutdown TODOs The commented-out ShutdownReport check and the ShieldedShutdownIncomplete bucket TODOs referenced platform-pr3968 (rs-platform-wallet-storage, an unrelated SQLite persistence PR). The actual upstream work that restores these types is the ThreadRegistry/shutdown-report PR, platform#3954. Co-Authored-By: Claude Sonnet 5 --- src/wallet_backend/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/wallet_backend/mod.rs b/src/wallet_backend/mod.rs index 546967902..98fa71bb7 100644 --- a/src/wallet_backend/mod.rs +++ b/src/wallet_backend/mod.rs @@ -1608,7 +1608,7 @@ impl WalletBackend { // flag a still-live worker or orphan, which teardown proceeds past // regardless — log it rather than surface it. // - // TODO(platform-pr3968): `shutdown()` returns `()` at this rev (no + // TODO(platform-pr3954): `shutdown()` returns `()` at this rev (no // clean-shutdown report type yet) — the report check below is // commented out rather than dropped outright; restore once platform // re-adds the report type. User-confirmed removal of the @@ -2756,7 +2756,7 @@ fn map_shielded_op_error(e: platform_wallet::error::PlatformWalletError) -> Task // Every remaining variant → generic WalletBackend wrapper. // - // TODO(platform-pr3968): `ShieldedShutdownIncomplete` doesn't exist on + // TODO(platform-pr3954): `ShieldedShutdownIncomplete` doesn't exist on // `PlatformWalletError` at this rev; it belongs in this bucket once // platform re-adds it. other @ (P::WalletCreation(_) @@ -2986,7 +2986,7 @@ fn identity_op_error_kind(e: &platform_wallet::error::PlatformWalletError) -> Id | P::TransactionBroadcastUnconfirmed(_) | P::ShieldedBroadcastUnconfirmed { .. } | P::ShieldedSpendUnconfirmed { .. } => IdentityOpErrorKind::Other, - // TODO(platform-pr3968): `ShieldedShutdownIncomplete` doesn't exist on + // TODO(platform-pr3954): `ShieldedShutdownIncomplete` doesn't exist on // `PlatformWalletError` at this rev; it belongs in the `Other` bucket // once platform re-adds it. } From 38dbb7895f5148c929e2fdd5b93d293c8e955534 Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Mon, 20 Jul 2026 21:45:02 +0000 Subject: [PATCH 07/10] fix(shutdown): keep wallet teardown behind task barrier Track timed network requests and MCP backend work before wallet teardown, and skip secret clearing whenever managed work fails to drain. Co-Authored-By: OpenAI Codex --- src/app.rs | 129 ++++-- src/backend_task/error.rs | 4 + src/backend_task/mod.rs | 53 ++- src/bin/det_cli/headless.rs | 6 +- src/context/wallet_lifecycle/bootstrap.rs | 9 +- src/context/wallet_lifecycle/registration.rs | 3 +- src/context/wallet_lifecycle/removal.rs | 3 +- src/context/wallet_lifecycle/spv.rs | 3 +- src/context/wallet_lifecycle/unlock.rs | 3 +- src/mcp/dispatch.rs | 122 ++++- src/mcp/server.rs | 27 +- src/utils/tasks.rs | 441 +++++++++++++++---- 12 files changed, 655 insertions(+), 148 deletions(-) diff --git a/src/app.rs b/src/app.rs index 03dee4b9e..2df78325b 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1590,7 +1590,7 @@ impl AppState { reason: BackendInitReason, start_spv: bool, ) { - subtasks.spawn_sync(reason.task_name(), async move { + let _ = subtasks.spawn_sync(reason.task_name(), async move { if start_spv { let already_running = app_ctx .wallet_backend() @@ -1613,7 +1613,7 @@ impl AppState { config: crate::mcp::McpConfig, ) { let cancel = subtasks.cancellation_token.clone(); - subtasks.spawn_sync("mcp-server", async move { + let _ = subtasks.spawn_sync("mcp-server", async move { if let Err(error) = crate::mcp::start_http_server(app_context, config, cancel).await { tracing::error!(%error, "MCP server failed"); } @@ -1637,7 +1637,7 @@ impl AppState { let watcher_context = context.clone(); let app_context = self.current_app_context().clone(); let handle = tokio::runtime::Handle::current(); - self.subtasks.spawn_blocking_sync( + let _ = self.subtasks.spawn_blocking_sync( "backend_task_join_watcher", move || { handle.block_on(async move { @@ -1672,7 +1672,7 @@ impl AppState { let app_context = self.current_app_context().clone(); let handle = tokio::runtime::Handle::current(); - self.subtasks.spawn_blocking_sync( + let _ = self.subtasks.spawn_blocking_sync( "backend_tasks_join_watcher", move || { handle.block_on(async move { @@ -2125,6 +2125,22 @@ impl AppState { Self::shutdown_wallet_backends(contexts).await } + async fn finish_shutdown_after_tasks( + task_shutdown_outcome: Option, + wallet_shutdown: F, + ) -> ShutdownOutcome + where + F: std::future::Future, + { + match task_shutdown_outcome { + Some(TaskShutdownOutcome::Complete) => wallet_shutdown.await, + Some(TaskShutdownOutcome::BackendTasksTimedOut) => { + ShutdownOutcome::BackendTasksTimedOut + } + None => ShutdownOutcome::TaskManagerFailed, + } + } + fn start_async_shutdown(&mut self) -> tokio::sync::oneshot::Receiver { let mut contexts = self.initial_shutdown_contexts(); let mut task_shutdown = self.subtasks.shutdown_async(); @@ -2148,20 +2164,16 @@ impl AppState { } }; - let wallet_outcome = Self::finish_wallet_shutdown( - contexts, - task_result_receiver, - #[cfg(feature = "mcp")] - mcp_app_context, + let outcome = Self::finish_shutdown_after_tasks( + task_shutdown_outcome, + Self::finish_wallet_shutdown( + contexts, + task_result_receiver, + #[cfg(feature = "mcp")] + mcp_app_context, + ), ) .await; - let outcome = match task_shutdown_outcome { - Some(TaskShutdownOutcome::Complete) => wallet_outcome, - Some(TaskShutdownOutcome::BackendTasksTimedOut) => { - ShutdownOutcome::BackendTasksTimedOut - } - None => ShutdownOutcome::TaskManagerFailed, - }; let _ = tx.send(outcome); }); @@ -2194,11 +2206,14 @@ impl AppState { let (tx, rx) = std::sync::mpsc::sync_channel(1); tokio::spawn(async move { - let outcome = Self::finish_wallet_shutdown( - contexts, - task_result_receiver, - #[cfg(feature = "mcp")] - mcp_app_context, + let outcome = Self::finish_shutdown_after_tasks( + task_shutdown_outcome, + Self::finish_wallet_shutdown( + contexts, + task_result_receiver, + #[cfg(feature = "mcp")] + mcp_app_context, + ), ) .await; let _ = tx.send(outcome); @@ -2849,7 +2864,7 @@ impl App for AppState { // only the winner spawns the async teardown. No banner is // needed for a user-initiated stop. if app_ctx.connection_status().begin_spv_stop() { - self.subtasks.spawn_sync("spv_manual_stop", async move { + let _ = self.subtasks.spawn_sync("spv_manual_stop", async move { app_ctx.stop_spv().await; }); } @@ -3443,13 +3458,18 @@ mod shutdown_tests { let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); - app.subtasks.spawn_blocking_sync( - "slow-shutdown-regression-task", - move || { - started_tx.send(()).expect("report slow task start"); - release_rx.recv().expect("wait for slow task release"); - }, - |_| async {}, + assert!( + app.subtasks + .spawn_blocking_sync( + "slow-shutdown-regression-task", + move || { + started_tx.send(()).expect("report slow task start"); + release_rx.recv().expect("wait for slow task release"); + }, + |_| async {}, + ) + .is_ok(), + "slow shutdown task is accepted" ); started_rx.recv().expect("slow task started"); @@ -3529,6 +3549,57 @@ mod shutdown_tests { ); } + #[tokio::test] + async fn backend_task_timeout_skips_wallet_teardown() { + let teardown_calls = Arc::new(AtomicUsize::new(0)); + let calls = Arc::clone(&teardown_calls); + + let outcome = AppState::finish_shutdown_after_tasks( + Some(TaskShutdownOutcome::BackendTasksTimedOut), + async move { + calls.fetch_add(1, Ordering::Relaxed); + ShutdownOutcome::Complete + }, + ) + .await; + + assert_eq!(outcome, ShutdownOutcome::BackendTasksTimedOut); + assert_eq!(teardown_calls.load(Ordering::Relaxed), 0); + } + + #[tokio::test] + async fn task_manager_failure_skips_wallet_teardown() { + let teardown_calls = Arc::new(AtomicUsize::new(0)); + let calls = Arc::clone(&teardown_calls); + + let outcome = AppState::finish_shutdown_after_tasks(None, async move { + calls.fetch_add(1, Ordering::Relaxed); + ShutdownOutcome::Complete + }) + .await; + + assert_eq!(outcome, ShutdownOutcome::TaskManagerFailed); + assert_eq!(teardown_calls.load(Ordering::Relaxed), 0); + } + + #[tokio::test] + async fn complete_task_shutdown_runs_wallet_teardown_once() { + let teardown_calls = Arc::new(AtomicUsize::new(0)); + let calls = Arc::clone(&teardown_calls); + + let outcome = AppState::finish_shutdown_after_tasks( + Some(TaskShutdownOutcome::Complete), + async move { + calls.fetch_add(1, Ordering::Relaxed); + ShutdownOutcome::Complete + }, + ) + .await; + + assert_eq!(outcome, ShutdownOutcome::Complete); + assert_eq!(teardown_calls.load(Ordering::Relaxed), 1); + } + #[test] fn collect_created_context_only_keeps_network_context_results() { let temp_dir = tempfile::tempdir().expect("temp dir"); diff --git a/src/backend_task/error.rs b/src/backend_task/error.rs index 50c343d57..98ad171c0 100644 --- a/src/backend_task/error.rs +++ b/src/backend_task/error.rs @@ -886,6 +886,10 @@ pub enum TaskError { source: BackendTaskJoinError, }, + /// A backend task reached the app after the shutdown admission barrier closed. + #[error("This action could not start because the app is closing. Reopen the app and try again.")] + TaskManagerShuttingDown, + /// DAPI node discovery or address resolution failed. #[error(transparent)] DapiDiscovery(#[from] crate::backend_task::dapi_discovery::DapiDiscoveryError), diff --git a/src/backend_task/mod.rs b/src/backend_task/mod.rs index 2c1a15b0a..bc65f8a36 100644 --- a/src/backend_task/mod.rs +++ b/src/backend_task/mod.rs @@ -88,20 +88,45 @@ pub(crate) async fn await_managed_network_request_with_timeout( where T: Send + 'static, { - let mut task = tokio::spawn(request); - match tokio::time::timeout(timeout_duration, &mut task).await { - Ok(result) => result.map_err(|source| TaskError::BackendTaskFailed { - source: source.into(), - }), - Err(source) => { - task_manager.spawn_sync(reaper_name, async move { - if let Err(source) = task.await { - let error = crate::backend_task::error::BackendTaskJoinError::from(source); - tracing::error!(?error, "Timed-out background request stopped unexpectedly"); - } - }); - Err(timeout_error(source)) - } + let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); + let reply = Arc::new(std::sync::Mutex::new(Some(reply_tx))); + let request_reply = Arc::clone(&reply); + let join_reply = Arc::clone(&reply); + let registration = task_manager.spawn_tracked_sync( + reaper_name, + async move { + let result = request.await; + if let Some(reply) = request_reply + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + { + let _ = reply.send(Ok(result)); + } + }, + move |join_result| async move { + if let Err(source) = join_result + && let Some(reply) = join_reply + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + { + let _ = reply.send(Err(crate::backend_task::error::BackendTaskJoinError::from( + source, + ))); + } + }, + ); + if registration.is_err() { + return Err(TaskError::TaskManagerShuttingDown); + } + drop(reply); + + match tokio::time::timeout(timeout_duration, reply_rx).await { + Ok(Ok(Ok(result))) => Ok(result), + Ok(Ok(Err(source))) => Err(TaskError::BackendTaskFailed { source }), + Ok(Err(_)) => Err(TaskError::TaskManagerShuttingDown), + Err(source) => Err(timeout_error(source)), } } diff --git a/src/bin/det_cli/headless.rs b/src/bin/det_cli/headless.rs index d92153410..5deb40b8b 100644 --- a/src/bin/det_cli/headless.rs +++ b/src/bin/det_cli/headless.rs @@ -9,7 +9,7 @@ /// for the race analysis. pub(super) fn run_headless() -> Result<(), Box> { use dash_evo_tool::logging::initialize_logger; - use dash_evo_tool::mcp::server::init_app_context; + use dash_evo_tool::mcp::server::{init_app_context, shutdown_app_context_wallet_backend}; use dash_evo_tool::mcp::{McpConfig, start_http_server}; // Require MCP_API_KEY -- headless without auth is not allowed. @@ -52,9 +52,7 @@ pub(super) fn run_headless() -> Result<(), Box> { // src/mcp/tools/network.rs), so only the current context needs // draining here. let current_ctx = swappable.load_full(); - if let Ok(backend) = current_ctx.wallet_backend() { - backend.shutdown().await; - } + shutdown_app_context_wallet_backend(¤t_ctx).await; result }); diff --git a/src/context/wallet_lifecycle/bootstrap.rs b/src/context/wallet_lifecycle/bootstrap.rs index 8f20d44c5..90e1bc1ac 100644 --- a/src/context/wallet_lifecycle/bootstrap.rs +++ b/src/context/wallet_lifecycle/bootstrap.rs @@ -442,7 +442,8 @@ impl AppContext { for wallet in open_wallets { let ctx = Arc::clone(self); - self.subtasks + let _ = self + .subtasks .spawn_sync("all_wallets_identity_discovery", async move { if let Err(error) = ctx .discover_identities_gap_limited(&wallet, 0, false, None) @@ -475,7 +476,8 @@ impl AppContext { ) { let ctx = Arc::clone(self); let wallet = Arc::clone(wallet); - self.subtasks + let _ = self + .subtasks .spawn_sync("unlocked_wallet_identity_discovery", async move { ctx.discover_unlocked_wallet_identities(&wallet).await; }); @@ -513,7 +515,8 @@ impl AppContext { ) { let ctx = Arc::clone(self); let wallet_clone = Arc::clone(wallet); - self.subtasks + let _ = self + .subtasks .spawn_sync("wallet_identity_discovery", async move { if let Err(error) = ctx .discover_identities_from_wallet(&wallet_clone, max_identity_index) diff --git a/src/context/wallet_lifecycle/registration.rs b/src/context/wallet_lifecycle/registration.rs index 3d46e82e3..0ab68eaeb 100644 --- a/src/context/wallet_lifecycle/registration.rs +++ b/src/context/wallet_lifecycle/registration.rs @@ -192,7 +192,8 @@ impl AppContext { }; let seed = zeroize::Zeroizing::new(*seed); let birth_height = registration_birth_height(origin); - self.subtasks + let _ = self + .subtasks .spawn_sync("wallet_upstream_registration", async move { if let Err(error) = backend .register_wallet_from_seed(&seed_hash, &seed, birth_height) diff --git a/src/context/wallet_lifecycle/removal.rs b/src/context/wallet_lifecycle/removal.rs index 803355868..b6c9e5625 100644 --- a/src/context/wallet_lifecycle/removal.rs +++ b/src/context/wallet_lifecycle/removal.rs @@ -52,7 +52,8 @@ impl AppContext { // sole async step; it carries no secret, so drive it off-thread. if let Some(wallet_id) = upstream_id { let backend = Arc::clone(&backend); - self.subtasks + let _ = self + .subtasks .spawn_sync("wallet_upstream_removal", async move { if let Err(error) = backend.remove_upstream_wallet(&wallet_id).await { tracing::warn!(%error, "Upstream wallet removal failed"); diff --git a/src/context/wallet_lifecycle/spv.rs b/src/context/wallet_lifecycle/spv.rs index 927916a44..3d1e76d09 100644 --- a/src/context/wallet_lifecycle/spv.rs +++ b/src/context/wallet_lifecycle/spv.rs @@ -40,7 +40,8 @@ impl AppContext { } = backend.forget_all_wallets_local(); for wallet_id in upstream_ids { let backend = Arc::clone(&backend); - self.subtasks + let _ = self + .subtasks .spawn_sync("wallet_upstream_removal", async move { if let Err(error) = backend.remove_upstream_wallet(&wallet_id).await { tracing::warn!(%error, "Upstream wallet removal failed during clear"); diff --git a/src/context/wallet_lifecycle/unlock.rs b/src/context/wallet_lifecycle/unlock.rs index bc5f1b394..9df39abac 100644 --- a/src/context/wallet_lifecycle/unlock.rs +++ b/src/context/wallet_lifecycle/unlock.rs @@ -124,7 +124,8 @@ impl AppContext { ) { let ctx = Arc::clone(self); let wallet = Arc::clone(wallet); - self.subtasks + let _ = self + .subtasks .spawn_sync("wallet_unlock_registration", async move { let lease = lease; ctx.bootstrap_wallet_addresses_jit(&wallet).await; diff --git a/src/mcp/dispatch.rs b/src/mcp/dispatch.rs index deb435cbb..9dacaefd9 100644 --- a/src/mcp/dispatch.rs +++ b/src/mcp/dispatch.rs @@ -6,6 +6,53 @@ use crate::backend_task::{BackendTask, BackendTaskSuccessResult}; use crate::context::AppContext; use std::sync::Arc; +async fn run_blocking_dispatch( + task_manager: Arc, + task: F, +) -> Result +where + T: Send + 'static, + F: FnOnce() -> T + Send + 'static, +{ + let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); + let reply = Arc::new(std::sync::Mutex::new(Some(reply_tx))); + let task_reply = Arc::clone(&reply); + let join_reply = Arc::clone(&reply); + let registration = task_manager.spawn_blocking_sync( + "mcp_backend_task", + move || { + let result = task(); + if let Some(reply) = task_reply + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + { + let _ = reply.send(Ok(result)); + } + }, + move |join_result| async move { + if let Err(source) = join_result + && let Some(reply) = join_reply + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + { + let _ = reply.send(Err(TaskError::BackendTaskFailed { + source: source.into(), + })); + } + }, + ); + if registration.is_err() { + return Err(TaskError::TaskManagerShuttingDown); + } + drop(reply); + + reply_rx + .await + .unwrap_or(Err(TaskError::TaskManagerShuttingDown)) +} + /// Run a single backend task and return its result. /// /// Creates a throwaway channel because `run_backend_task` requires a sender @@ -13,16 +60,16 @@ use std::sync::Arc; /// receiver, so it is intentionally dropped. The task result is returned /// directly from the async call rather than through the channel. /// -/// Uses `spawn_blocking` + `block_on` to avoid `Send` bound issues with -/// platform SDK types (`DataContract`/`Sdk` references across await points). -/// Same pattern as `AppState::handle_backend_task` in `app.rs`. +/// Uses TaskManager-tracked blocking work plus `block_on` to avoid `Send` +/// bound issues with platform SDK types across await points. This matches +/// `AppState::handle_backend_task` and participates in the same shutdown barrier. pub(crate) async fn dispatch_task( app_context: &Arc, task: BackendTask, ) -> Result { let app_context = app_context.clone(); let handle = tokio::runtime::Handle::current(); - tokio::task::spawn_blocking(move || { + run_blocking_dispatch(app_context.subtasks.clone(), move || { handle.block_on(async move { let (tx, _) = tokio::sync::mpsc::channel::(32); let sender = crate::utils::egui_mpsc::SenderAsync::new(tx, egui::Context::default()); @@ -31,3 +78,70 @@ pub(crate) async fn dispatch_task( }) .await? } + +#[cfg(test)] +mod tests { + use super::*; + use crate::utils::tasks::{TaskManager, TaskShutdownOutcome}; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::Duration; + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn shutdown_waits_for_in_flight_mcp_dispatch() { + let task_manager = Arc::new(TaskManager::new()); + let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); + let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + let dispatch_manager = Arc::clone(&task_manager); + let dispatch = tokio::spawn(async move { + run_blocking_dispatch(dispatch_manager, move || { + started_tx.send(()).expect("report dispatch start"); + release_rx.recv().expect("wait for dispatch release"); + 42 + }) + .await + }); + started_rx.recv().expect("dispatch started"); + + let mut shutdown = task_manager.shutdown_async(); + let shutdown_stayed_pending = + tokio::time::timeout(Duration::from_millis(50), &mut shutdown) + .await + .is_err(); + + release_tx.send(()).expect("release dispatch"); + assert_eq!( + dispatch.await.expect("dispatch task").expect("dispatch"), + 42 + ); + assert!( + shutdown_stayed_pending, + "shutdown must remain pending while MCP dispatch is running" + ); + assert_eq!( + shutdown.await.expect("shutdown outcome"), + TaskShutdownOutcome::Complete + ); + } + + #[tokio::test] + async fn dispatch_after_shutdown_is_rejected_without_running() { + let task_manager = Arc::new(TaskManager::new()); + assert_eq!( + task_manager + .shutdown_async() + .await + .expect("shutdown outcome"), + TaskShutdownOutcome::Complete + ); + let ran = Arc::new(AtomicBool::new(false)); + let task_ran = Arc::clone(&ran); + + let result = run_blocking_dispatch(task_manager, move || { + task_ran.store(true, Ordering::Release); + }) + .await; + + assert!(matches!(result, Err(TaskError::TaskManagerShuttingDown))); + assert!(!ran.load(Ordering::Acquire)); + } +} diff --git a/src/mcp/server.rs b/src/mcp/server.rs index 1af2aa387..2aeaf0d10 100644 --- a/src/mcp/server.rs +++ b/src/mcp/server.rs @@ -8,6 +8,27 @@ use rmcp::model::*; use rmcp::{ErrorData as McpError, RoleServer, ServerHandler, service::RequestContext}; use std::sync::Arc; +#[cfg(feature = "cli")] +use crate::utils::tasks::TaskShutdownOutcome; + +/// Drain managed backend work before tearing down a standalone wallet backend. +#[cfg(feature = "cli")] +pub async fn shutdown_app_context_wallet_backend(ctx: &Arc) { + match ctx.subtasks.shutdown_async().await { + Ok(TaskShutdownOutcome::Complete) => { + if let Ok(backend) = ctx.wallet_backend() { + backend.shutdown().await; + } + } + Ok(TaskShutdownOutcome::BackendTasksTimedOut) => { + tracing::warn!("Managed backend work timed out; skipping standalone wallet teardown") + } + Err(_) => { + tracing::warn!("Managed task shutdown failed; skipping standalone wallet teardown") + } + } +} + /// Abstracts how the MCP service stores and swaps its AppContext. /// Both variants support `load` and `store` for network switching. #[derive(Clone)] @@ -209,11 +230,7 @@ impl DashMcpService { #[cfg(feature = "cli")] pub async fn shutdown_wallet_backend(&self) { let Some(ctx) = self.ctx.load() else { return }; - let Ok(backend) = ctx.wallet_backend() else { - return; - }; - // Drain in-flight persister writes. Does not join coordinator threads. - backend.shutdown().await; + shutdown_app_context_wallet_backend(&ctx).await; } /// Build the tool router using trait-based tool composition. diff --git a/src/utils/tasks.rs b/src/utils/tasks.rs index 1d29af129..350c837b9 100644 --- a/src/utils/tasks.rs +++ b/src/utils/tasks.rs @@ -7,7 +7,7 @@ use tokio_util::sync::CancellationToken; pub const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); // `Shared` needs a cloneable output; the mutex-wrapped option lets one observer take `JoinError`. -type BlockingTaskCompletion = futures::future::Shared< +type TaskCompletion = futures::future::Shared< futures::future::BoxFuture<'static, Arc>>>>, >; @@ -35,15 +35,15 @@ impl Drop for ActiveTaskGuard { } #[derive(Clone)] -struct TrackedBlockingTask { +struct TrackedTask { name: &'static str, - completion: BlockingTaskCompletion, + completion: TaskCompletion, } -impl std::fmt::Debug for TrackedBlockingTask { +impl std::fmt::Debug for TrackedTask { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter - .debug_tuple("TrackedBlockingTask") + .debug_tuple("TrackedTask") .field(&self.name) .finish() } @@ -52,9 +52,9 @@ impl std::fmt::Debug for TrackedBlockingTask { /// Terminal state of managed task shutdown. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum TaskShutdownOutcome { - /// Ordinary and blocking tasks reached their bounded shutdown points. + /// Ordinary tasks and tracked joins reached their bounded shutdown points. Complete, - /// Blocking work was still running when its shutdown wait expired. + /// Tracked work was still running when its shutdown wait expired. BackendTasksTimedOut, } @@ -68,14 +68,15 @@ pub struct TaskManager { #[derive(Debug)] struct TaskState { accepting: bool, + tracking_joins: bool, tasks: tokio::task::JoinSet<&'static str>, - blocking_tasks: Vec, + tracked_tasks: Vec, } #[derive(Debug)] struct ShutdownTasks { tasks: tokio::task::JoinSet<&'static str>, - blocking_tasks: Vec, + task_state: Arc>, } /// TaskManager tracks spawned subtasks and allows for graceful shutdown of all tasks. @@ -87,8 +88,9 @@ impl TaskManager { cancellation_token, task_state: Arc::new(Mutex::new(TaskState { accepting: true, + tracking_joins: true, tasks: tokio::task::JoinSet::new(), - blocking_tasks: Vec::new(), + tracked_tasks: Vec::new(), })), active_names: Arc::new(Mutex::new(Vec::new())), } @@ -97,8 +99,9 @@ impl TaskManager { /// Spawn a named future as a subtask, to be used in synchronous context. /// /// The `name` label is logged during shutdown to identify slow tasks. + /// Returns the original future when the shutdown admission barrier is closed. #[inline(always)] - pub fn spawn_sync(&self, name: &'static str, future: F) + pub fn spawn_sync(&self, name: &'static str, future: F) -> Result<(), F> where F: std::future::Future + Send + 'static, F::Output: Send + 'static, @@ -106,7 +109,7 @@ impl TaskManager { let mut state = self.task_state.lock().unwrap_or_else(|e| e.into_inner()); if !state.accepting { tracing::debug!(task = name, "Rejected task registration during shutdown"); - return; + return Err(future); } self.active_names .lock() @@ -118,13 +121,20 @@ impl TaskManager { future.await; name }); + Ok(()) } /// Spawn blocking work and retain its real task handle through shutdown. /// /// The async join observer remains an ordinary abortable subtask, while a /// separate completion handle lets shutdown await non-cancellable blocking work. - pub fn spawn_blocking_sync(&self, name: &'static str, task: F, on_join: C) + /// Returns the work and callback when shutdown has stopped accepting tasks. + pub fn spawn_blocking_sync( + &self, + name: &'static str, + task: F, + on_join: C, + ) -> Result<(), (F, C)> where F: FnOnce() + Send + 'static, C: FnOnce(Result<(), tokio::task::JoinError>) -> Fut + Send + 'static, @@ -136,20 +146,64 @@ impl TaskManager { task = name, "Rejected blocking task registration during shutdown" ); - return; + return Err((task, on_join)); } let join_handle = tokio::task::spawn_blocking(task); - let completion = async move { Arc::new(Mutex::new(Some(join_handle.await))) } - .boxed() - .shared(); - state - .blocking_tasks - .retain(|task| task.completion.peek().is_none()); - state.blocking_tasks.push(TrackedBlockingTask { + let completion = task_completion(join_handle); + self.register_tracked_completion(&mut state, name, completion, on_join); + Ok(()) + } + + /// Spawn async work whose join must survive ordinary-task shutdown. + pub fn spawn_tracked_sync( + &self, + name: &'static str, + task: F, + on_join: C, + ) -> Result<(), (F, C)> + where + F: std::future::Future + Send + 'static, + T: Send + 'static, + C: FnOnce(Result<(), tokio::task::JoinError>) -> Fut + Send + 'static, + Fut: std::future::Future + Send + 'static, + { + let mut state = self.task_state.lock().unwrap_or_else(|e| e.into_inner()); + if !state.tracking_joins { + tracing::debug!(task = name, "Rejected task join tracking after shutdown"); + return Err((task, on_join)); + } + + let completion = task_completion(tokio::spawn(task)); + self.register_tracked_completion(&mut state, name, completion, on_join); + Ok(()) + } + + fn register_tracked_completion( + &self, + state: &mut TaskState, + name: &'static str, + completion: TaskCompletion, + on_join: C, + ) where + C: FnOnce(Result<(), tokio::task::JoinError>) -> Fut + Send + 'static, + Fut: std::future::Future + Send + 'static, + { + state.tracked_tasks.retain(|task| { + task.completion.peek().is_none_or(|completion| { + completion + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_some() + }) + }); + state.tracked_tasks.push(TrackedTask { name, completion: completion.clone(), }); + if !state.accepting { + return; + } self.active_names .lock() .unwrap_or_else(|e| e.into_inner()) @@ -169,7 +223,7 @@ impl TaskManager { }); } - /// Maximum time used by the ordinary-task and blocking-task shutdown phases. + /// Maximum time used by the ordinary-task and tracked-join shutdown phases. pub const fn graceful_shutdown_budget() -> Duration { SHUTDOWN_TIMEOUT.saturating_add(SHUTDOWN_TIMEOUT) } @@ -177,7 +231,7 @@ impl TaskManager { /// Start an asynchronous graceful shutdown of all subtasks. /// /// Cancels ordinary tasks and returns a receiver that resolves after both - /// ordinary and blocking task waits reach a bounded outcome. This does + /// ordinary-task and tracked-join waits reach a bounded outcome. This does /// **not** block the calling thread, so the UI can keep repainting. pub fn shutdown_async(&self) -> tokio::sync::oneshot::Receiver { let tasks = self.begin_shutdown(); @@ -208,22 +262,28 @@ impl TaskManager { } fn begin_shutdown(&self) -> ShutdownTasks { - let (tasks, blocking_tasks) = { + let tasks = { let mut state = self.task_state.lock().unwrap_or_else(|e| e.into_inner()); state.accepting = false; - ( - std::mem::take(&mut state.tasks), - std::mem::take(&mut state.blocking_tasks), - ) + std::mem::take(&mut state.tasks) }; self.cancellation_token.cancel(); ShutdownTasks { tasks, - blocking_tasks, + task_state: Arc::clone(&self.task_state), } } } +fn task_completion(task: tokio::task::JoinHandle) -> TaskCompletion +where + T: Send + 'static, +{ + async move { Arc::new(Mutex::new(Some(task.await.map(|_| ())))) } + .boxed() + .shared() +} + async fn shutdown_all_inner( tasks: ShutdownTasks, active_names: &Arc>>, @@ -232,9 +292,9 @@ async fn shutdown_all_inner( blocking_task_timeout: Duration, ) -> (usize, TaskShutdownOutcome) { let completed = shutdown_inner(tasks.tasks, active_names, label, task_timeout).await; - let blocking_tasks_completed = - shutdown_blocking_tasks(tasks.blocking_tasks, label, blocking_task_timeout).await; - let outcome = if blocking_tasks_completed { + let tracked_tasks_completed = + shutdown_tracked_tasks(tasks.task_state, label, blocking_task_timeout).await; + let outcome = if tracked_tasks_completed { TaskShutdownOutcome::Complete } else { TaskShutdownOutcome::BackendTasksTimedOut @@ -273,7 +333,8 @@ async fn shutdown_inner( task_num = completed, total, elapsed_ms = start.elapsed().as_millis() as u64, - error = %e, + cancelled = e.is_cancelled(), + panicked = e.is_panic(), "{label}: task joined with error" ), } @@ -310,35 +371,65 @@ async fn shutdown_inner( completed } -async fn shutdown_blocking_tasks( - tasks: Vec, +async fn shutdown_tracked_tasks( + task_state: Arc>, label: &str, shutdown_timeout: Duration, ) -> bool { - let total = tasks.len(); - if total == 0 { - return true; - } - - tracing::trace!(total, "{label}: joining backend task blocking work"); - let joined = timeout( - shutdown_timeout, - futures::future::join_all(tasks.into_iter().map(|task| async move { - task.completion.await; - task.name - })), - ) + let mut total = 0; + let joined = timeout(shutdown_timeout, async { + loop { + let tasks = { + let mut state = task_state.lock().unwrap_or_else(|error| error.into_inner()); + if state.tracked_tasks.is_empty() { + state.tracking_joins = false; + break; + } + std::mem::take(&mut state.tracked_tasks) + }; + total += tasks.len(); + tracing::trace!( + batch = tasks.len(), + total, + "{label}: joining tracked task work" + ); + let completions = futures::future::join_all( + tasks + .into_iter() + .map(|task| async move { (task.name, task.completion.await) }), + ) + .await; + for (name, completion) in completions { + let result = completion + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take(); + if let Some(Err(error)) = result { + tracing::warn!( + task = name, + cancelled = error.is_cancelled(), + panicked = error.is_panic(), + "Tracked task stopped unexpectedly after its completion observer ended" + ); + } + } + } + }) .await; if joined.is_err() { + task_state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .tracking_joins = false; tracing::warn!( total, timeout_secs = shutdown_timeout.as_secs(), - "Backend task blocking work exceeded shutdown wait; continuing with degraded teardown" + "Tracked backend work exceeded shutdown wait; continuing with degraded teardown" ); false } else { - tracing::trace!(total, "{label}: backend task blocking work joined"); + tracing::trace!(total, "{label}: tracked task work joined"); true } } @@ -358,29 +449,39 @@ mod tests { async fn completed_blocking_entry_is_pruned_on_next_registration() { let manager = TaskManager::new(); let (joined_tx, joined_rx) = tokio::sync::oneshot::channel(); - manager.spawn_blocking_sync( - "completed-backend-task", - || {}, - move |_| async move { - let _ = joined_tx.send(()); - }, + assert!( + manager + .spawn_blocking_sync( + "completed-backend-task", + || {}, + move |_| async move { + let _ = joined_tx.send(()); + }, + ) + .is_ok(), + "completed task is accepted" ); joined_rx.await.expect("completed task observer ran"); let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); - manager.spawn_blocking_sync( - "pending-backend-task", - move || { - release_rx.recv().expect("wait for task release"); - }, - |_| async {}, + assert!( + manager + .spawn_blocking_sync( + "pending-backend-task", + move || { + release_rx.recv().expect("wait for task release"); + }, + |_| async {}, + ) + .is_ok(), + "pending task is accepted" ); let tracked = manager .task_state .lock() .unwrap_or_else(|error| error.into_inner()) - .blocking_tasks + .tracked_tasks .len(); release_tx.send(()).expect("release pending task"); @@ -390,10 +491,15 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn panicking_on_join_callback_is_removed_from_active_names() { let manager = TaskManager::new(); - manager.spawn_blocking_sync( - "panicking-on-join", - || {}, - |_| async { panic!("on_join panic for regression coverage") }, + assert!( + manager + .spawn_blocking_sync( + "panicking-on-join", + || {}, + |_| async { panic!("on_join panic for regression coverage") }, + ) + .is_ok(), + "panicking callback task is accepted" ); let _ = shutdown_all_inner( @@ -420,15 +526,20 @@ mod tests { let manager = TaskManager::new(); let accepted_task_ran = Arc::new(AtomicBool::new(false)); let ran = Arc::clone(&accepted_task_ran); - manager.spawn_sync("accepted-task", async move { - ran.store(true, Ordering::Release); - }); + assert!( + manager + .spawn_sync("accepted-task", async move { + ran.store(true, Ordering::Release); + }) + .is_ok(), + "task is accepted before shutdown" + ); let shutdown = manager.shutdown_async(); let late_task_ran = Arc::new(AtomicBool::new(false)); let ran = Arc::clone(&late_task_ran); - manager.spawn_sync("late-task", async move { + let late_task = manager.spawn_sync("late-task", async move { ran.store(true, Ordering::Release); }); @@ -436,6 +547,10 @@ mod tests { tokio::task::yield_now().await; assert!(accepted_task_ran.load(Ordering::Acquire)); assert!(!late_task_ran.load(Ordering::Acquire)); + late_task + .expect_err("late task is returned to the caller") + .await; + assert!(late_task_ran.load(Ordering::Acquire)); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -443,13 +558,18 @@ mod tests { let manager = TaskManager::new(); let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); - manager.spawn_blocking_sync( - "backend-task", - move || { - started_tx.send(()).expect("report task start"); - release_rx.recv().expect("wait for task release"); - }, - |_| async {}, + assert!( + manager + .spawn_blocking_sync( + "backend-task", + move || { + started_tx.send(()).expect("report task start"); + release_rx.recv().expect("wait for task release"); + }, + |_| async {}, + ) + .is_ok(), + "backend task is accepted" ); started_rx.recv().expect("blocking task started"); @@ -482,13 +602,18 @@ mod tests { let manager = TaskManager::new(); let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); - manager.spawn_blocking_sync( - "stuck-backend-task", - move || { - started_tx.send(()).expect("report task start"); - release_rx.recv().expect("wait for task release"); - }, - |_| async {}, + assert!( + manager + .spawn_blocking_sync( + "stuck-backend-task", + move || { + started_tx.send(()).expect("report task start"); + release_rx.recv().expect("wait for task release"); + }, + |_| async {}, + ) + .is_ok(), + "stuck backend task is accepted" ); started_rx.recv().expect("blocking task started"); @@ -505,4 +630,150 @@ mod tests { assert_eq!(outcome, TaskShutdownOutcome::BackendTasksTimedOut); } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn request_is_tracked_before_ordinary_shutdown_aborts_its_caller() { + let manager = Arc::new(TaskManager::new()); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (_release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + let (result_tx, result_rx) = tokio::sync::oneshot::channel(); + let task_manager = Arc::clone(&manager); + assert!( + manager + .spawn_sync("managed-request", async move { + let result = crate::backend_task::await_managed_network_request_with_timeout( + task_manager, + "request-reaper", + Duration::from_secs(1), + async move { + let _ = started_tx.send(()); + let _ = release_rx.await; + }, + |source| { + crate::backend_task::error::TaskError::TokenBalanceRefreshTimeout { + source, + } + }, + ) + .await; + let _ = result_tx.send(result); + }) + .is_ok(), + "managed request is accepted" + ); + started_rx.await.expect("request started"); + + let outcome = shutdown_all_inner( + manager.begin_shutdown(), + &manager.active_names, + "test", + Duration::from_millis(10), + Duration::from_millis(25), + ) + .await + .1; + + assert!( + result_rx.await.is_err(), + "ordinary shutdown aborts the request caller before its timeout" + ); + assert_eq!(outcome, TaskShutdownOutcome::BackendTasksTimedOut); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn accepted_timeout_reaper_cannot_yield_complete_while_request_runs() { + let manager = Arc::new(TaskManager::new()); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (_release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + let result = crate::backend_task::await_managed_network_request_with_timeout( + Arc::clone(&manager), + "request-reaper", + Duration::from_millis(10), + async move { + let _ = started_tx.send(()); + let _ = release_rx.await; + }, + |source| crate::backend_task::error::TaskError::TokenBalanceRefreshTimeout { source }, + ); + let (_, result) = tokio::join!(started_rx, result); + assert!(result.is_err()); + + let outcome = shutdown_all_inner( + manager.begin_shutdown(), + &manager.active_names, + "test", + Duration::from_millis(10), + Duration::from_millis(25), + ) + .await + .1; + + assert_eq!(outcome, TaskShutdownOutcome::BackendTasksTimedOut); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn shutdown_consumes_late_blocking_join_error() { + let manager = TaskManager::new(); + let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); + let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + assert!( + manager + .spawn_blocking_sync( + "late-panic", + move || { + started_tx.send(()).expect("report task start"); + release_rx.recv().expect("wait for task release"); + panic!("synthetic panic without secret material"); + }, + |_| async {}, + ) + .is_ok(), + "late panic task is accepted" + ); + started_rx.recv().expect("blocking task started"); + let tasks = manager.begin_shutdown(); + let completion = tasks + .task_state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .tracked_tasks[0] + .completion + .clone(); + + shutdown_inner( + tasks.tasks, + &manager.active_names, + "test", + Duration::from_millis(10), + ) + .await; + release_tx.send(()).expect("release blocking task"); + + let ready = tokio::time::timeout(Duration::from_secs(1), completion.clone()) + .await + .expect("late join error becomes ready"); + assert!( + ready + .lock() + .unwrap_or_else(|error| error.into_inner()) + .as_ref() + .is_some_and(Result::is_err), + "late join error remains available for shutdown" + ); + assert!( + manager + .spawn_tracked_sync("later-tracked-task", async {}, |_| async {}) + .is_ok(), + "tracked registration remains open during the join phase" + ); + + assert!(shutdown_tracked_tasks(tasks.task_state, "test", Duration::from_secs(1)).await); + + let stored = completion + .peek() + .expect("blocking completion") + .lock() + .unwrap_or_else(|error| error.into_inner()); + assert!(stored.is_none(), "shutdown consumes the late join error"); + } } From 26cb9d6dc5270512b1caaae38262963be495ba76 Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Tue, 21 Jul 2026 08:00:45 +0000 Subject: [PATCH 08/10] fix(shutdown): close remaining barrier gaps Keep join-error callbacks and MCP network-switch lifecycle work inside the bounded shutdown barrier. Co-Authored-By: OpenAI Codex GPT-5 --- src/mcp/dispatch.rs | 60 ++++++++++++++++++++++- src/mcp/tools/network.rs | 54 ++++++++++---------- src/utils/tasks.rs | 103 +++++++++++++++++++++++++++++++-------- 3 files changed, 168 insertions(+), 49 deletions(-) diff --git a/src/mcp/dispatch.rs b/src/mcp/dispatch.rs index 9dacaefd9..2b0fa4ab9 100644 --- a/src/mcp/dispatch.rs +++ b/src/mcp/dispatch.rs @@ -67,13 +67,28 @@ pub(crate) async fn dispatch_task( app_context: &Arc, task: BackendTask, ) -> Result { + dispatch_task_with(app_context, task, std::future::ready).await +} + +/// Run a backend task and keep its lifecycle tail inside the shutdown barrier. +pub(crate) async fn dispatch_task_with( + app_context: &Arc, + task: BackendTask, + lifecycle_tail: F, +) -> Result +where + T: Send + 'static, + F: FnOnce(BackendTaskSuccessResult) -> Fut + Send + 'static, + Fut: std::future::Future, +{ let app_context = app_context.clone(); let handle = tokio::runtime::Handle::current(); run_blocking_dispatch(app_context.subtasks.clone(), move || { handle.block_on(async move { let (tx, _) = tokio::sync::mpsc::channel::(32); let sender = crate::utils::egui_mpsc::SenderAsync::new(tx, egui::Context::default()); - app_context.run_backend_task(task, sender).await + let result = app_context.run_backend_task(task, sender).await?; + Ok(lifecycle_tail(result).await) }) }) .await? @@ -144,4 +159,47 @@ mod tests { assert!(matches!(result, Err(TaskError::TaskManagerShuttingDown))); assert!(!ran.load(Ordering::Acquire)); } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn shutdown_waits_for_post_dispatch_lifecycle_tail() { + let temp_dir = tempfile::tempdir().expect("tempdir"); + let app_context = crate::context::test_support::test_app_context(temp_dir.path()); + let task_manager = Arc::clone(&app_context.subtasks); + let dispatch_context = Arc::clone(&app_context); + let (tail_started_tx, tail_started_rx) = tokio::sync::oneshot::channel(); + let (release_tail_tx, release_tail_rx) = tokio::sync::oneshot::channel(); + let dispatch = tokio::spawn(async move { + dispatch_task_with( + &dispatch_context, + BackendTask::None, + move |result| async move { + let _ = tail_started_tx.send(()); + let _ = release_tail_rx.await; + result + }, + ) + .await + }); + tail_started_rx.await.expect("lifecycle tail started"); + + let mut shutdown = task_manager.shutdown_async(); + let shutdown_stayed_pending = + tokio::time::timeout(Duration::from_millis(50), &mut shutdown) + .await + .is_err(); + + release_tail_tx.send(()).expect("release lifecycle tail"); + assert!(matches!( + dispatch.await.expect("dispatch task").expect("dispatch"), + BackendTaskSuccessResult::None + )); + assert!( + shutdown_stayed_pending, + "shutdown must remain pending until post-dispatch lifecycle work finishes" + ); + assert_eq!( + shutdown.await.expect("shutdown outcome"), + TaskShutdownOutcome::Complete + ); + } } diff --git a/src/mcp/tools/network.rs b/src/mcp/tools/network.rs index 34d4cff09..4781af68e 100644 --- a/src/mcp/tools/network.rs +++ b/src/mcp/tools/network.rs @@ -1,6 +1,6 @@ //! Network MCP tools. -use std::borrow::Cow; +use std::{borrow::Cow, sync::Arc}; use rmcp::handler::server::router::tool::{AsyncTool, ToolBase}; use rmcp::model::ToolAnnotations; @@ -8,7 +8,7 @@ use rmcp::schemars; use serde::{Deserialize, Serialize}; use crate::backend_task::{BackendTask, BackendTaskSuccessResult}; -use crate::mcp::dispatch::dispatch_task; +use crate::mcp::dispatch::{dispatch_task, dispatch_task_with}; use crate::mcp::error::McpToolError; use crate::mcp::resolve; use crate::mcp::server::{DashMcpService, collect_available, network_display_name}; @@ -228,34 +228,30 @@ impl AsyncTool for NetworkSwitch { network: target, start_spv: true, }; - let result = dispatch_task(&ctx, task) - .await - .map_err(McpToolError::TaskFailed)?; - - match result { - BackendTaskSuccessResult::NetworkContextCreated { - context, - spv_started, - .. - } => { - // S5: drain the OUTGOING context's wallet backend before - // replacing it. The old `ctx` (still in scope above) is the - // context that is being evicted; the new `context` is the one - // being installed. Draining here — at the swap callsite — - // ensures every network switch leaves no orphaned WalletBackend, - // regardless of how many switches happen in a session. - if let Ok(backend) = ctx.wallet_backend() { - backend.shutdown().await; - } - service.swap_context(context); - Ok(NetworkSwitchOutput { - active: network_display_name(target).to_owned(), + let outgoing_context = Arc::clone(&ctx); + let switch_service = service.clone(); + dispatch_task_with(&ctx, task, move |result| async move { + match result { + BackendTaskSuccessResult::NetworkContextCreated { + context, spv_started, - }) + .. + } => { + if let Ok(backend) = outgoing_context.wallet_backend() { + backend.shutdown().await; + } + switch_service.swap_context(context); + Ok(NetworkSwitchOutput { + active: network_display_name(target).to_owned(), + spv_started, + }) + } + other => Err(McpToolError::Internal(format!( + "Unexpected task result: {other:?}" + ))), } - other => Err(McpToolError::Internal(format!( - "Unexpected task result: {other:?}" - ))), - } + }) + .await + .map_err(McpToolError::TaskFailed)? } } diff --git a/src/utils/tasks.rs b/src/utils/tasks.rs index 350c837b9..e79deaf0a 100644 --- a/src/utils/tasks.rs +++ b/src/utils/tasks.rs @@ -189,6 +189,16 @@ impl TaskManager { C: FnOnce(Result<(), tokio::task::JoinError>) -> Fut + Send + 'static, Fut: std::future::Future + Send + 'static, { + let callback_completion = task_completion(tokio::spawn(async move { + let result = completion + .await + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take(); + if let Some(result) = result { + on_join(result).await; + } + })); state.tracked_tasks.retain(|task| { task.completion.peek().is_none_or(|completion| { completion @@ -199,7 +209,7 @@ impl TaskManager { }); state.tracked_tasks.push(TrackedTask { name, - completion: completion.clone(), + completion: callback_completion.clone(), }); if !state.accepting { return; @@ -211,13 +221,18 @@ impl TaskManager { let active_names = Arc::clone(&self.active_names); state.tasks.spawn(async move { let _active_task = ActiveTaskGuard::new(active_names, name); - let result = completion + let result = callback_completion .await .lock() .unwrap_or_else(|error| error.into_inner()) .take(); - if let Some(result) = result { - on_join(result).await; + if let Some(Err(error)) = result { + tracing::warn!( + task = name, + cancelled = error.is_cancelled(), + panicked = error.is_panic(), + "Tracked task completion callback stopped unexpectedly" + ); } name }); @@ -409,7 +424,7 @@ async fn shutdown_tracked_tasks( task = name, cancelled = error.is_cancelled(), panicked = error.is_panic(), - "Tracked task stopped unexpectedly after its completion observer ended" + "Tracked task completion callback stopped unexpectedly" ); } } @@ -461,7 +476,19 @@ mod tests { .is_ok(), "completed task is accepted" ); - joined_rx.await.expect("completed task observer ran"); + joined_rx.await.expect("completed task callback ran"); + tokio::time::timeout(Duration::from_secs(1), async { + while !manager + .active_names + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_empty() + { + tokio::task::yield_now().await; + } + }) + .await + .expect("completed task observer ran"); let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); assert!( @@ -716,6 +743,8 @@ mod tests { let manager = TaskManager::new(); let (started_tx, started_rx) = std::sync::mpsc::sync_channel(1); let (release_tx, release_rx) = std::sync::mpsc::sync_channel(1); + let join_error_observed = Arc::new(AtomicBool::new(false)); + let observed = Arc::clone(&join_error_observed); assert!( manager .spawn_blocking_sync( @@ -725,7 +754,9 @@ mod tests { release_rx.recv().expect("wait for task release"); panic!("synthetic panic without secret material"); }, - |_| async {}, + move |result| async move { + observed.store(result.is_err(), Ordering::Release); + }, ) .is_ok(), "late panic task is accepted" @@ -749,17 +780,6 @@ mod tests { .await; release_tx.send(()).expect("release blocking task"); - let ready = tokio::time::timeout(Duration::from_secs(1), completion.clone()) - .await - .expect("late join error becomes ready"); - assert!( - ready - .lock() - .unwrap_or_else(|error| error.into_inner()) - .as_ref() - .is_some_and(Result::is_err), - "late join error remains available for shutdown" - ); assert!( manager .spawn_tracked_sync("later-tracked-task", async {}, |_| async {}) @@ -768,12 +788,57 @@ mod tests { ); assert!(shutdown_tracked_tasks(tasks.task_state, "test", Duration::from_secs(1)).await); + assert!( + join_error_observed.load(Ordering::Acquire), + "the tracked callback consumes the late join error" + ); let stored = completion .peek() .expect("blocking completion") .lock() .unwrap_or_else(|error| error.into_inner()); - assert!(stored.is_none(), "shutdown consumes the late join error"); + assert!( + stored.is_none(), + "shutdown consumes the callback completion" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn shutdown_finishes_join_callback_consumed_before_ordinary_abort() { + let manager = TaskManager::new(); + let (callback_started_tx, callback_started_rx) = tokio::sync::oneshot::channel(); + let (release_callback_tx, release_callback_rx) = tokio::sync::oneshot::channel(); + let callback_finished = Arc::new(AtomicBool::new(false)); + let finished = Arc::clone(&callback_finished); + assert!( + manager + .spawn_tracked_sync( + "callback-in-progress", + async { panic!("synthetic panic without secret material") }, + move |result| async move { + assert!(result.is_err(), "task panic reaches its join callback"); + let _ = callback_started_tx.send(()); + let _ = release_callback_rx.await; + finished.store(true, Ordering::Release); + }, + ) + .is_ok(), + "tracked task is accepted" + ); + + callback_started_rx.await.expect("join callback started"); + let tasks = manager.begin_shutdown(); + shutdown_inner(tasks.tasks, &manager.active_names, "test", Duration::ZERO).await; + let _ = release_callback_tx.send(()); + + assert!( + shutdown_tracked_tasks(tasks.task_state, "test", Duration::from_secs(1)).await, + "tracked callback finishes within its shutdown phase" + ); + assert!( + callback_finished.load(Ordering::Acquire), + "tracked shutdown finishes a callback after ordinary observers are aborted" + ); } } From a3a8da6e3dd7bc9cec2449ca0027320b5a0eafdd Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Tue, 21 Jul 2026 10:35:53 +0000 Subject: [PATCH 09/10] fix(app): store the secret prompt host on AppState MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The base-branch merge and this PR's NetworkContextRegistered handler both landed changes to AppState without a textual conflict, but the handler referenced a `self.secret_prompt_host` field that never existed — the Arc built in AppState::new() was only installed on the initial network contexts and then dropped. This broke compilation (E0609) for both the Test Suite and Clippy CI jobs. Store the host as an AppState field so contexts registered later (via an in-flight network switch surviving into NetworkContextRegistered) can also have the secret prompt installed on them. Co-Authored-By: Claude Opus 4.6 --- src/app.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/app.rs b/src/app.rs index bc3a3201a..274bf5012 100644 --- a/src/app.rs +++ b/src/app.rs @@ -834,6 +834,10 @@ pub struct AppState { /// MCP configuration held until a required boot-time network selection succeeds. #[cfg(feature = "mcp")] mcp_server_pending_config: Option, + /// The egui just-in-time secret prompt host, installed on every network + /// context at construction time. Kept here so it can also be installed on + /// contexts created later (e.g. [`BackendTaskSuccessResult::NetworkContextRegistered`]). + secret_prompt_host: Arc, /// Receives just-in-time passphrase requests enqueued by the egui secret /// prompt host. Drained once per frame in [`Self::update`]; the active /// request becomes [`Self::active_secret_prompt`]. @@ -1504,6 +1508,7 @@ impl AppState { mcp_app_context, #[cfg(feature = "mcp")] mcp_server_pending_config, + secret_prompt_host, secret_prompt_receiver, active_secret_prompt: None, prompt_was_blocking: false, From 7603be6c0a531b83d03e59145aabbdf826ae0ecd Mon Sep 17 00:00:00 2001 From: Lukasz Klimek <842586+lklimek@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:54:54 +0000 Subject: [PATCH 10/10] fix(shutdown): forget wallet secrets on degraded shutdown outcomes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two independent reviewers (CodeRabbit and a Claude Code automated review) found that finish_shutdown_after_tasks (src/app.rs) and shutdown_app_context_wallet_backend (src/mcp/server.rs) only cleared session-cached wallet secrets on the Complete outcome. On BackendTasksTimedOut or a failed task manager, the wallet-backend Arc set was collected and then dropped unpolled, so forget_all_secrets() never ran — contradicting this PR's own CHANGELOG entry promising secrets are cleared on close. Split the two concerns: forget_all_secrets() now runs unconditionally on every collected wallet backend regardless of outcome (cheap, synchronous, no dependency on any task finishing), while the full coordinator shutdown() join stays gated on Complete only, since backend tasks may still be using those coordinators on a degraded outcome. Applies to both the async and blocking-fallback GUI shutdown paths, and to the standalone MCP shutdown helper. Adds regression tests proving secrets are forgotten with zero coordinator-shutdown calls on both degraded outcomes, in both src/app.rs::shutdown_tests and the new src/mcp/server.rs::tests. Co-Authored-By: Codex Sol --- src/app.rs | 101 +++++++++++++++++++++++++++++++--------------- src/mcp/server.rs | 100 ++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 164 insertions(+), 37 deletions(-) diff --git a/src/app.rs b/src/app.rs index 274bf5012..e045780ee 100644 --- a/src/app.rs +++ b/src/app.rs @@ -2058,9 +2058,7 @@ impl AppState { wallet_backends: &[Arc], shutdown_timeout: Duration, ) -> ShutdownOutcome { - for backend in wallet_backends { - backend.forget_all_secrets(); - } + Self::forget_wallet_backend_secrets(wallet_backends); let shutdowns = futures::future::join_all( wallet_backends @@ -2080,14 +2078,21 @@ impl AppState { ShutdownOutcome::Complete }; + Self::forget_wallet_backend_secrets(wallet_backends); + + outcome + } + + fn forget_wallet_backend_secrets(wallet_backends: &[Arc]) { for backend in wallet_backends { backend.forget_all_secrets(); } - - outcome } - async fn shutdown_wallet_backends(contexts: Vec>) -> ShutdownOutcome { + async fn shutdown_wallet_backends( + task_shutdown_outcome: Option, + contexts: Vec>, + ) -> ShutdownOutcome { let mut wallet_backends = Vec::new(); for context in contexts { if let Ok(backend) = context.wallet_backend() @@ -2099,8 +2104,15 @@ impl AppState { } } - Self::shutdown_wallet_backend_instances(&wallet_backends, WALLET_BACKEND_SHUTDOWN_TIMEOUT) - .await + Self::finish_shutdown_after_tasks( + task_shutdown_outcome, + &wallet_backends, + Self::shutdown_wallet_backend_instances( + &wallet_backends, + WALLET_BACKEND_SHUTDOWN_TIMEOUT, + ), + ) + .await } fn initial_shutdown_contexts(&self) -> Vec> { @@ -2112,6 +2124,7 @@ impl AppState { } async fn finish_wallet_shutdown( + task_shutdown_outcome: Option, mut contexts: Vec>, mut task_result_receiver: tokiompsc::Receiver, #[cfg(feature = "mcp")] mcp_app_context: Option>>, @@ -2125,16 +2138,20 @@ impl AppState { Self::push_unique_context(&mut contexts, mcp_app_context.load_full()); } - Self::shutdown_wallet_backends(contexts).await + Self::shutdown_wallet_backends(task_shutdown_outcome, contexts).await } - async fn finish_shutdown_after_tasks( + async fn finish_shutdown_after_tasks( task_shutdown_outcome: Option, + wallet_backends: &[Arc], wallet_shutdown: F, ) -> ShutdownOutcome where + B: ShutdownWalletBackend, F: std::future::Future, { + Self::forget_wallet_backend_secrets(wallet_backends); + match task_shutdown_outcome { Some(TaskShutdownOutcome::Complete) => wallet_shutdown.await, Some(TaskShutdownOutcome::BackendTasksTimedOut) => { @@ -2167,14 +2184,12 @@ impl AppState { } }; - let outcome = Self::finish_shutdown_after_tasks( + let outcome = Self::finish_wallet_shutdown( task_shutdown_outcome, - Self::finish_wallet_shutdown( - contexts, - task_result_receiver, - #[cfg(feature = "mcp")] - mcp_app_context, - ), + contexts, + task_result_receiver, + #[cfg(feature = "mcp")] + mcp_app_context, ) .await; let _ = tx.send(outcome); @@ -2209,14 +2224,12 @@ impl AppState { let (tx, rx) = std::sync::mpsc::sync_channel(1); tokio::spawn(async move { - let outcome = Self::finish_shutdown_after_tasks( + let outcome = Self::finish_wallet_shutdown( task_shutdown_outcome, - Self::finish_wallet_shutdown( - contexts, - task_result_receiver, - #[cfg(feature = "mcp")] - mcp_app_context, - ), + contexts, + task_result_receiver, + #[cfg(feature = "mcp")] + mcp_app_context, ) .await; let _ = tx.send(outcome); @@ -3550,12 +3563,18 @@ mod shutdown_tests { } #[tokio::test] - async fn backend_task_timeout_skips_wallet_teardown() { + async fn backend_task_timeout_forgets_secrets_without_wallet_teardown() { + let backend = Arc::new(MockShutdownBackend { + forget_calls: AtomicUsize::new(0), + secret_cached: AtomicBool::new(true), + shutdown_completes: true, + }); let teardown_calls = Arc::new(AtomicUsize::new(0)); let calls = Arc::clone(&teardown_calls); let outcome = AppState::finish_shutdown_after_tasks( Some(TaskShutdownOutcome::BackendTasksTimedOut), + &[Arc::clone(&backend)], async move { calls.fetch_add(1, Ordering::Relaxed); ShutdownOutcome::Complete @@ -3564,31 +3583,47 @@ mod shutdown_tests { .await; assert_eq!(outcome, ShutdownOutcome::BackendTasksTimedOut); + assert_eq!(backend.forget_calls.load(Ordering::Relaxed), 1); + assert!(!backend.secret_cached.load(Ordering::Acquire)); assert_eq!(teardown_calls.load(Ordering::Relaxed), 0); } #[tokio::test] - async fn task_manager_failure_skips_wallet_teardown() { + async fn task_manager_failure_forgets_secrets_without_wallet_teardown() { + let backend = Arc::new(MockShutdownBackend { + forget_calls: AtomicUsize::new(0), + secret_cached: AtomicBool::new(true), + shutdown_completes: true, + }); let teardown_calls = Arc::new(AtomicUsize::new(0)); let calls = Arc::clone(&teardown_calls); - let outcome = AppState::finish_shutdown_after_tasks(None, async move { - calls.fetch_add(1, Ordering::Relaxed); - ShutdownOutcome::Complete - }) - .await; + let outcome = + AppState::finish_shutdown_after_tasks(None, &[Arc::clone(&backend)], async move { + calls.fetch_add(1, Ordering::Relaxed); + ShutdownOutcome::Complete + }) + .await; assert_eq!(outcome, ShutdownOutcome::TaskManagerFailed); + assert_eq!(backend.forget_calls.load(Ordering::Relaxed), 1); + assert!(!backend.secret_cached.load(Ordering::Acquire)); assert_eq!(teardown_calls.load(Ordering::Relaxed), 0); } #[tokio::test] - async fn complete_task_shutdown_runs_wallet_teardown_once() { + async fn complete_task_shutdown_forgets_secrets_and_runs_wallet_teardown_once() { + let backend = Arc::new(MockShutdownBackend { + forget_calls: AtomicUsize::new(0), + secret_cached: AtomicBool::new(true), + shutdown_completes: true, + }); let teardown_calls = Arc::new(AtomicUsize::new(0)); let calls = Arc::clone(&teardown_calls); let outcome = AppState::finish_shutdown_after_tasks( Some(TaskShutdownOutcome::Complete), + &[Arc::clone(&backend)], async move { calls.fetch_add(1, Ordering::Relaxed); ShutdownOutcome::Complete @@ -3597,6 +3632,8 @@ mod shutdown_tests { .await; assert_eq!(outcome, ShutdownOutcome::Complete); + assert_eq!(backend.forget_calls.load(Ordering::Relaxed), 1); + assert!(!backend.secret_cached.load(Ordering::Acquire)); assert_eq!(teardown_calls.load(Ordering::Relaxed), 1); } diff --git a/src/mcp/server.rs b/src/mcp/server.rs index 2aeaf0d10..d294654d6 100644 --- a/src/mcp/server.rs +++ b/src/mcp/server.rs @@ -11,19 +11,51 @@ use std::sync::Arc; #[cfg(feature = "cli")] use crate::utils::tasks::TaskShutdownOutcome; +#[cfg(feature = "cli")] +trait StandaloneShutdownWalletBackend: Send + Sync { + fn forget_all_secrets(&self); + fn shutdown(&self) -> futures::future::BoxFuture<'_, ()>; +} + +#[cfg(feature = "cli")] +impl StandaloneShutdownWalletBackend for crate::wallet_backend::WalletBackend { + fn forget_all_secrets(&self) { + crate::wallet_backend::WalletBackend::forget_all_secrets(self); + } + + fn shutdown(&self) -> futures::future::BoxFuture<'_, ()> { + Box::pin(crate::wallet_backend::WalletBackend::shutdown(self)) + } +} + /// Drain managed backend work before tearing down a standalone wallet backend. #[cfg(feature = "cli")] pub async fn shutdown_app_context_wallet_backend(ctx: &Arc) { - match ctx.subtasks.shutdown_async().await { - Ok(TaskShutdownOutcome::Complete) => { - if let Ok(backend) = ctx.wallet_backend() { + let task_shutdown_outcome = ctx.subtasks.shutdown_async().await.ok(); + let wallet_backend = ctx.wallet_backend().ok(); + finish_app_context_wallet_backend_shutdown(task_shutdown_outcome, wallet_backend).await; +} + +#[cfg(feature = "cli")] +async fn finish_app_context_wallet_backend_shutdown( + task_shutdown_outcome: Option, + wallet_backend: Option>, +) { + if let Some(backend) = wallet_backend.as_ref() { + backend.forget_all_secrets(); + } + + match task_shutdown_outcome { + Some(TaskShutdownOutcome::Complete) => { + if let Some(backend) = wallet_backend { backend.shutdown().await; + backend.forget_all_secrets(); } } - Ok(TaskShutdownOutcome::BackendTasksTimedOut) => { + Some(TaskShutdownOutcome::BackendTasksTimedOut) => { tracing::warn!("Managed backend work timed out; skipping standalone wallet teardown") } - Err(_) => { + None => { tracing::warn!("Managed task shutdown failed; skipping standalone wallet teardown") } } @@ -464,3 +496,61 @@ fn available_network_names(config: &crate::config::Config) -> String { names.join(", ") } } + +#[cfg(all(test, feature = "cli"))] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + struct MockStandaloneShutdownBackend { + forget_calls: AtomicUsize, + secret_cached: AtomicBool, + shutdown_calls: AtomicUsize, + } + + impl StandaloneShutdownWalletBackend for MockStandaloneShutdownBackend { + fn forget_all_secrets(&self) { + self.forget_calls.fetch_add(1, Ordering::Relaxed); + self.secret_cached.store(false, Ordering::Release); + } + + fn shutdown(&self) -> futures::future::BoxFuture<'_, ()> { + self.shutdown_calls.fetch_add(1, Ordering::Relaxed); + Box::pin(async {}) + } + } + + fn mock_wallet_backend() -> Arc { + Arc::new(MockStandaloneShutdownBackend { + forget_calls: AtomicUsize::new(0), + secret_cached: AtomicBool::new(true), + shutdown_calls: AtomicUsize::new(0), + }) + } + + #[tokio::test] + async fn backend_task_timeout_forgets_standalone_wallet_secrets_without_teardown() { + let backend = mock_wallet_backend(); + + finish_app_context_wallet_backend_shutdown( + Some(TaskShutdownOutcome::BackendTasksTimedOut), + Some(Arc::clone(&backend)), + ) + .await; + + assert_eq!(backend.forget_calls.load(Ordering::Relaxed), 1); + assert!(!backend.secret_cached.load(Ordering::Acquire)); + assert_eq!(backend.shutdown_calls.load(Ordering::Relaxed), 0); + } + + #[tokio::test] + async fn task_manager_failure_forgets_standalone_wallet_secrets_without_teardown() { + let backend = mock_wallet_backend(); + + finish_app_context_wallet_backend_shutdown(None, Some(Arc::clone(&backend))).await; + + assert_eq!(backend.forget_calls.load(Ordering::Relaxed), 1); + assert!(!backend.secret_cached.load(Ordering::Acquire)); + assert_eq!(backend.shutdown_calls.load(Ordering::Relaxed), 0); + } +}