diff --git a/soulseek-rs-lib/src/client.rs b/soulseek-rs-lib/src/client.rs index b4c1a93..0f6f064 100644 --- a/soulseek-rs-lib/src/client.rs +++ b/soulseek-rs-lib/src/client.rs @@ -1,5 +1,6 @@ use crate::actor::ActorHandle; use crate::actor::server_actor::{PeerAddress, ServerActor, ServerMessage}; +use crate::download_store::{DownloadStore, collect_failed_tokens}; use crate::types::DownloadStatus; use crate::utils::logger; use crate::{ @@ -87,10 +88,10 @@ pub enum ClientOperation { } pub struct ClientContext { pub peer_registry: Option, + pub downloads: DownloadStore, sender: Option>, server_sender: Option>, searches: HashMap, - downloads: Vec, actor_system: Arc, } impl Default for ClientContext { @@ -98,64 +99,51 @@ impl Default for ClientContext { Self::new() } } + +// Thin delegating shims so existing callers (peer/listen, peer/download_peer, +// tests) keep working while download state lives in DownloadStore. impl ClientContext { pub fn add_download(&mut self, download: Download) { - self.downloads.push(download); + self.downloads.add(download); } pub fn remove_download(&mut self, token: u32) { - self.downloads.retain(|d| d.token != token); + self.downloads.remove(token); } pub fn get_download_by_token(&self, token: u32) -> Option<&Download> { - self.downloads.iter().find(|d| d.token == token) + self.downloads.get_by_token(token) } - pub fn get_download_by_token_mut( &mut self, token: u32, ) -> Option<&mut Download> { - self.downloads.iter_mut().find(|d| d.token == token) + self.downloads.get_by_token_mut(token) } pub fn get_download_by_file_mut( &mut self, username: &str, filename: &str, ) -> Option<&mut Download> { - self.downloads - .iter_mut() - .find(|d| d.username == username && d.filename == filename) + self.downloads.get_by_file_mut(username, filename) } pub fn get_download_tokens(&self) -> Vec { - self.downloads.iter().map(|d| d.token).collect() + self.downloads.tokens() } pub fn get_downloads(&self) -> &Vec { - &self.downloads + self.downloads.list() } - pub fn update_download_with_status( &mut self, token: u32, status: DownloadStatus, ) { - if let Some(download) = self.get_download_by_token_mut(token) { - download.status = status; - } + self.downloads.update_status(token, status); } - pub fn remove_queued_download_by_file( &mut self, username: &str, filename: &str, ) -> bool { - let Some(index) = self.downloads.iter().position(|download| { - download.username == username - && download.filename == filename - && matches!(download.status, DownloadStatus::Queued) - }) else { - return false; - }; - - self.downloads.remove(index); - true + self.downloads.remove_queued_by_file(username, filename) } } #[test] @@ -300,7 +288,7 @@ impl ClientContext { sender: None, server_sender: None, searches: HashMap::new(), - downloads: Vec::new(), + downloads: DownloadStore::new(), actor_system, } } @@ -527,67 +515,23 @@ impl Client { } pub fn pause_download(&self, username: &str, filename: &str) -> bool { - let mut context = match self.context.write_safe() { - Ok(c) => c, + match self.context.write_safe() { + Ok(mut ctx) => ctx.downloads.pause_by_file(username, filename), Err(e) => { error!("[client] pause_download: {}", e); - return false; + false } - }; - let Some(download) = - context.get_download_by_file_mut(username, filename) - else { - return false; - }; - - let paused_status = match &download.status { - DownloadStatus::InProgress { - bytes_downloaded, - total_bytes, - .. - } => DownloadStatus::Paused { - bytes_downloaded: *bytes_downloaded, - total_bytes: *total_bytes, - }, - DownloadStatus::Paused { .. } => return true, - _ => return false, - }; - - download.status = paused_status.clone(); - let _ = download.sender.send(paused_status); - true + } } pub fn resume_download(&self, username: &str, filename: &str) -> bool { - let mut context = match self.context.write_safe() { - Ok(c) => c, + match self.context.write_safe() { + Ok(mut ctx) => ctx.downloads.resume_by_file(username, filename), Err(e) => { error!("[client] resume_download: {}", e); - return false; + false } - }; - let Some(download) = - context.get_download_by_file_mut(username, filename) - else { - return false; - }; - - let resumed_status = match &download.status { - DownloadStatus::Paused { - bytes_downloaded, - total_bytes, - } => DownloadStatus::InProgress { - bytes_downloaded: *bytes_downloaded, - total_bytes: *total_bytes, - speed_bytes_per_sec: 0.0, - }, - DownloadStatus::InProgress { .. } => return true, - _ => return false, - }; - - download.status = resumed_status.clone(); - let _ = download.sender.send(resumed_status); - true + } } pub fn remove_queued_download( @@ -597,7 +541,7 @@ impl Client { ) -> bool { match self.context.write_safe() { Ok(mut ctx) => { - ctx.remove_queued_download_by_file(username, filename) + ctx.downloads.remove_queued_by_file(username, filename) } Err(e) => { error!("[client] remove_queued_download: {}", e); @@ -662,39 +606,33 @@ impl Client { username: &str, filename: Option<&str>, ) { - let failed_downloads: Vec<_> = match client_context.read_safe() { - Ok(context) => context - .get_downloads() - .iter() - .filter(|download| { - download.username == username - && (filename.is_none_or(|f| download.filename == *f)) - }) - .map(|download| { - let _ = download.sender.send(DownloadStatus::Failed); - download.token - }) - .collect(), + let failed_tokens = match client_context.read_safe() { + Ok(context) => { + collect_failed_tokens(&context.downloads, username, filename) + } Err(e) => { error!("[client] process_failed_uploads read: {}", e); return; } }; - failed_downloads.iter().for_each(|token| { - match client_context.write_safe() { - Ok(mut context) => { - context.update_download_with_status( - *token, - DownloadStatus::Failed, - ); - context.remove_download(*token); - } - Err(e) => { - error!("[client] process_failed_uploads write: {}", e); + if failed_tokens.is_empty() { + return; + } + + match client_context.write_safe() { + Ok(mut context) => { + for token in failed_tokens { + context + .downloads + .update_status(token, DownloadStatus::Failed); + context.downloads.remove(token); } } - }); + Err(e) => { + error!("[client] process_failed_uploads write: {}", e); + } + } } fn listen_to_client_operations( diff --git a/soulseek-rs-lib/src/download_store.rs b/soulseek-rs-lib/src/download_store.rs new file mode 100644 index 0000000..cfba0a6 --- /dev/null +++ b/soulseek-rs-lib/src/download_store.rs @@ -0,0 +1,270 @@ +use crate::types::{Download, DownloadStatus}; + +#[derive(Default)] +pub struct DownloadStore { + downloads: Vec, +} + +impl DownloadStore { + pub fn new() -> Self { + Self::default() + } + + pub fn add(&mut self, download: Download) { + self.downloads.push(download); + } + + pub fn remove(&mut self, token: u32) { + self.downloads.retain(|d| d.token != token); + } + + pub fn get_by_token(&self, token: u32) -> Option<&Download> { + self.downloads.iter().find(|d| d.token == token) + } + + pub fn get_by_token_mut(&mut self, token: u32) -> Option<&mut Download> { + self.downloads.iter_mut().find(|d| d.token == token) + } + + pub fn get_by_file_mut( + &mut self, + username: &str, + filename: &str, + ) -> Option<&mut Download> { + self.downloads + .iter_mut() + .find(|d| d.username == username && d.filename == filename) + } + + pub fn tokens(&self) -> Vec { + self.downloads.iter().map(|d| d.token).collect() + } + + pub fn list(&self) -> &Vec { + &self.downloads + } + + pub fn update_status(&mut self, token: u32, status: DownloadStatus) { + if let Some(download) = self.get_by_token_mut(token) { + download.status = status; + } + } + + pub fn remove_queued_by_file( + &mut self, + username: &str, + filename: &str, + ) -> bool { + let Some(index) = self.downloads.iter().position(|download| { + download.username == username + && download.filename == filename + && matches!(download.status, DownloadStatus::Queued) + }) else { + return false; + }; + + self.downloads.remove(index); + true + } + + pub fn pause_by_file(&mut self, username: &str, filename: &str) -> bool { + let Some(download) = self.get_by_file_mut(username, filename) else { + return false; + }; + + let paused_status = match &download.status { + DownloadStatus::InProgress { + bytes_downloaded, + total_bytes, + .. + } => DownloadStatus::Paused { + bytes_downloaded: *bytes_downloaded, + total_bytes: *total_bytes, + }, + DownloadStatus::Paused { .. } => return true, + _ => return false, + }; + + download.status = paused_status.clone(); + let _ = download.sender.send(paused_status); + true + } + + pub fn resume_by_file(&mut self, username: &str, filename: &str) -> bool { + let Some(download) = self.get_by_file_mut(username, filename) else { + return false; + }; + + let resumed_status = match &download.status { + DownloadStatus::Paused { + bytes_downloaded, + total_bytes, + } => DownloadStatus::InProgress { + bytes_downloaded: *bytes_downloaded, + total_bytes: *total_bytes, + speed_bytes_per_sec: 0.0, + }, + DownloadStatus::InProgress { .. } => return true, + _ => return false, + }; + + download.status = resumed_status.clone(); + let _ = download.sender.send(resumed_status); + true + } +} + +/// Returns the tokens of downloads matching `username` (and optionally a +/// `filename`) after notifying their senders of `Failed`. Caller is +/// responsible for then calling `update_status` and `remove` for each token, +/// typically under a write lock. +pub fn collect_failed_tokens( + store: &DownloadStore, + username: &str, + filename: Option<&str>, +) -> Vec { + store + .list() + .iter() + .filter(|d| { + d.username == username && filename.is_none_or(|f| d.filename == *f) + }) + .map(|d| { + let _ = d.sender.send(DownloadStatus::Failed); + d.token + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::mpsc; + + fn make_download(token: u32, status: DownloadStatus) -> Download { + Download { + username: "peer".to_string(), + filename: format!("file-{}.mp3", token), + token, + size: 100, + download_directory: "test".to_string(), + status, + sender: mpsc::channel().0, + } + } + + #[test] + fn add_get_remove_roundtrip() { + let mut store = DownloadStore::new(); + store.add(make_download(123, DownloadStatus::Queued)); + + assert!(store.get_by_token(123).is_some()); + assert_eq!(store.tokens(), vec![123]); + assert_eq!(store.list().len(), 1); + + store.remove(123); + assert!(store.get_by_token(123).is_none()); + assert!(store.list().is_empty()); + } + + #[test] + fn pause_then_resume_in_progress_download() { + let mut store = DownloadStore::new(); + let (tx, rx) = mpsc::channel(); + let mut download = make_download( + 1, + DownloadStatus::InProgress { + bytes_downloaded: 25, + total_bytes: 100, + speed_bytes_per_sec: 10.0, + }, + ); + download.sender = tx; + store.add(download); + + assert!(store.pause_by_file("peer", "file-1.mp3")); + assert!(matches!( + store.get_by_token(1).unwrap().status, + DownloadStatus::Paused { + bytes_downloaded: 25, + total_bytes: 100 + } + )); + assert!(matches!( + rx.try_recv().unwrap(), + DownloadStatus::Paused { + bytes_downloaded: 25, + total_bytes: 100 + } + )); + + assert!(store.resume_by_file("peer", "file-1.mp3")); + assert!(matches!( + store.get_by_token(1).unwrap().status, + DownloadStatus::InProgress { + bytes_downloaded: 25, + total_bytes: 100, + speed_bytes_per_sec: 0.0 + } + )); + } + + #[test] + fn remove_queued_skips_active_downloads() { + let mut store = DownloadStore::new(); + store.add(make_download(123, DownloadStatus::Queued)); + store.add(make_download( + 456, + DownloadStatus::InProgress { + bytes_downloaded: 25, + total_bytes: 100, + speed_bytes_per_sec: 10.0, + }, + )); + // Override second download's filename so they don't collide + store.get_by_token_mut(456).unwrap().filename = + "active.mp3".to_string(); + store.get_by_token_mut(123).unwrap().filename = + "queued.mp3".to_string(); + + assert!(store.remove_queued_by_file("peer", "queued.mp3")); + assert!(!store.remove_queued_by_file("peer", "active.mp3")); + assert!(store.get_by_token(123).is_none()); + assert!(store.get_by_token(456).is_some()); + } + + #[test] + fn collect_failed_tokens_notifies_and_lists_matching() { + let mut store = DownloadStore::new(); + let (tx_match, rx_match) = mpsc::channel(); + let (tx_other_user, _rx_other_user) = mpsc::channel(); + let (tx_other_file, _rx_other_file) = mpsc::channel(); + + let mut a = make_download(1, DownloadStatus::Queued); + a.sender = tx_match; + a.username = "peer".to_string(); + a.filename = "song.mp3".to_string(); + + let mut b = make_download(2, DownloadStatus::Queued); + b.sender = tx_other_user; + b.username = "other".to_string(); + b.filename = "song.mp3".to_string(); + + let mut c = make_download(3, DownloadStatus::Queued); + c.sender = tx_other_file; + c.username = "peer".to_string(); + c.filename = "different.mp3".to_string(); + + store.add(a); + store.add(b); + store.add(c); + + let tokens = collect_failed_tokens(&store, "peer", Some("song.mp3")); + + assert_eq!(tokens, vec![1]); + assert!(matches!( + rx_match.try_recv().unwrap(), + DownloadStatus::Failed + )); + } +} diff --git a/soulseek-rs-lib/src/lib.rs b/soulseek-rs-lib/src/lib.rs index d4fd935..b35e2af 100644 --- a/soulseek-rs-lib/src/lib.rs +++ b/soulseek-rs-lib/src/lib.rs @@ -2,6 +2,7 @@ pub mod actor; pub mod client; pub mod dispatcher; +pub mod download_store; pub mod error; pub mod message; pub mod peer;