chore: refactor downloads

This commit is contained in:
Michel de Graaf
2025-11-18 14:11:25 +01:00
parent 467de9cb91
commit 79da19ed0f
5 changed files with 114 additions and 59 deletions
+99 -36
View File
@@ -90,7 +90,7 @@ pub struct ClientContext {
sender: Option<Sender<ClientOperation>>, sender: Option<Sender<ClientOperation>>,
server_sender: Option<Sender<ServerMessage>>, server_sender: Option<Sender<ServerMessage>>,
searches: HashMap<String, Search>, searches: HashMap<String, Search>,
pub download_tokens: HashMap<u32, Download>, downloads: Vec<Download>,
actor_system: Arc<ActorSystem>, actor_system: Arc<ActorSystem>,
} }
impl Default for ClientContext { impl Default for ClientContext {
@@ -98,6 +98,68 @@ impl Default for ClientContext {
Self::new() Self::new()
} }
} }
impl ClientContext {
pub fn add_download(&mut self, download: Download) {
self.downloads.push(download);
}
pub fn remove_download(&mut self, token: u32) {
self.downloads.retain(|d| d.token != token);
}
pub fn get_download_by_token(&self, token: u32) -> Option<&Download> {
self.downloads.iter().find(|d| d.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)
}
pub fn get_download_tokens(&self) -> Vec<u32> {
self.downloads.iter().map(|d| d.token).collect()
}
pub fn get_downloads(&self) -> &Vec<Download> {
&self.downloads
}
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;
}
}
}
#[test]
fn test_client_context_downloads() {
let mut context = ClientContext::new();
let token = 123;
let new_token = 1234;
let download = Download {
username: "test".to_string(),
filename: "test.txt".to_string(),
token,
size: 100,
download_directory: "test".to_string(),
status: DownloadStatus::Queued,
sender: mpsc::channel().0,
};
context.add_download(download);
assert!(context.get_download_by_token(123).is_some());
assert_eq!(context.get_download_tokens(), vec![123]);
assert_eq!(context.get_downloads().len(), 1);
if let Some(download) = context.get_download_by_token_mut(token) {
assert_eq!(download.token, token);
download.token = new_token
}
assert!(context.get_download_by_token(new_token).is_some());
assert_eq!(context.get_download_tokens(), vec![new_token]);
context.remove_download(new_token);
assert_eq!(context.get_downloads().len(), 0);
assert!(context.get_download_by_token(1234).is_none());
}
impl ClientContext { impl ClientContext {
#[must_use] #[must_use]
@@ -114,7 +176,7 @@ impl ClientContext {
sender: None, sender: None,
server_sender: None, server_sender: None,
searches: HashMap::new(), searches: HashMap::new(),
download_tokens: HashMap::new(), downloads: Vec::new(),
actor_system, actor_system,
} }
} }
@@ -335,11 +397,12 @@ impl Client {
token, token,
size, size,
download_directory, download_directory,
status: DownloadStatus::Queued,
sender: download_sender, sender: download_sender,
}; };
let mut context = self.context.write().unwrap(); let mut context = self.context.write().unwrap();
context.download_tokens.insert(token, download.clone()); context.add_download(download.clone());
let download_initiated = let download_initiated =
if let Some(ref registry) = context.peer_registry { if let Some(ref registry) = context.peer_registry {
@@ -354,6 +417,7 @@ impl Client {
if !download_initiated { if !download_initiated {
let _ = download.sender.send(DownloadStatus::Failed); let _ = download.sender.send(DownloadStatus::Failed);
self.context.write().unwrap().update_download_with_status(token, DownloadStatus::Failed);
} }
Ok(download_receiver) Ok(download_receiver)
@@ -365,25 +429,24 @@ impl Client {
filename: Option<&str>, filename: Option<&str>,
) { ) {
let context = client_context.read().unwrap(); let context = client_context.read().unwrap();
let download_tokens = &context.download_tokens; let failed_downloads: Vec<_> = context
let failed_downloads: Vec<_> = download_tokens .get_downloads()
.iter() .iter()
.filter(|(_, download)| { .filter(|download| {
download.username == username download.username == username
&& (filename.is_none_or(|f| download.filename == *f)) && (filename.is_none_or(|f| download.filename == *f))
}) })
.map(|(token, download)| { .map(|download| {
let _ = download.sender.send(DownloadStatus::Failed); let _ = download.sender.send(DownloadStatus::Failed);
*token download.token
}) })
.collect(); .collect();
drop(context);
failed_downloads.iter().for_each(|token| { failed_downloads.iter().for_each(|token| {
client_context let mut context = client_context.write().unwrap();
.write() context.update_download_with_status(*token, DownloadStatus::Failed);
.unwrap() context.remove_download(*token);
.download_tokens
.remove(token);
}); });
} }
@@ -459,8 +522,7 @@ impl Client {
let client_context = let client_context =
client_context.read().unwrap(); client_context.read().unwrap();
client_context client_context
.download_tokens .get_download_by_token(token)
.get(&token)
.cloned() .cloned()
}; };
let own_username = own_username.clone(); let own_username = own_username.clone();
@@ -489,16 +551,18 @@ impl Client {
match filename { match filename {
Some(filename) => { Some(filename) => {
match download_peer.download_file( match download_peer.download_file(
client_context_clone, client_context_clone.clone(),
Some(download.clone()), Some(download.clone()),
None None
) { ) {
Ok((download, filename)) => { Ok((download, filename)) => {
let _ = download.sender.send(DownloadStatus::Completed); let _ = download.sender.send(DownloadStatus::Completed);
client_context_clone.write().unwrap().update_download_with_status(download.token, DownloadStatus::Completed);
info!("Successfully downloaded {} bytes to {}", download.size, filename); info!("Successfully downloaded {} bytes to {}", download.size, filename);
} }
Err(e) => { Err(e) => {
let _ = download.sender.send(DownloadStatus::Failed); let _ = download.sender.send(DownloadStatus::Failed);
client_context_clone.write().unwrap().update_download_with_status(download.token, DownloadStatus::Failed);
error!( error!(
"Failed to download file '{}' from {}:{} (token: {}) - Error: {}", "Failed to download file '{}' from {}:{} (token: {}) - Error: {}",
filename, peer.host, peer.port, download.token, e filename, peer.host, peer.port, download.token, e
@@ -631,38 +695,36 @@ impl Client {
) => { ) => {
let mut context = client_context.write().unwrap(); let mut context = client_context.write().unwrap();
let key_to_remove = context let download_to_update =
.download_tokens context.get_downloads().iter().find_map(|d| {
.iter()
.find_map(|(key, d)| {
if d.username == username if d.username == username
&& d.filename == transfer.filename && d.filename == transfer.filename
{ {
Some((*key, d.clone())) Some((d.token, d.clone()))
} else { } else {
None None
} }
}); });
if let Some((key, download)) = key_to_remove { if let Some((old_token, download)) =
download_to_update
{
trace!( trace!(
"[client] UpdateDownloadTokens found {key}, transfer: {:?}", "[client] UpdateDownloadTokens found {old_token}, transfer: {:?}",
transfer transfer
); );
context.download_tokens.insert( context.add_download(Download {
transfer.token, username: username.clone(),
Download { filename: transfer.filename,
username: username.clone(), token: transfer.token,
filename: transfer.filename, size: transfer.size,
token: transfer.token, download_directory: download
size: transfer.size, .download_directory,
download_directory: download status: download.status.clone(),
.download_directory, sender: download.sender.clone(),
sender: download.sender.clone(), });
}, context.remove_download(old_token);
);
context.download_tokens.remove(&key);
} }
} }
ClientOperation::UploadFailed(username, filename) => { ClientOperation::UploadFailed(username, filename) => {
@@ -745,6 +807,7 @@ impl Client {
download.size download.size
); );
let _ = download.sender.send(DownloadStatus::Completed); let _ = download.sender.send(DownloadStatus::Completed);
client_context.write().unwrap().update_download_with_status(download.token, DownloadStatus::Completed);
} }
Err(e) => { Err(e) => {
trace!("[client] failed to download: {}", e); trace!("[client] failed to download: {}", e);
+2 -2
View File
@@ -13,7 +13,7 @@ pub mod utils;
pub mod prelude { pub mod prelude {
pub use crate::actor::server_actor::PeerAddress; pub use crate::actor::server_actor::PeerAddress;
pub use crate::types::{ pub use crate::types::{
DownloadResult, DownloadStatus, File, Search, SearchResult, Transfer, DownloadStatus, File, Search, SearchResult, Transfer,
}; };
pub use crate::{debug, error, info, trace, warn}; pub use crate::{debug, error, info, trace, warn};
} }
@@ -23,5 +23,5 @@ pub use actor::server_actor::PeerAddress;
pub use client::{Client, ClientSettings}; pub use client::{Client, ClientSettings};
pub use error::{Result, SoulseekRs}; pub use error::{Result, SoulseekRs};
pub use types::{ pub use types::{
DownloadResult, DownloadStatus, File, Search, SearchResult, Transfer, DownloadStatus, File, Search, SearchResult, Transfer,
}; };
+9 -9
View File
@@ -242,8 +242,7 @@ impl DownloadPeer {
let client_guard = client_context let client_guard = client_context
.read() .read()
.map_err(|_| DownloadError::LockPoisoned)?; .map_err(|_| DownloadError::LockPoisoned)?;
let download_info = let download_info = client_guard.get_download_by_token(token_u32).cloned();
client_guard.download_tokens.get(&token_u32).cloned();
drop(client_guard); drop(client_guard);
download_info.ok_or(DownloadError::TokenNotFound(token_u32)) download_info.ok_or(DownloadError::TokenNotFound(token_u32))
@@ -318,13 +317,13 @@ impl DownloadPeer {
0.0 0.0
}; };
let _ = let status = DownloadStatus::InProgress {
dl.sender.send(DownloadStatus::InProgress { bytes_downloaded: processor.total_bytes as u64,
bytes_downloaded: processor.total_bytes total_bytes: dl.size,
as u64, speed_bytes_per_sec: speed,
total_bytes: dl.size, };
speed_bytes_per_sec: speed, let _ = dl.sender.send(status.clone());
}); client_context.write().unwrap().update_download_with_status(dl.token, status);
last_update_time = Instant::now(); last_update_time = Instant::now();
} }
@@ -422,6 +421,7 @@ impl DownloadPeer {
if let Some(ref dl) = download { if let Some(ref dl) = download {
let _ = dl.sender.send(DownloadStatus::Queued); let _ = dl.sender.send(DownloadStatus::Queued);
client_context.write().unwrap().update_download_with_status(dl.token, DownloadStatus::Queued);
} }
let mut stream = match stream { let mut stream = match stream {
+3 -3
View File
@@ -84,11 +84,10 @@ fn extract_download_from_buffer(
); );
let context = client_context.read().unwrap(); let context = client_context.read().unwrap();
let download = context.download_tokens.get(&token).cloned(); let download = context.get_download_by_token(token).cloned();
if download.is_none() { if download.is_none() {
let download_tokens: Vec<u32> = let download_tokens = context.get_download_tokens();
context.download_tokens.keys().cloned().collect();
trace!("[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", token, download_tokens); trace!("[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}", token, download_tokens);
} }
@@ -158,6 +157,7 @@ fn handle_file_connection(
) { ) {
Ok((download, filename)) => { Ok((download, filename)) => {
let _ = download.sender.send(DownloadStatus::Completed); let _ = download.sender.send(DownloadStatus::Completed);
context.client_context.write().unwrap().update_download_with_status(download.token, DownloadStatus::Completed);
info!( info!(
"Successfully downloaded {} bytes to {}", "Successfully downloaded {} bytes to {}",
download.size, filename download.size, filename
+1 -9
View File
@@ -94,16 +94,8 @@ pub struct Download {
pub token: u32, pub token: u32,
pub size: u64, pub size: u64,
pub download_directory: String, pub download_directory: String,
pub sender: Sender<DownloadStatus>,
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct DownloadResult {
pub filename: String,
pub username: String,
pub status: DownloadStatus, pub status: DownloadStatus,
pub elapsed_time: std::time::Duration, pub sender: Sender<DownloadStatus>,
} }
#[derive(Debug, Clone)] #[derive(Debug, Clone)]