use crate::ipc::{DownloadStateEvent, DownloadStatus, QueueDirection}; use crate::retry::{backoff_and_emit, is_transient_network_error, BackoffOutcome}; use log; use serde::Deserialize; use serde_json; use std::collections::{HashMap, HashSet, VecDeque}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use tauri::{AppHandle, Manager}; use tokio::sync::{Mutex, Notify, OwnedSemaphorePermit, Semaphore}; use ts_rs::TS; /// Default capacity when no setting is read yet. pub const DEFAULT_MAX_CONCURRENT: usize = 3; pub const MEDIA_RUN_CANCELLED: &str = "__firelink_media_run_cancelled__"; /// Outcome of an aria2 completion that arrived before its gid was stored. /// Carries the outcome so the correct state emit survives the race. #[derive(Debug, Clone)] pub enum PendingOutcome { Complete, Error(String), } /// What kind of sidecar a queued task spawns. Drives which runner the /// dispatcher invokes. #[derive(Debug, Clone)] pub enum TaskKind { Aria2, Media, } /// Everything needed to start a sidecar, captured at enqueue time so the /// dispatcher can spawn it later without round-tripping back to the frontend. #[derive(Debug, Clone)] pub struct QueuedTask { pub id: String, pub queue_id: String, pub kind: TaskKind, pub payload: SpawnPayload, } /// Args mirroring start_download / start_media_download. Kept untyped-loose /// (String/Option) to match the existing command signatures exactly. #[derive(Debug, Clone, Default)] pub struct SpawnPayload { pub url: String, pub destination: String, pub filename: String, pub connections: Option, pub speed_limit: Option, pub username: Option, pub password: Option, pub headers: Option, pub checksum: Option, pub cookies: Option, pub mirrors: Option, pub user_agent: Option, pub max_tries: Option, pub proxy: Option, pub format_selector: Option, pub cookie_source: Option, pub is_media: bool, } /// A sidecar spawner. In production this calls the real aria2/yt-dlp /// runners; in tests it is replaced with a fake that records calls and /// optionally hangs to simulate a long-running download. #[async_trait::async_trait] pub trait SidecarSpawner: Send + Sync + 'static { /// Spawn an aria2 download. Returns the gid. Must return quickly (the /// permit is already parked before this is called). async fn add_uri(&self, id: &str, payload: &SpawnPayload) -> Result; /// Force-remove an aria2 gid created by a retry that raced with user /// cancellation. async fn remove_uri(&self, gid: &str) -> Result<(), String>; /// Run a media download to completion. The permit is parked for the full /// duration; release is handled by QueueManager on the runner's exit. async fn run_media(&self, id: &str, payload: &SpawnPayload) -> Result<(), String>; } /// The centralized concurrency gatekeeper. One instance lives in AppState. pub struct QueueManager { registered_ids: Mutex>, enqueue_cancellations: Mutex>, enqueue_generations: Mutex>, pending: Mutex>, semaphore: Arc, active_permits: Mutex>, active_kinds: Mutex>, target_capacity: AtomicUsize, slots_to_retire: AtomicUsize, notify: Notify, notify_permit_released: Notify, /// aria2 gid -> download id map (shared with the WS poller). pub aria2_gids: Arc>>, /// gid -> buffered (id_placeholder, outcome) for completions that arrived /// before the gid was stored. Drained by `remember_gid`. pub pending_completion: Arc>>, /// download id -> spawn payload for aria2 transient-error re-addUri retries. aria2_payloads: Mutex>, /// 0-based transient-error strike counter per aria2 download id. aria2_retry_strikes: Mutex>, /// Download ids whose aria2 retry loop must not create another job. aria2_retry_cancelled: Mutex>, /// Monotonic per-download aria2 control generation. Long-running queued /// resume tasks capture this and abort when a later pause/remove wins. aria2_control_epochs: Mutex>, spawner: Arc, app_handle: AppHandle, } impl QueueManager { /// Production constructor. Wired up in lib.rs setup(). pub fn new(app_handle: AppHandle, capacity: usize) -> Self { let spawner: Arc = Arc::new(ProductionSpawner::new(app_handle.clone())); Self::test_new(app_handle, capacity, spawner) } } impl QueueManager { /// Test-only constructor injecting a fake spawner. pub fn test_new( app_handle: AppHandle, capacity: usize, spawner: Arc, ) -> Self { Self { registered_ids: Mutex::new(HashSet::new()), enqueue_cancellations: Mutex::new(HashMap::new()), enqueue_generations: Mutex::new(HashMap::new()), pending: Mutex::new(VecDeque::new()), semaphore: Arc::new(Semaphore::new(capacity)), active_permits: Mutex::new(HashMap::new()), active_kinds: Mutex::new(HashMap::new()), target_capacity: AtomicUsize::new(capacity), slots_to_retire: AtomicUsize::new(0), notify: Notify::new(), notify_permit_released: Notify::new(), aria2_gids: Arc::new(std::sync::RwLock::new(HashMap::new())), pending_completion: Arc::new(Mutex::new(HashMap::new())), aria2_payloads: Mutex::new(HashMap::new()), aria2_retry_strikes: Mutex::new(HashMap::new()), aria2_retry_cancelled: Mutex::new(HashSet::new()), aria2_control_epochs: Mutex::new(HashMap::new()), spawner, app_handle, } } /// Current pending order, as id list. Returned by move_in_queue. pub async fn pending_order(&self, queue_id: Option<&str>) -> Vec { self.pending .lock() .await .iter() .filter(|task| queue_id.is_none_or(|queue_id| task.queue_id == queue_id)) .map(|t| t.id.clone()) .collect() } /// Explicitly release a backend registry id (e.g. on un-resumable false paths, removals, or detach). pub async fn release_registered_id(&self, id: &str) { self.registered_ids.lock().await.remove(id); } pub async fn is_registered(&self, id: &str) -> bool { self.registered_ids.lock().await.contains(id) } /// Reject an in-flight enqueue generation if a newer UI action supersedes it. pub async fn cancel_enqueue_generation(&self, id: &str, generation: u64) { let mut cancellations = self.enqueue_cancellations.lock().await; cancellations .entry(id.to_string()) .and_modify(|current| *current = (*current).max(generation)) .or_insert(generation); } /// Atomically reserve an ID after rejecting cancelled or replayed generations. /// The returned watermark must be passed to `rollback_enqueue_reservation` /// if ownership registration fails before the task is committed. pub async fn reserve_enqueue_generation( &self, id: &str, generation: u64, ) -> Result, String> { let cancellations = self.enqueue_cancellations.lock().await; if cancellations .get(id) .is_some_and(|cancelled| *cancelled >= generation) { return Err("Download enqueue was superseded by a newer user action".to_string()); } let mut generations = self.enqueue_generations.lock().await; let previous_generation = generations.get(id).copied(); if previous_generation.is_some_and(|seen| seen >= generation) { return Err("Download enqueue was superseded by a newer user action".to_string()); } let mut registered = self.registered_ids.lock().await; if registered.contains(id) { return Err("Duplicate task".to_string()); } registered.insert(id.to_string()); generations.insert(id.to_string(), generation); Ok(previous_generation) } pub async fn rollback_enqueue_reservation( &self, id: &str, generation: u64, previous_generation: Option, ) { let mut generations = self.enqueue_generations.lock().await; let mut registered = self.registered_ids.lock().await; if generations.get(id).copied() != Some(generation) { return; } registered.remove(id); match previous_generation { Some(previous) => { generations.insert(id.to_string(), previous); } None => { generations.remove(id); } } } pub async fn commit_reserved_enqueue( &self, task: QueuedTask, generation: u64, ) -> Result<(), String> { let id = task.id.clone(); let cancellations = self.enqueue_cancellations.lock().await; if cancellations .get(&id) .is_some_and(|cancelled| *cancelled >= generation) { return Err("Download enqueue was superseded by a newer user action".to_string()); } self.pending.lock().await.push_back(task); self.emit_state(id, DownloadStatus::Queued); self.notify.notify_one(); Ok(()) } /// Atomically checks the generation watermark before registering a task. pub async fn push_with_generation( &self, task: QueuedTask, generation: u64, ) -> Result<(), String> { let id = task.id.clone(); let previous_generation = self.reserve_enqueue_generation(&id, generation).await?; if let Err(error) = self.commit_reserved_enqueue(task, generation).await { self.rollback_enqueue_reservation(&id, generation, previous_generation) .await; return Err(error); } Ok(()) } /// Enqueue a task without a frontend lifecycle token. This is retained for /// internal/test callers and still gets replay protection at generation 0. pub async fn push(&self, task: QueuedTask) -> Result<(), String> { self.push_with_generation(task, 0).await } pub async fn next_aria2_control_epoch(&self, id: &str) -> u64 { let mut epochs = self.aria2_control_epochs.lock().await; let epoch = epochs.get(id).copied().unwrap_or_default().wrapping_add(1); epochs.insert(id.to_string(), epoch); epoch } pub async fn is_aria2_control_epoch_current(&self, id: &str, epoch: u64) -> bool { self.aria2_control_epochs .lock() .await .get(id) .copied() .unwrap_or_default() == epoch } pub async fn is_aria2_retry_cancelled(&self, id: &str) -> bool { self.aria2_retry_cancelled.lock().await.contains(id) } pub async fn has_aria2_retry_state(&self, id: &str) -> bool { self.aria2_retry_strikes.lock().await.contains_key(id) } /// Pop the next task, or None if empty. pub async fn pop_front(&self) -> Option { self.pending.lock().await.pop_front() } /// Acquire a permit from the semaphore (blocks until one is available). pub async fn acquire_permit(&self) -> Option { self.semaphore.clone().acquire_owned().await.ok() } async fn acquire_permit_after_retirement(&self) -> Option { loop { let permit = self.acquire_permit().await?; if self.retire_slot_if_needed() { permit.forget(); continue; } return Some(permit); } } fn retire_slot_if_needed(&self) -> bool { let mut debt = self.slots_to_retire.load(Ordering::Relaxed); while debt > 0 { match self.slots_to_retire.compare_exchange_weak( debt, debt - 1, Ordering::Relaxed, Ordering::Relaxed, ) { Ok(_) => return true, Err(actual) => debt = actual, } } false } /// Park an already-acquired permit under `id`. pub async fn park_permit(&self, id: &str, permit: OwnedSemaphorePermit) { self.active_permits .lock() .await .insert(id.to_string(), permit); } pub async fn active_kind(&self, id: &str) -> Option { self.active_kinds.lock().await.get(id).cloned() } /// Ensure an aria2 transfer owns exactly one queue permit. Returns true /// when this call acquired and parked the permit, false when one was /// already parked. pub async fn ensure_aria2_permit(&self, id: &str) -> bool { if self.active_permits.lock().await.contains_key(id) { return false; } let permit = match self.acquire_permit_after_retirement().await { Some(p) => p, None => return false, }; let mut permits = self.active_permits.lock().await; if permits.contains_key(id) { drop(permits); drop(permit); return false; } permits.insert(id.to_string(), permit); drop(permits); self.active_kinds .lock() .await .insert(id.to_string(), TaskKind::Aria2); true } pub async fn release_permit(&self, id: &str) { let removed = self.active_permits.lock().await.remove(id).is_some(); self.active_kinds.lock().await.remove(id); if removed { self.notify_permit_released.notify_waiters(); self.notify.notify_one(); } } /// Clear all permits belonging to aria2. Useful when aria2 WS connection drops. pub async fn clear_aria2_permits(&self) { let ids_to_fail: Vec = { let kinds = self.active_kinds.lock().await; kinds .iter() .filter(|(_, kind)| matches!(kind, TaskKind::Aria2)) .map(|(id, _)| id.clone()) .collect() }; for id in ids_to_fail { self.apply_completion( &id, PendingOutcome::Error("Aria2 WebSocket connection lost".to_string()), ) .await; } } /// Number of un-acquired permits currently in the semaphore pool. pub fn available_permits(&self) -> usize { self.semaphore.available_permits() } fn emit_state(&self, id: impl Into, status: DownloadStatus) { use tauri::Emitter; let _ = self .app_handle .emit("download-state", DownloadStateEvent::new(id, status)); } /// Resize the global concurrency limit. Grow adds permits immediately; /// shrink records a retirement debt honored lazily by the dispatcher. pub fn set_capacity(&self, new_target: usize) { let prev_target = self.target_capacity.swap(new_target, Ordering::Relaxed); if new_target == prev_target { return; } if new_target > prev_target { let mut delta = new_target - prev_target; loop { let debt = self.slots_to_retire.load(Ordering::Relaxed); let to_deduct = std::cmp::min(debt, delta); if self .slots_to_retire .compare_exchange_weak( debt, debt - to_deduct, Ordering::Relaxed, Ordering::Relaxed, ) .is_ok() { delta -= to_deduct; break; } } if delta > 0 { self.semaphore.add_permits(delta); } self.notify.notify_one(); } else { let delta = prev_target - new_target; self.slots_to_retire.fetch_add(delta, Ordering::Relaxed); } } /// Test accessor for the retirement debt counter. pub fn slots_to_retire_load(&self) -> usize { self.slots_to_retire.load(Ordering::Relaxed) } /// The long-running dispatcher. One instance is spawned in setup(). /// Idle-parks on Notify; CAS-honors retirement debt; re-pops under lock. pub async fn run_dispatcher(self: Arc) { loop { // (1) Idle-park: avoid busy-spin when pending is empty. if self.pending.lock().await.is_empty() { self.notify.notified().await; continue; } // (2) Acquire a slot. let permit = match self.acquire_permit_after_retirement().await { Some(p) => p, None => break, // Semaphore closed, exit dispatcher }; // (4) Re-pop under lock — guards against racing removals between // waking from Notify and acquiring the permit. let task = match self.pending.lock().await.pop_front() { Some(t) => t, None => { drop(permit); continue; } }; Arc::clone(&self).dispatch_one(permit, task).await; } } async fn dispatch_one(self: Arc, permit: OwnedSemaphorePermit, task: QueuedTask) { let id = task.id.clone(); // Park the permit BEFORE spawning. Uniform parking: // aria2's RPC returns instantly, so the permit must outlive the // dispatch_one call. Media runners release on exit. self.park_permit(&id, permit).await; self.active_kinds .lock() .await .insert(id.clone(), task.kind.clone()); self.emit_state(&id, DownloadStatus::Downloading); match task.kind { TaskKind::Aria2 => { self.aria2_retry_cancelled.lock().await.remove(&id); self.aria2_payloads .lock() .await .insert(id.clone(), task.payload.clone()); self.aria2_retry_strikes.lock().await.remove(&id); match self.spawner.add_uri(&id, &task.payload).await { Ok(gid) => { let cancelled = self.aria2_retry_cancelled.lock().await.contains(&id); if cancelled || !self.is_registered(&id).await { log::info!( "aria2 dispatch cancellation [{}]: removing late gid {}", id, gid ); if let Err(error) = self.spawner.remove_uri(&gid).await { log::warn!( "aria2 dispatch cancellation [{}]: failed to remove late gid {}: {}", id, gid, error ); } self.clear_aria2_retry_state(&id).await; self.release_permit(&id).await; return; } self.remember_gid(id.clone(), gid).await; } Err(error) => { self.clear_aria2_retry_state(&id).await; self.emit_failed(&id, error); self.release_permit(&id).await; } } } TaskKind::Media => { let this = Arc::clone(&self); let payload = task.payload.clone(); let id_for_task = id.clone(); tauri::async_runtime::spawn(async move { let outcome = this.spawner.run_media(&id_for_task, &payload).await; this.finish_runner(&id_for_task, outcome).await; }); } } } /// Terminal handler for non-aria2 transfers. Emits state and frees the permit. /// Does not emit or release anything on intentional MEDIA_RUN_CANCELLED. /// Note: `id` is the frontend download UUID, which survives indefinitely as /// the terminal state. async fn finish_runner(self: Arc, id: &str, outcome: Result<(), String>) { match outcome { Ok(()) => { self.emit_state(id, DownloadStatus::Completed); self.release_registered_id(id).await; } Err(error) if error == MEDIA_RUN_CANCELLED => {} Err(error) => { self.emit_failed(id, error); self.release_registered_id(id).await; } } self.release_permit(id).await; } fn emit_failed(&self, id: &str, error: String) { use tauri::Emitter; let _ = self .app_handle .emit("download-state", DownloadStateEvent::failed(id, error)); } /// Store gid -> id, then reconcile any buffered completion for that gid. pub async fn remember_gid(&self, id: String, gid: String) { { let mut gids = self.aria2_gids.write().unwrap(); gids.retain(|existing_gid, existing_id| { let keep = existing_id != &id || existing_gid == &gid; if !keep { log::warn!( "aria2 gid transition [{}]: dropping stale mapping {} before storing {}", id, existing_gid, gid ); } keep }); gids.insert(gid.clone(), id.clone()); } log::info!("aria2 gid transition [{}]: mapped {}", id, gid); let buffered = self.pending_completion.lock().await.remove(&gid); if let Some((_buf_id, outcome)) = buffered { self.apply_completion(&id, outcome).await; } } /// Apply an aria2 completion outcome: release permit + emit state. pub async fn apply_completion(&self, id: &str, outcome: PendingOutcome) { match outcome { PendingOutcome::Complete => { self.clear_aria2_retry_state(id).await; self.forget_aria2_gid(id).await; self.emit_state(id, DownloadStatus::Completed); self.release_registered_id(id).await; } PendingOutcome::Error(error) => { if error.to_ascii_lowercase().contains("checksum") { log::warn!("Checksum error detected for {}, cleaning up assets", id); if let Ok(primary_path) = crate::download_ownership::primary_path_for_id(&self.app_handle, id) { if let Some(path) = primary_path.as_deref() { let _ = crate::remove_download_assets(path, &self.app_handle).await; } } } log::error!("aria2 download {} failed: {}", id, error); self.clear_aria2_retry_state(id).await; self.forget_aria2_gid(id).await; self.emit_failed(id, error); self.release_registered_id(id).await; } } self.release_permit(id).await; } pub async fn clear_aria2_retry_state(&self, id: &str) { self.aria2_payloads.lock().await.remove(id); self.aria2_retry_strikes.lock().await.remove(id); } pub async fn cancel_aria2_retries(&self, id: &str) { self.aria2_retry_cancelled .lock() .await .insert(id.to_string()); } pub async fn allow_aria2_retries(&self, id: &str) { self.aria2_retry_cancelled.lock().await.remove(id); } pub fn aria2_gid_for_download(&self, id: &str) -> Option { self.aria2_gids .read() .unwrap() .iter() .find_map(|(gid, download_id)| (download_id == id).then(|| gid.clone())) } /// Remove every gid mapping for a download and discard buffered terminal /// events for those gids. Returns the most recently encountered gid. pub async fn forget_aria2_gid(&self, id: &str) -> Option { let removed = { let mut gids = self.aria2_gids.write().unwrap(); let removed: Vec = gids .iter() .filter(|(_, download_id)| *download_id == id) .map(|(gid, _)| gid.clone()) .collect(); for gid in &removed { gids.remove(gid); } removed }; if removed.is_empty() { return None; } let mut buffered = self.pending_completion.lock().await; for gid in &removed { buffered.remove(gid); log::info!("aria2 gid transition [{}]: forgot {}", id, gid); } removed.last().cloned() } async fn wait_permit_released(self: &Arc, id: &str) { loop { if !self.active_permits.lock().await.contains_key(id) { return; } let notified = self.notify_permit_released.notified(); if !self.active_permits.lock().await.contains_key(id) { return; } tokio::select! { _ = notified => {} _ = tokio::time::sleep(std::time::Duration::from_secs(5)) => {} } } } /// Intercept transient `onDownloadError` events: backoff, re-issue /// `addUri`, and rotate the gid mapping. Permanent errors and exhausted /// strikes fall through to a hard `Failed` state. async fn handle_aria2_download_error(self: &Arc, gid: &str, error: String) { let id = { let gids = self.aria2_gids.read().unwrap(); gids.get(gid).cloned() }; let id = match id { Some(id) => id, None => { self.pending_completion.lock().await.insert( gid.to_string(), (String::new(), PendingOutcome::Error(error)), ); return; } }; if self.aria2_retry_cancelled.lock().await.contains(&id) { log::info!( "aria2 retry cancellation [{}]: ignoring error for gid {} during removal", id, gid ); return; } let payload = self.aria2_payloads.lock().await.get(&id).cloned(); if payload.is_none() { self.apply_completion(&id, PendingOutcome::Error(error)) .await; return; } let mut payload = payload.unwrap(); let strike = { let mut strikes = self.aria2_retry_strikes.lock().await; let entry = strikes.entry(id.clone()).or_insert(0); *entry }; let transient = is_retryable_aria2_error(&error); let strikes_left = strike < automatic_retry_limit(payload.max_tries); if !(transient && strikes_left) { self.apply_completion(&id, PendingOutcome::Error(error)) .await; return; } if is_aria2_range_mode_error(&error) { log::warn!( "aria2 range mode [{}]: server rejected bounded chunk ranges; restarting with a single connection", id ); payload.connections = Some(1); if let Err(cleanup_error) = remove_incompatible_aria2_range_state(self, &id).await { log::warn!( "aria2 range mode [{}]: failed to remove incompatible partial state: {}", id, cleanup_error ); } self.aria2_payloads .lock() .await .insert(id.clone(), payload.clone()); } let this = Arc::clone(self); let id_for_task = id.clone(); let error_for_emit = error.clone(); tauri::async_runtime::spawn(async move { let outcome = backoff_and_emit( strike, error_for_emit, this.wait_permit_released(&id_for_task), |reason| { use tauri::Emitter; let _ = this.app_handle.emit( "download-state", DownloadStateEvent::retrying(&id_for_task, reason), ); }, ) .await; if outcome == BackoffOutcome::Aborted { return; } if !this.active_permits.lock().await.contains_key(&id_for_task) { return; } if this .aria2_retry_cancelled .lock() .await .contains(&id_for_task) { return; } match this.spawner.add_uri(&id_for_task, &payload).await { Ok(new_gid) => { if this .aria2_retry_cancelled .lock() .await .contains(&id_for_task) { if let Err(error) = this.spawner.remove_uri(&new_gid).await { log::error!( "aria2 retry cancellation [{}]: failed to remove late gid {}: {}", id_for_task, new_gid, error ); } else { log::info!( "aria2 retry cancellation [{}]: removed late gid {}", id_for_task, new_gid ); return; } let retained_gid = new_gid.clone(); this.remember_gid(id_for_task.clone(), new_gid).await; log::warn!( "aria2 retry cancellation [{}]: retained late gid {} mapping for remove retry", id_for_task, retained_gid ); return; } this.aria2_retry_strikes .lock() .await .insert(id_for_task.clone(), strike + 1); this.emit_state(&id_for_task, DownloadStatus::Downloading); this.remember_gid(id_for_task.clone(), new_gid).await; } Err(retry_error) => { this.apply_completion(&id_for_task, PendingOutcome::Error(retry_error)) .await; } } }); } /// Entry point for the aria2 WS poller. Resolves gid -> id; if not yet /// stored, buffers the outcome for reconciliation by remember_gid. pub async fn handle_aria2_event(self: &Arc, gid: &str, outcome: PendingOutcome) { match outcome { PendingOutcome::Error(error) => { self.handle_aria2_download_error(gid, error).await; } other => { let id_opt = { let gids = self.aria2_gids.read().unwrap(); gids.get(gid).cloned() }; match id_opt { Some(id) => { self.apply_completion(&id, other).await; } None => { self.pending_completion .lock() .await .insert(gid.to_string(), (String::new(), other)); } } } } } /// Reorder a pending task up or down. Returns the new pending order. /// No-op at boundaries. Does not emit (membership unchanged); the caller /// (Tauri command) returns the order to the frontend. pub async fn move_in_queue( &self, id: &str, queue_id: &str, direction: QueueDirection, ) -> Vec { let mut pending = self.pending.lock().await; let queue_positions = pending .iter() .enumerate() .filter_map(|(index, task)| (task.queue_id == queue_id).then_some(index)) .collect::>(); let queue_pos = queue_positions .iter() .position(|index| pending[*index].id == id); if let Some(queue_pos) = queue_pos { let target = match direction { QueueDirection::Up => queue_pos.checked_sub(1), QueueDirection::Down => { if queue_pos + 1 < queue_positions.len() { Some(queue_pos + 1) } else { None } } }; if let Some(target) = target { pending.swap(queue_positions[queue_pos], queue_positions[target]); } } pending .iter() .filter(|task| task.queue_id == queue_id) .map(|task| task.id.clone()) .collect() } /// Remove a task from pending if present (used by remove_download). /// Does NOT release a permit (the caller handles active permits via /// release_permit if the task was already dispatched). pub async fn remove_from_pending(&self, id: &str) -> bool { let mut pending = self.pending.lock().await; let before = pending.len(); pending.retain(|t| t.id != id); let removed = pending.len() < before; if removed { self.notify.notify_one(); } removed } } fn automatic_retry_limit(max_tries: Option) -> usize { max_tries.unwrap_or(0).max(0) as usize } fn aria2_attempt_limit(max_tries: Option) -> u32 { (automatic_retry_limit(max_tries) + 1) as u32 } fn is_retryable_aria2_error(error: &str) -> bool { is_transient_network_error(error) || is_aria2_range_mode_error(error) } fn is_aria2_rpc_unavailable(error: &str) -> bool { let lower = error.to_ascii_lowercase(); is_transient_network_error(error) || lower.contains("aria2 did not become ready") || lower.contains("connection refused") || lower.contains("failed to connect") || lower.contains("error trying to connect") || lower.contains("connection closed") || lower.contains("connection reset") } fn is_aria2_range_mode_error(error: &str) -> bool { let lower = error.to_ascii_lowercase(); lower.contains("invalid range header") || lower.contains("aria2 error code 8") || lower.contains("errorcode=8") } async fn remove_incompatible_aria2_range_state( manager: &QueueManager, id: &str, ) -> Result<(), String> { let Some(primary_path) = crate::download_ownership::primary_path_for_id(&manager.app_handle, id)? else { return Ok(()); }; crate::remove_download_assets(&primary_path, &manager.app_handle).await } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum BoundedRangeSupport { Supported, Unsupported, Unknown, } async fn effective_aria2_connections(id: &str, payload: &SpawnPayload) -> i32 { let requested = payload.connections.unwrap_or(1).max(1); if requested <= 1 { return requested; } for uri in crate::collect_download_uris(&payload.url, payload.mirrors.as_deref()) { if !is_http_uri(&uri) { continue; } match probe_bounded_range_support(&uri, payload).await { Ok(BoundedRangeSupport::Unsupported) => { log::warn!( "aria2 range probe [{}]: {} does not honor bounded byte ranges; using one connection", id, uri_host_for_log(&uri) ); return 1; } Ok(BoundedRangeSupport::Supported) => {} Ok(BoundedRangeSupport::Unknown) => { log::debug!( "aria2 range probe [{}]: {} range support unknown; keeping {} connections", id, uri_host_for_log(&uri), requested ); } Err(error) => { log::debug!( "aria2 range probe [{}]: {} probe failed: {}; keeping {} connections", id, uri_host_for_log(&uri), error, requested ); } } } requested } fn is_http_uri(uri: &str) -> bool { reqwest::Url::parse(uri) .ok() .is_some_and(|url| matches!(url.scheme(), "http" | "https")) } fn uri_host_for_log(uri: &str) -> String { reqwest::Url::parse(uri) .ok() .and_then(|url| url.host_str().map(str::to_string)) .unwrap_or_else(|| "".to_string()) } fn proxy_scheme(proxy: &str) -> Option { proxy .split_once("://") .map(|(scheme, _)| scheme.trim().to_ascii_lowercase()) } fn aria2_all_proxy_value(proxy: &str) -> Result, String> { let proxy = proxy.trim(); if proxy.is_empty() { return Ok(None); } if proxy.eq_ignore_ascii_case("none") { return Ok(Some(String::new())); } if proxy_scheme(proxy).is_some_and(|scheme| scheme.starts_with("socks")) { return Err( "SOCKS system proxies are not supported for normal file downloads because aria2 only accepts HTTP/HTTPS/FTP proxy URLs. Use an HTTP proxy endpoint for normal downloads, or use media downloads where yt-dlp supports SOCKS.".to_string(), ); } Ok(Some(proxy.to_string())) } async fn probe_bounded_range_support( uri: &str, payload: &SpawnPayload, ) -> Result { crate::ensure_reqwest_crypto_provider(); let mut builder = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::limited(5)) .timeout(std::time::Duration::from_secs(10)); if let Some(proxy) = payload .proxy .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) { if proxy.eq_ignore_ascii_case("none") { builder = builder.no_proxy(); } else { builder = builder.proxy(reqwest::Proxy::all(proxy).map_err(|error| error.to_string())?); } } let client = builder.build().map_err(|error| error.to_string())?; let request = client .get(uri) .header(reqwest::header::RANGE, "bytes=0-0") .header(reqwest::header::ACCEPT_ENCODING, "identity"); let request = apply_payload_headers(request, payload); let response = request.send().await.map_err(|error| error.to_string())?; let content_range = response .headers() .get(reqwest::header::CONTENT_RANGE) .and_then(|value| value.to_str().ok()); Ok(classify_bounded_range_response( response.status(), content_range, )) } fn apply_payload_headers( mut request: reqwest::RequestBuilder, payload: &SpawnPayload, ) -> reqwest::RequestBuilder { if let Some(user_agent) = payload .user_agent .as_deref() .filter(|value| !value.is_empty()) { request = request.header(reqwest::header::USER_AGENT, user_agent); } if let Some(cookies) = payload.cookies.as_deref().filter(|value| !value.is_empty()) { request = request.header(reqwest::header::COOKIE, cookies); } if let Some(headers) = payload.headers.as_deref() { for line in headers .lines() .map(str::trim) .filter(|line| !line.is_empty()) { let Some((name, value)) = line.split_once(':') else { continue; }; if name.trim().eq_ignore_ascii_case("range") { continue; } let Ok(name) = reqwest::header::HeaderName::from_bytes(name.trim().as_bytes()) else { continue; }; let Ok(value) = reqwest::header::HeaderValue::from_str(value.trim()) else { continue; }; request = request.header(name, value); } } if let Some(username) = payload .username .as_deref() .filter(|value| !value.is_empty()) { request = request.basic_auth(username, payload.password.as_deref()); } request } fn classify_bounded_range_response( status: reqwest::StatusCode, content_range: Option<&str>, ) -> BoundedRangeSupport { if status == reqwest::StatusCode::PARTIAL_CONTENT { return match content_range.and_then(parse_content_range_bounds) { Some((0, 0)) => BoundedRangeSupport::Supported, Some((0, _)) => BoundedRangeSupport::Unsupported, Some(_) => BoundedRangeSupport::Unknown, None => BoundedRangeSupport::Unknown, }; } if status.is_success() { BoundedRangeSupport::Unsupported } else { BoundedRangeSupport::Unknown } } fn parse_content_range_bounds(value: &str) -> Option<(u64, u64)> { let value = value.trim(); let (unit, range) = value.split_once(' ')?; if !unit.eq_ignore_ascii_case("bytes") { return None; } let (bounds, _) = range.split_once('/')?; let (start, end) = bounds.split_once('-')?; Some((start.trim().parse().ok()?, end.trim().parse().ok()?)) } /// Production spawner that delegates to the real aria2 RPC and yt-dlp runners. pub struct ProductionSpawner { app_handle: AppHandle, } impl ProductionSpawner { pub fn new(app_handle: AppHandle) -> Self { Self { app_handle } } async fn add_uri_rpc( &self, state: &crate::AppState, params: &serde_json::Value, ) -> Result { let deadline = std::time::Instant::now() + std::time::Duration::from_secs(15); loop { match crate::rpc_call( state.aria2_port.load(std::sync::atomic::Ordering::Relaxed), &state.aria2_secret, "aria2.addUri", params.clone(), ) .await { Ok(result) => return Ok(result), Err(error) => { if !is_aria2_rpc_unavailable(&error) || std::time::Instant::now() >= deadline { return Err(error); } tokio::time::sleep(std::time::Duration::from_millis(250)).await; } } } } } #[async_trait::async_trait] impl SidecarSpawner for ProductionSpawner { async fn add_uri(&self, id: &str, payload: &SpawnPayload) -> Result { let state = self.app_handle.state::(); let mut options = serde_json::Map::new(); let resolved_dest = crate::resolve_path(&payload.destination, &self.app_handle); if !crate::is_safe_path(&resolved_dest, &self.app_handle) { return Err("Path traversal blocked".to_string()); } let proxy_value = payload .proxy .as_deref() .map(aria2_all_proxy_value) .transpose()? .flatten(); options.insert( "dir".to_string(), serde_json::json!(resolved_dest.to_string_lossy().to_string()), ); let safe_filename = crate::download_ownership::canonical_download_filename(&payload.filename); options.insert("out".to_string(), serde_json::json!(safe_filename)); let conn = effective_aria2_connections(id, payload).await; options.insert("split".to_string(), serde_json::json!(conn.to_string())); options.insert( "max-connection-per-server".to_string(), serde_json::json!(conn.to_string()), ); let mt = aria2_attempt_limit(payload.max_tries); options.insert("max-tries".to_string(), serde_json::json!(mt.to_string())); options.insert("retry-wait".to_string(), serde_json::json!("2")); options.insert("continue".to_string(), serde_json::json!("true")); if let Some(speed) = payload .speed_limit .as_deref() .and_then(crate::normalize_speed_limit_for_aria2) { options.insert("max-download-limit".to_string(), serde_json::json!(speed)); } if let Some(user) = &payload.username { options.insert("http-user".to_string(), serde_json::json!(user)); } if let Some(pass) = &payload.password { options.insert("http-passwd".to_string(), serde_json::json!(pass)); } if let Some(chk) = &payload.checksum { let formatted_chk = if let Some((algo, digest)) = chk.split_once('=') { format!("{}={}", algo.to_ascii_lowercase(), digest) } else { chk.clone() }; options.insert("checksum".to_string(), serde_json::json!(formatted_chk)); } if let Some(ua) = &payload.user_agent { options.insert("user-agent".to_string(), serde_json::json!(ua)); } let mut header_list = Vec::new(); if let Some(cook) = &payload.cookies { header_list.push(format!("Cookie: {}", cook)); } if let Some(hdrs) = &payload.headers { for line in hdrs.lines() { if !line.trim().is_empty() { header_list.push(line.trim().to_string()); } } } if !header_list.is_empty() { options.insert("header".to_string(), serde_json::json!(header_list)); } if let Some(prox) = proxy_value { options.insert("all-proxy".to_string(), serde_json::json!(prox)); } let uris = crate::collect_download_uris(&payload.url, payload.mirrors.as_deref()); let params = serde_json::json!([uris, options]); match self.add_uri_rpc(&state, ¶ms).await { Ok(result) => { let gid = result.as_str().unwrap_or("").to_string(); if gid.is_empty() { Err("aria2.addUri returned an empty gid".to_string()) } else { log::info!("aria2 addUri [{}]: created gid {}", id, gid); Ok(gid) } } Err(e) => { log::error!("aria2 addUri [{}] failed: {}", id, e); Err(format!("aria2 addUri failed: {e}")) } } } async fn remove_uri(&self, gid: &str) -> Result<(), String> { let state = self.app_handle.state::(); let result = crate::rpc_call( state.aria2_port.load(std::sync::atomic::Ordering::Relaxed), &state.aria2_secret, "aria2.forceRemove", serde_json::json!([gid]), ) .await?; match result.as_str() { Some(returned_gid) if returned_gid == gid => Ok(()), Some(returned_gid) => Err(format!( "aria2.forceRemove returned unexpected gid {returned_gid}, expected {gid}" )), None => Err("aria2.forceRemove returned a non-string result".to_string()), } } async fn run_media(&self, id: &str, payload: &SpawnPayload) -> Result<(), String> { let state = self.app_handle.state::(); let mut cancel_rx = state .download_coordinator .register_media(id.to_string()) .await?; let outcome = crate::start_media_download_internal( self.app_handle.clone(), id, payload.url.clone(), payload.destination.clone(), payload.filename.clone(), payload.format_selector.clone(), payload.cookie_source.clone(), payload.speed_limit.clone(), payload.username.clone(), payload.password.clone(), payload.headers.clone(), payload.cookies.clone(), payload.proxy.clone(), payload.user_agent.clone(), payload.max_tries, &mut cancel_rx, ) .await; if let Ok(path) = outcome.as_ref() { let _ = crate::download_ownership::set_primary_path(&self.app_handle, id, path); if let Some(file_name) = path.file_name().and_then(|name| name.to_str()) { use tauri::Emitter; let _ = self.app_handle.emit( "download-state", crate::ipc::DownloadStateEvent::completed_with_file(id, file_name), ); } } let _ = state .download_coordinator .finish_media(id.to_string()) .await; outcome.map(|_| ()) } } #[derive(Debug, Clone, Deserialize, TS)] #[ts(export, export_to = "../../src/bindings/")] pub struct EnqueueItem { pub id: String, pub queue_id: String, pub url: String, pub destination: String, pub filename: String, pub connections: Option, pub speed_limit: Option, pub username: Option, pub password: Option, pub headers: Option, pub checksum: Option, pub cookies: Option, pub mirrors: Option, pub user_agent: Option, pub max_tries: Option, pub proxy: Option, pub format_selector: Option, pub cookie_source: Option, pub is_media: Option, #[serde(default)] #[ts(optional)] pub lifecycle_generation: Option, } impl EnqueueItem { pub fn into_task(self) -> QueuedTask { let media = self.is_media.unwrap_or(false); let kind = if media { TaskKind::Media } else { TaskKind::Aria2 }; let id = self.id.clone(); QueuedTask { id, queue_id: self.queue_id, kind, payload: SpawnPayload { url: self.url, destination: self.destination, filename: self.filename, connections: self.connections, speed_limit: self.speed_limit, username: self.username, password: self.password, headers: self.headers, checksum: self.checksum, cookies: self.cookies, mirrors: self.mirrors, user_agent: self.user_agent, max_tries: self.max_tries, proxy: self.proxy, format_selector: self.format_selector, cookie_source: self.cookie_source, is_media: media, }, } } } #[cfg(test)] mod tests { use super::*; #[test] fn bounded_range_probe_accepts_exact_requested_byte() { assert_eq!( classify_bounded_range_response( reqwest::StatusCode::PARTIAL_CONTENT, Some("bytes 0-0/383882118"), ), BoundedRangeSupport::Supported ); } #[test] fn bounded_range_probe_accepts_case_insensitive_content_range_unit() { assert_eq!( classify_bounded_range_response( reqwest::StatusCode::PARTIAL_CONTENT, Some("Bytes 0-0/383882118"), ), BoundedRangeSupport::Supported ); } #[test] fn aria2_proxy_value_rejects_socks_proxies() { assert_eq!(aria2_all_proxy_value("none").unwrap().as_deref(), Some("")); assert_eq!( aria2_all_proxy_value("http://127.0.0.1:8080") .unwrap() .as_deref(), Some("http://127.0.0.1:8080") ); assert!(aria2_all_proxy_value("socks5://127.0.0.1:1080") .unwrap_err() .contains("SOCKS system proxies are not supported")); } #[test] fn bounded_range_probe_rejects_server_that_expands_to_end() { assert_eq!( classify_bounded_range_response( reqwest::StatusCode::PARTIAL_CONTENT, Some("bytes 0-383882117/383882118"), ), BoundedRangeSupport::Unsupported ); } #[test] fn bounded_range_probe_rejects_ignored_range_request() { assert_eq!( classify_bounded_range_response(reqwest::StatusCode::OK, None), BoundedRangeSupport::Unsupported ); } #[test] fn aria2_range_code_is_retryable_without_global_no_uri_retry() { assert!(is_retryable_aria2_error( "aria2 error code 8: No URI available." )); assert!(!is_retryable_aria2_error("No URI available.")); } #[test] fn aria2_startup_rpc_errors_are_retryable() { assert!(is_aria2_rpc_unavailable( "error trying to connect: tcp connect error: Connection refused" )); assert!(is_aria2_rpc_unavailable( "aria2 did not become ready: connection refused" )); assert!(!is_aria2_rpc_unavailable( "aria2 error code 3: Resource not found" )); } }