diff --git a/CHANGELOG.md b/CHANGELOG.md index 4de229b17..48d500acc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -202,6 +202,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. + - **Token balance refresh status**: requesting a refresh while token balances are already updating now shows a brief informational note instead of a red error banner that must be dismissed. diff --git a/src/app.rs b/src/app.rs index 1e525fbb5..e045780ee 100644 --- a/src/app.rs +++ b/src/app.rs @@ -38,8 +38,8 @@ 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::wallet_backend::DetScope; +use crate::utils::tasks::{TaskManager, TaskShutdownOutcome}; +use crate::wallet_backend::{DetScope, WalletBackend}; use dash_sdk::dpp::dashcore::Network; use dash_sdk::platform::Identifier; use eframe::{App, egui}; @@ -81,6 +81,38 @@ 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, + BackendTasksTimedOut, + WalletBackendTimedOut, +} + +fn shutdown_hard_deadline() -> Duration { + 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, @@ -304,8 +336,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 @@ -334,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 }, @@ -370,7 +407,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, @@ -657,12 +694,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(), }; @@ -781,12 +818,14 @@ 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. - shutdown_receiver: Option>, + /// 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. 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`. @@ -795,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`]. @@ -1459,11 +1502,13 @@ 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, #[cfg(feature = "mcp")] mcp_server_pending_config, + secret_prompt_host, secret_prompt_receiver, active_secret_prompt: None, prompt_was_blocking: false, @@ -1546,7 +1591,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() @@ -1569,7 +1614,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"); } @@ -1593,25 +1638,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( + let _ = 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, + ) + }, ); } @@ -1626,39 +1673,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 - } - }; + let _ = 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, + ) + }, ); } @@ -1983,6 +2032,223 @@ 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 { + match *result { + BackendTaskSuccessResult::NetworkContextRegistered { context, .. } + | BackendTaskSuccessResult::NetworkContextCreated { context, .. } => { + Self::push_unique_context(contexts, context); + } + _ => {} + } + } + } + + async fn shutdown_wallet_backend_instances( + wallet_backends: &[Arc], + shutdown_timeout: Duration, + ) -> ShutdownOutcome { + Self::forget_wallet_backend_secrets(wallet_backends); + + let shutdowns = futures::future::join_all( + wallet_backends + .iter() + .map(|wallet_backend| wallet_backend.shutdown()), + ); + let outcome = if tokio::time::timeout(shutdown_timeout, shutdowns) + .await + .is_err() + { + tracing::warn!( + timeout_secs = shutdown_timeout.as_secs(), + "Wallet backend shutdown timed out; closing with degraded teardown" + ); + ShutdownOutcome::WalletBackendTimedOut + } else { + 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(); + } + } + + 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() + && !wallet_backends + .iter() + .any(|existing| Arc::ptr_eq(existing, &backend)) + { + wallet_backends.push(backend); + } + } + + 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> { + let mut contexts = Vec::new(); + for context in self.network_contexts.values() { + Self::push_unique_context(&mut contexts, Arc::clone(context)); + } + contexts + } + + 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>>, + ) -> 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(task_shutdown_outcome, contexts).await + } + + 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) => { + 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(); + 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 task_shutdown_outcome = loop { + tokio::select! { + result = &mut task_shutdown => break result.ok(), + task_result = task_result_receiver.recv() => { + let Some(task_result) = task_result else { + break task_shutdown.await.ok(); + }; + Self::collect_created_context(&mut contexts, task_result); + } + } + }; + + let outcome = Self::finish_wallet_shutdown( + task_shutdown_outcome, + contexts, + task_result_receiver, + #[cfg(feature = "mcp")] + mcp_app_context, + ) + .await; + 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 (_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) = task_result_receiver.try_recv() { + Self::collect_created_context(&mut contexts, task_result); + } + + match task_shutdown.try_recv() { + 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); + } + } + }; + + let (tx, rx) = std::sync::mpsc::sync_channel(1); + tokio::spawn(async move { + let outcome = Self::finish_wallet_shutdown( + task_shutdown_outcome, + contexts, + task_result_receiver, + #[cfg(feature = "mcp")] + mcp_app_context, + ) + .await; + let _ = tx.send(outcome); + }); + + match rx.recv_timeout(WALLET_BACKEND_SHUTDOWN_TIMEOUT + SHUTDOWN_DEADLINE_MARGIN) { + Ok(ShutdownOutcome::Complete) + if task_shutdown_outcome == Some(TaskShutdownOutcome::Complete) => {} + Ok(outcome) => tracing::warn!( + ?outcome, + ?task_shutdown_outcome, + "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 { @@ -1993,37 +2259,54 @@ 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(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::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" + ), + } + 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 = shutdown_hard_deadline(); + 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(); @@ -2043,7 +2326,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; @@ -2228,6 +2511,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 @@ -2590,7 +2877,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; }); } @@ -2629,18 +2916,13 @@ 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() { - tracing::debug!("on_exit: async shutdown was initiated, 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"); - if let Err(e) = self.subtasks.shutdown() { - tracing::error!("Error during task shutdown: {}", e); - } + self.run_blocking_shutdown_fallback(); tracing::debug!("App shutdown complete"); } } @@ -3100,3 +3382,318 @@ 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); + 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"); + + 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" + ); + } + + #[tokio::test] + 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 + }, + ) + .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_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, &[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_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 + }, + ) + .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); + } + + #[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::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 { + 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_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/error.rs b/src/backend_task/error.rs index c6b2a55c1..62db9e18f 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 08ddcbf82..6c5b99592 100644 --- a/src/backend_task/mod.rs +++ b/src/backend_task/mod.rs @@ -87,20 +87,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)), } } @@ -729,6 +754,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, @@ -1015,28 +1046,68 @@ impl AppContext { .ok_or(TaskError::NetworkContextCreationFailed { network })?; new_ctx.install_secret_prompt(self.secret_prompt()); - // 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 { - 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 cancellation_token.is_cancelled() + && let Ok(backend) = new_ctx.wallet_backend() + { + backend.forget_all_secrets(); + backend.shutdown().await; + } Ok(BackendTaskSuccessResult::NetworkContextCreated { network, context: new_ctx, @@ -1183,6 +1254,82 @@ mod tests { use super::*; use crate::context::feature_gate::FeatureGate; + #[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/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..2b0fa4ab9 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,21 +60,146 @@ 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 { + 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(); - 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()); - app_context.run_backend_task(task, sender).await + let result = app_context.run_backend_task(task, sender).await?; + Ok(lifecycle_tail(result).await) }) }) .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)); + } + + #[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/server.rs b/src/mcp/server.rs index 1af2aa387..d294654d6 100644 --- a/src/mcp/server.rs +++ b/src/mcp/server.rs @@ -8,6 +8,59 @@ use rmcp::model::*; use rmcp::{ErrorData as McpError, RoleServer, ServerHandler, service::RequestContext}; 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) { + 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(); + } + } + Some(TaskShutdownOutcome::BackendTasksTimedOut) => { + tracing::warn!("Managed backend work timed out; skipping standalone wallet teardown") + } + None => { + 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 +262,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. @@ -447,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); + } +} diff --git a/src/mcp/tools/network.rs b/src/mcp/tools/network.rs index e832ef034..bb61e010c 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,35 +228,31 @@ 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 14193e83d..e79deaf0a 100644 --- a/src/utils/tasks.rs +++ b/src/utils/tasks.rs @@ -1,26 +1,97 @@ -use std::sync::{Arc, Mutex, atomic::AtomicUsize}; +use futures::FutureExt; +use std::sync::{Arc, Mutex}; use tokio::time::{Duration, timeout}; 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 TaskCompletion = 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 TrackedTask { + name: &'static str, + completion: TaskCompletion, +} + +impl std::fmt::Debug for TrackedTask { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_tuple("TrackedTask") + .field(&self.name) + .finish() + } +} + +/// Terminal state of managed task shutdown. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TaskShutdownOutcome { + /// Ordinary tasks and tracked joins reached their bounded shutdown points. + Complete, + /// Tracked 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 - 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, + tracking_joins: bool, + tasks: tokio::task::JoinSet<&'static str>, + tracked_tasks: Vec, +} + +#[derive(Debug)] +struct ShutdownTasks { + tasks: tokio::task::JoinSet<&'static str>, + task_state: Arc>, +} + /// 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, + tracking_joins: true, + tasks: tokio::task::JoinSet::new(), + tracked_tasks: Vec::new(), + })), active_names: Arc::new(Mutex::new(Vec::new())), } } @@ -28,133 +99,257 @@ 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, { - 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 Err(future); } - let subtasks = self.tasks.clone(); - tokio::spawn(spawn_subtask(subtasks, name, future)); + 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 + }); + Ok(()) } - /// Start an asynchronous graceful shutdown of all subtasks. + /// Spawn blocking work and retain its real task handle through shutdown. /// - /// 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<()> { - let cancel = self.cancellation_token.clone(); - let subtasks = self.tasks.clone(); - let active_names = self.active_names.clone(); + /// The async join observer remains an ordinary abortable subtask, while a + /// separate completion handle lets shutdown await non-cancellable blocking work. + /// 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, + 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 Err((task, on_join)); + } - let (tx, rx) = tokio::sync::oneshot::channel::<()>(); + let join_handle = tokio::task::spawn_blocking(task); + let completion = task_completion(join_handle); + self.register_tracked_completion(&mut state, name, completion, on_join); + Ok(()) + } - tokio::task::spawn(async move { - let completed = shutdown_inner(&cancel, &subtasks, &active_names, "async").await; + /// 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)); + } - tracing::debug!( - "Async shutdown complete, {} subtasks finished cleanly", - completed - ); + let completion = task_completion(tokio::spawn(task)); + self.register_tracked_completion(&mut state, name, completion, on_join); + Ok(()) + } - let _ = tx.send(()); + 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, + { + 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 + .lock() + .unwrap_or_else(|error| error.into_inner()) + .is_some() + }) + }); + state.tracked_tasks.push(TrackedTask { + name, + completion: callback_completion.clone(), + }); + if !state.accepting { + return; + } + 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 = callback_completion + .await + .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 completion callback stopped unexpectedly" + ); + } + name }); + } - rx + /// 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) } - /// 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. + /// Start an asynchronous graceful shutdown of all subtasks. /// - /// 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(); + /// Cancels ordinary tasks and returns a receiver that resolves after both + /// 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(); 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, rx) = tokio::sync::oneshot::channel::(); - // 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, outcome) = shutdown_all_inner( + tasks, + &active_names, + "async", + SHUTDOWN_TIMEOUT, + SHUTDOWN_TIMEOUT, + ) + .await; - // notify that shutdown is complete - if tx.send(completed).is_err() { - tracing::error!("Failed to send shutdown completion signal"); - } + tracing::debug!( + ?outcome, + "Async shutdown complete, {} subtasks finished cleanly", + completed + ); + + let _ = tx.send(outcome); }); - // 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() { - completed = count; - break; - } - // wait for a short time to avoid busy waiting - std::thread::sleep(WAIT_TIME); + rx + } + + fn begin_shutdown(&self) -> ShutdownTasks { + 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(); + ShutdownTasks { + tasks, + task_state: Arc::clone(&self.task_state), } + } +} - tracing::debug!("Shutdown complete, {} subtasks finished cleanly", completed); +fn task_completion(task: tokio::task::JoinHandle) -> TaskCompletion +where + T: Send + 'static, +{ + async move { Arc::new(Mutex::new(Some(task.await.map(|_| ())))) } + .boxed() + .shared() +} - Ok(()) - } +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 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 + }; + (completed, outcome) } -/// 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, + shutdown_timeout: Duration, ) -> 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 names_for_join = active_names.clone(); - let timed_out = timeout(SHUTDOWN_TIMEOUT, async move { - let mut tasks = tasks_list.lock().await; + let mut completed = 0; + 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 - 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 = 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, + cancelled = e.is_cancelled(), + panicked = e.is_panic(), "{label}: task joined with error" ), } @@ -163,10 +358,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 +381,72 @@ async fn shutdown_inner( } // Abort all remaining tasks - subtasks.lock().await.shutdown().await; + tasks.shutdown().await; - completed.load(std::sync::atomic::Ordering::Relaxed) + completed } -#[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 - }); +async fn shutdown_tracked_tasks( + task_state: Arc>, + label: &str, + shutdown_timeout: Duration, +) -> bool { + 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 completion callback stopped unexpectedly" + ); + } + } + } + }) + .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(), + "Tracked backend work exceeded shutdown wait; continuing with degraded teardown" + ); + false + } else { + tracing::trace!(total, "{label}: tracked task work joined"); + true + } } impl Default for TaskManager { @@ -213,3 +454,391 @@ impl Default for TaskManager { TaskManager::new() } } + +#[cfg(test)] +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(); + 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 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!( + 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()) + .tracked_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(); + 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( + 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(); + let accepted_task_ran = Arc::new(AtomicBool::new(false)); + let ran = Arc::clone(&accepted_task_ran); + 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); + + let late_task = 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)); + 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)] + 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); + 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"); + + 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); + 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"); + + 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); + } + + #[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); + let join_error_observed = Arc::new(AtomicBool::new(false)); + let observed = Arc::clone(&join_error_observed); + 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"); + }, + move |result| async move { + observed.store(result.is_err(), Ordering::Release); + }, + ) + .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"); + + 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); + 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 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" + ); + } +} 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. }