From 79ce0c18a1e772d18f5348580f0cfb54a644e57e Mon Sep 17 00:00:00 2001 From: NimBold Date: Thu, 16 Jul 2026 17:56:00 +0330 Subject: [PATCH] fix(downloads): harden aria2 recovery lifecycle --- src-tauri/src/lib.rs | 156 +++++++++++++++++++++++++++++-- src-tauri/src/queue.rs | 123 +++++++++++++++++++++--- src-tauri/tests/queue_manager.rs | 100 +++++++++++++++++++- src/components/DownloadTable.tsx | 13 ++- 4 files changed, 362 insertions(+), 30 deletions(-) diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 67450d4..804fde7 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1814,6 +1814,7 @@ const MEDIA_METADATA_CACHE_TTL: Duration = Duration::from_secs(60); const MEDIA_METADATA_TIMEOUT: Duration = Duration::from_secs(55); const MEDIA_METADATA_CACHE_MAX_ENTRIES: usize = 128; const FILE_METADATA_TIMEOUT: Duration = Duration::from_secs(20); +const MAX_SHELL_OUTPUT_BYTES: usize = 32 * 1024 * 1024; struct ShellCommandOutput { status_code: Option, @@ -1836,16 +1837,37 @@ async fn shell_command_output_with_timeout( stdout: Vec::new(), stderr: Vec::new(), }; + let mut output_bytes = 0usize; while let Some(event) = events.recv().await { match event { tauri_plugin_shell::process::CommandEvent::Stdout(bytes) => { + let next_size = output_bytes + .saturating_add(bytes.len()) + .saturating_add(1); + if next_size > MAX_SHELL_OUTPUT_BYTES { + return Err(format!( + "{operation} produced more than {} MiB of output", + MAX_SHELL_OUTPUT_BYTES / (1024 * 1024) + )); + } output.stdout.extend(bytes); output.stdout.push(b'\n'); + output_bytes = next_size; } tauri_plugin_shell::process::CommandEvent::Stderr(bytes) => { + let next_size = output_bytes + .saturating_add(bytes.len()) + .saturating_add(1); + if next_size > MAX_SHELL_OUTPUT_BYTES { + return Err(format!( + "{operation} produced more than {} MiB of output", + MAX_SHELL_OUTPUT_BYTES / (1024 * 1024) + )); + } output.stderr.extend(bytes); output.stderr.push(b'\n'); + output_bytes = next_size; } tauri_plugin_shell::process::CommandEvent::Terminated(payload) => { output.status_code = payload.code; @@ -1855,26 +1877,35 @@ async fn shell_command_output_with_timeout( } } - output + Ok(output) }; match tokio::time::timeout(timeout, collect_output).await { - Ok(output) if output.status_code.is_some() => Ok(output), - Ok(_) => match child.kill() { + Ok(Ok(output)) if output.status_code.is_some() => Ok(output), + Ok(Ok(_)) => match child.kill() { Ok(()) => Err(format!( "{operation} ended without a process exit status" )), Err(error) => Err(format!( - "{operation} ended without a process exit status and failed to terminate yt-dlp: {error}" + "{operation} ended without a process exit status and failed to terminate the child process: {error}" )), }, + Ok(Err(error)) => { + let kill_error = child.kill().err(); + match kill_error { + Some(kill_error) => Err(format!( + "{error}; failed to terminate the child process: {kill_error}" + )), + None => Err(error), + } + } Err(_) => match child.kill() { Ok(()) => Err(format!( "{operation} timed out after {}s", timeout.as_secs() )), Err(error) => Err(format!( - "{operation} timed out after {}s and failed to terminate yt-dlp: {error}", + "{operation} timed out after {}s and failed to terminate the child process: {error}", timeout.as_secs() )), }, @@ -6028,7 +6059,8 @@ mod tests { should_send_metadata_credentials, collect_log_files, FirelinkDeepLink, percent_decode_metadata_value, MediaProgress, MediaProgressEmitterState, MediaSpeedSampler, MEDIA_PROGRESS_PREFIX, - observe_aria2_connections, Aria2ConnectionObservation, Aria2RecoveryReason, + observe_aria2_connections, observe_aria2_connections_with_epoch, + Aria2ConnectionObservation, Aria2RecoveryReason, parse_media_playlist_metadata, }; use serde_json::json; @@ -6179,6 +6211,66 @@ mod tests { ); } + #[test] + fn aria2_recovery_observation_resets_after_a_control_epoch_change() { + let start = Instant::now(); + let mut observation = Aria2ConnectionObservation::default(); + + for (offset, completed, speed) in [ + (0, 10 * 1024 * 1024, 10.0 * 1024.0 * 1024.0), + (1, 11 * 1024 * 1024, 10.0 * 1024.0 * 1024.0), + (2, 12 * 1024 * 1024, 10.0 * 1024.0 * 1024.0), + ] { + observe_aria2_connections_with_epoch( + &mut observation, + "gid-epoch", + 1, + "active", + 100 * 1024 * 1024, + completed, + speed, + 16, + 16, + false, + start + Duration::from_secs(offset), + ); + } + observe_aria2_connections_with_epoch( + &mut observation, + "gid-epoch", + 1, + "active", + 100 * 1024 * 1024, + 13 * 1024 * 1024, + 1024.0 * 1024.0, + 16, + 16, + false, + start + Duration::from_secs(31), + ); + assert!(observation.saw_multiple_connections); + + assert_eq!( + observe_aria2_connections_with_epoch( + &mut observation, + "gid-epoch", + 2, + "active", + 100 * 1024 * 1024, + 13 * 1024 * 1024, + 1024.0 * 1024.0, + 1, + 16, + false, + start + Duration::from_secs(62), + ), + None + ); + assert_eq!(observation.control_epoch, 2); + assert!(!observation.saw_multiple_connections); + assert_eq!(observation.healthy_speed_samples, 0); + } + #[test] fn slow_recovery_requires_a_real_multi_connection_transfer() { let start = Instant::now(); @@ -7345,6 +7437,7 @@ impl Aria2RecoveryReason { #[derive(Default)] struct Aria2ConnectionObservation { gid: String, + control_epoch: u64, saw_multiple_connections: bool, healthy_speed_samples: u8, degraded_since: Option, @@ -7361,6 +7454,7 @@ const ARIA2_MIN_PEAK_SPEED_FOR_DEGRADED_RECOVERY: f64 = 64.0 * 1024.0; const ARIA2_DEGRADED_SPEED_FRACTION: f64 = 0.20; const ARIA2_MIN_HEALTHY_SPEED_SAMPLES: u8 = 3; +#[cfg(test)] fn observe_aria2_connections( observation: &mut Aria2ConnectionObservation, gid: &str, @@ -7373,9 +7467,38 @@ fn observe_aria2_connections( speed_limited: bool, now: Instant, ) -> Option { - if observation.gid != gid { + observe_aria2_connections_with_epoch( + observation, + gid, + 0, + status, + total, + completed, + speed_bytes, + active_connections, + requested_connections, + speed_limited, + now, + ) +} + +fn observe_aria2_connections_with_epoch( + observation: &mut Aria2ConnectionObservation, + gid: &str, + control_epoch: u64, + status: &str, + total: u64, + completed: u64, + speed_bytes: f64, + active_connections: i32, + requested_connections: i32, + speed_limited: bool, + now: Instant, +) -> Option { + if observation.gid != gid || observation.control_epoch != control_epoch { *observation = Aria2ConnectionObservation { gid: gid.to_string(), + control_epoch, last_completed: completed, ..Default::default() }; @@ -7912,6 +8035,15 @@ pub fn run() { let mut observations: HashMap = HashMap::new(); loop { interval.tick().await; + // Terminal cleanup removes a download's GID mapping. Do + // not retain one observation per historical download for + // the lifetime of the poller. + let mapped_ids: HashSet = poll_mgr + .aria2_gid_mappings() + .into_iter() + .map(|(_, id)| id) + .collect(); + observations.retain(|id, _| mapped_ids.contains(id)); let params = serde_json::json!([["gid", "status", "totalLength", "completedLength", "downloadSpeed", "connections", "errorMessage"]]); if let Ok(active_list) = rpc_call(poll_port.load(std::sync::atomic::Ordering::Relaxed), &poll_secret, "aria2.tellActive", params).await { if let Some(active_arr) = active_list.as_array() { @@ -7939,11 +8071,14 @@ pub fn run() { .unwrap_or(1) .max(1); let speed_limited = poll_mgr.aria2_speed_limited(&id).await; + let control_epoch = + poll_mgr.current_aria2_control_epoch(&id).await; let now = Instant::now(); let observation = observations.entry(id.clone()).or_default(); - let recovery_reason = observe_aria2_connections( + let recovery_reason = observe_aria2_connections_with_epoch( observation, gid, + control_epoch, status, total, completed, @@ -7990,7 +8125,10 @@ pub fn run() { active_connections, requested_connections ); - if let Err(error) = poll_mgr.refresh_aria2_connections(&id, gid).await { + if let Err(error) = poll_mgr + .refresh_aria2_connections(&id, gid, control_epoch) + .await + { log::warn!( "aria2 connection recovery [{}] for gid {} failed: {}", id, diff --git a/src-tauri/src/queue.rs b/src-tauri/src/queue.rs index ac7fba1..18ea852 100644 --- a/src-tauri/src/queue.rs +++ b/src-tauri/src/queue.rs @@ -57,6 +57,16 @@ pub enum PendingOutcome { Error(String), } +/// Result of recycling an aria2 transfer's connections. A refresh can race +/// with daemon completion or leave the transfer paused after an ambiguous +/// unpause failure, so callers must handle the verified daemon outcome. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Aria2RefreshOutcome { + Resumed, + Paused, + Complete, +} + /// What kind of sidecar a queued task spawns. Drives which runner the /// dispatcher invokes. #[derive(Debug, Clone)] @@ -117,7 +127,7 @@ pub trait SidecarSpawner: Send + Sync + 'static { /// Recycle the connections for an active aria2 transfer without changing /// its gid or releasing its queue permit. Production uses forcePause / /// unpause; test spawners can leave this unsupported. - async fn refresh_uri(&self, _gid: &str) -> Result<(), String> { + async fn refresh_uri(&self, _gid: &str) -> Result { Err("aria2 connection refresh is unavailable".to_string()) } @@ -1112,22 +1122,27 @@ impl QueueManager { /// Recycle an active transfer's connections after the poller observes a /// persistent connection-pool collapse or a true zero-progress stall. - /// The transfer keeps its gid, partial file, and queue permit. - pub async fn refresh_aria2_connections(&self, id: &str, gid: &str) -> Result<(), String> { + /// The observed epoch must still own the GID before the refresh can act. + pub async fn refresh_aria2_connections( + &self, + id: &str, + gid: &str, + observed_epoch: u64, + ) -> Result<(), String> { let _control_guard = self.acquire_aria2_control(id).await; if self.aria2_gid_for_download(id).as_deref() != Some(gid) || !self.is_registered(id).await || self.is_aria2_retry_cancelled(id).await + || !self.is_aria2_control_epoch_current(id, observed_epoch).await { return Ok(()); } - let epoch = self.current_aria2_control_epoch(id).await; - self.spawner.refresh_uri(gid).await?; + let outcome = self.spawner.refresh_uri(gid).await?; let still_current = self.is_registered(id).await && !self.is_aria2_retry_cancelled(id).await - && self.is_aria2_control_epoch_current(id, epoch).await + && self.is_aria2_control_epoch_current(id, observed_epoch).await && self.aria2_gid_for_download(id).as_deref() == Some(gid); if !still_current { log::info!( @@ -1135,6 +1150,31 @@ impl QueueManager { id, gid ); + return Ok(()); + } + + match outcome { + Aria2RefreshOutcome::Resumed => {} + Aria2RefreshOutcome::Paused => { + self.next_aria2_control_epoch(id).await; + self.cancel_aria2_retries(id).await; + self.release_permit(id).await; + self.emit_state(id, DownloadStatus::Paused); + log::warn!( + "aria2 connection refresh [{}]: gid {} remained paused; released its queue permit", + id, + gid + ); + } + Aria2RefreshOutcome::Complete => { + self.apply_completion_locked(id, PendingOutcome::Complete) + .await; + log::info!( + "aria2 connection refresh [{}]: gid {} completed during recovery", + id, + gid + ); + } } Ok(()) } @@ -1961,19 +2001,72 @@ impl SidecarSpawner for ProductionSpawner { } } - async fn refresh_uri(&self, gid: &str) -> Result<(), String> { + async fn refresh_uri(&self, gid: &str) -> Result { let state = self.app_handle.state::(); let port = state.aria2_port.load(std::sync::atomic::Ordering::Relaxed); let secret = &state.aria2_secret; - let paused = crate::rpc_call(port, secret, "aria2.forcePause", serde_json::json!([gid])) - .await - .map_err(|error| format!("failed to refresh aria2 gid {gid}: {error}"))?; - crate::ensure_aria2_gid_result("forcePause", gid, &paused)?; + let pause_error = match crate::rpc_call( + port, + secret, + "aria2.forcePause", + serde_json::json!([gid]), + ) + .await + { + Ok(result) => crate::ensure_aria2_gid_result("forcePause", gid, &result) + .err() + .map(|error| error.to_string()), + Err(error) => Some(format!("failed to refresh aria2 gid {gid}: {error}")), + }; - let resumed = crate::rpc_call(port, secret, "aria2.unpause", serde_json::json!([gid])) - .await - .map_err(|error| format!("failed to refresh aria2 gid {gid}: {error}"))?; - crate::ensure_aria2_gid_result("unpause", gid, &resumed) + if let Some(error) = pause_error { + match crate::aria2_download_status(port, secret, gid).await { + Ok(status) if status == "paused" => { + log::warn!( + "aria2 connection refresh: forcePause for gid {} failed after the daemon paused it; continuing with unpause", + gid + ); + } + Ok(status) if status == "complete" => { + return Ok(Aria2RefreshOutcome::Complete); + } + Ok(status) => { + return Err(format!("{error}; aria2 gid {gid} is still {status}")); + } + Err(status_error) => { + return Err(format!( + "{error}; failed to verify aria2 gid {gid}: {status_error}" + )); + } + } + } + + let unpause_error = match crate::rpc_call( + port, + secret, + "aria2.unpause", + serde_json::json!([gid]), + ) + .await + { + Ok(result) => crate::ensure_aria2_gid_result("unpause", gid, &result) + .err() + .map(|error| error.to_string()), + Err(error) => Some(format!("failed to refresh aria2 gid {gid}: {error}")), + }; + + let Some(error) = unpause_error else { + return Ok(Aria2RefreshOutcome::Resumed); + }; + let status = crate::aria2_download_status(port, secret, gid).await?; + match status.as_str() { + "active" | "waiting" => Ok(Aria2RefreshOutcome::Resumed), + "paused" => Ok(Aria2RefreshOutcome::Paused), + "complete" => Ok(Aria2RefreshOutcome::Complete), + _ => Err(format!( + "{error}; aria2 gid {gid} reports terminal or unknown state {status}" + )), + } } async fn run_media( diff --git a/src-tauri/tests/queue_manager.rs b/src-tauri/tests/queue_manager.rs index 3de002b..1089532 100644 --- a/src-tauri/tests/queue_manager.rs +++ b/src-tauri/tests/queue_manager.rs @@ -1,5 +1,6 @@ use firelink_lib::queue::{ - QueueManager, QueuedTask, SidecarSpawner, SpawnPayload, TaskKind, MEDIA_RUN_CANCELLED, + Aria2RefreshOutcome, QueueManager, QueuedTask, SidecarSpawner, SpawnPayload, TaskKind, + MEDIA_RUN_CANCELLED, }; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; @@ -25,6 +26,11 @@ struct FailFirstAria2Spawner { fail_first: std::sync::atomic::AtomicBool, } +struct RefreshOutcomeSpawner { + outcome: Aria2RefreshOutcome, + refresh_calls: AtomicUsize, +} + impl FailFirstAria2Spawner { fn new() -> Self { Self { @@ -100,6 +106,31 @@ impl CountingSpawner { } } +#[async_trait::async_trait] +impl SidecarSpawner for RefreshOutcomeSpawner { + async fn add_uri(&self, _id: &str, _payload: &SpawnPayload) -> Result { + unreachable!("refresh outcome tests do not spawn aria2") + } + + async fn remove_uri(&self, _gid: &str) -> Result<(), String> { + unreachable!("refresh outcome tests do not remove aria2") + } + + async fn refresh_uri(&self, _gid: &str) -> Result { + self.refresh_calls.fetch_add(1, Ordering::SeqCst); + Ok(self.outcome) + } + + async fn run_media( + &self, + _id: &str, + _payload: &SpawnPayload, + _generation: u64, + ) -> Result<(), String> { + unreachable!("refresh outcome tests do not run media") + } +} + #[async_trait::async_trait] impl firelink_lib::queue::SidecarSpawner for CountingSpawner { async fn add_uri(&self, _id: &str, _payload: &SpawnPayload) -> Result { @@ -648,6 +679,73 @@ async fn aria2_permit_survives_rpc_return() { handle.abort(); } +#[tokio::test] +async fn failed_refresh_that_leaves_gid_paused_releases_permit_but_keeps_resume_mapping() { + let app = mock_builder() + .build(mock_context(noop_assets())) + .expect("mock app"); + let spawner = Arc::new(RefreshOutcomeSpawner { + outcome: Aria2RefreshOutcome::Paused, + refresh_calls: AtomicUsize::new(0), + }); + let manager = QueueManager::test_new( + app.handle().clone(), + 1, + Arc::clone(&spawner) as Arc, + ); + manager.push(aria2_task("refresh-paused")).await.unwrap(); + assert!(manager.ensure_aria2_permit("refresh-paused").await); + manager + .remember_gid("refresh-paused".to_string(), "gid-refresh-paused".to_string()) + .await; + + let epoch = manager.current_aria2_control_epoch("refresh-paused").await; + manager + .refresh_aria2_connections("refresh-paused", "gid-refresh-paused", epoch) + .await + .unwrap(); + + assert_eq!(spawner.refresh_calls.load(Ordering::SeqCst), 1); + assert_eq!(manager.available_permits(), 1); + assert!(!manager.has_active_permit("refresh-paused").await); + assert_eq!( + manager.aria2_gid_for_download("refresh-paused").as_deref(), + Some("gid-refresh-paused") + ); + assert!(manager.is_registered("refresh-paused").await); +} + +#[tokio::test] +async fn stale_refresh_observation_cannot_touch_a_newer_control_epoch() { + let app = mock_builder() + .build(mock_context(noop_assets())) + .expect("mock app"); + let spawner = Arc::new(RefreshOutcomeSpawner { + outcome: Aria2RefreshOutcome::Resumed, + refresh_calls: AtomicUsize::new(0), + }); + let manager = QueueManager::test_new( + app.handle().clone(), + 1, + Arc::clone(&spawner) as Arc, + ); + manager.push(aria2_task("refresh-stale")).await.unwrap(); + assert!(manager.ensure_aria2_permit("refresh-stale").await); + manager + .remember_gid("refresh-stale".to_string(), "gid-refresh-stale".to_string()) + .await; + let stale_epoch = manager.current_aria2_control_epoch("refresh-stale").await; + manager.next_aria2_control_epoch("refresh-stale").await; + + manager + .refresh_aria2_connections("refresh-stale", "gid-refresh-stale", stale_epoch) + .await + .unwrap(); + + assert_eq!(spawner.refresh_calls.load(Ordering::SeqCst), 0); + assert!(manager.has_active_permit("refresh-stale").await); +} + #[tokio::test] async fn transient_aria2_error_reissues_after_backoff() { use firelink_lib::queue::PendingOutcome; diff --git a/src/components/DownloadTable.tsx b/src/components/DownloadTable.tsx index 1ffbd65..3554091 100644 --- a/src/components/DownloadTable.tsx +++ b/src/components/DownloadTable.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect, useMemo, useRef, useCallback } from 'react'; -import { useDownloadStore, DownloadItem } from '../store/useDownloadStore'; +import { useDownloadStore, DownloadItem, MAIN_QUEUE_ID } from '../store/useDownloadStore'; import { useToast } from '../contexts/ToastContext'; import { useSettingsStore } from '../store/useSettingsStore'; import { SidebarFilter } from './Sidebar'; @@ -240,7 +240,7 @@ export const DownloadTable: React.FC = ({ filter }) => { // on every download update. Compute the same queue membership once and pass // the resulting position to rows instead. const queuePositionsByDownloadId = useMemo(() => { - const grouped = new Map(); + const grouped = new Map(); for (const download of downloads) { if ( download.status === 'completed' @@ -248,15 +248,18 @@ export const DownloadTable: React.FC = ({ filter }) => { ) { continue; } - const queueItems = grouped.get(download.queueId) || []; + const queueId = download.queueId || MAIN_QUEUE_ID; + const queueItems = grouped.get(queueId) || []; queueItems.push(download); - grouped.set(download.queueId, queueItems); + grouped.set(queueId, queueItems); } const positions = new Map(); for (const queueItems of grouped.values()) { queueItems.sort((left, right) => - (left.queuePosition ?? 0) - (right.queuePosition ?? 0) + (left.queuePosition ?? Number.MAX_SAFE_INTEGER) - + (right.queuePosition ?? Number.MAX_SAFE_INTEGER) || + left.id.localeCompare(right.id) ); queueItems.forEach((download, index) => { positions.set(download.id, { index, length: queueItems.length });