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>>,
server_sender: Option<Sender<ServerMessage>>,
searches: HashMap<String, Search>,
pub download_tokens: HashMap<u32, Download>,
downloads: Vec<Download>,
actor_system: Arc<ActorSystem>,
}
impl Default for ClientContext {
@@ -98,6 +98,68 @@ impl Default for ClientContext {
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 {
#[must_use]
@@ -114,7 +176,7 @@ impl ClientContext {
sender: None,
server_sender: None,
searches: HashMap::new(),
download_tokens: HashMap::new(),
downloads: Vec::new(),
actor_system,
}
}
@@ -335,11 +397,12 @@ impl Client {
token,
size,
download_directory,
status: DownloadStatus::Queued,
sender: download_sender,
};
let mut context = self.context.write().unwrap();
context.download_tokens.insert(token, download.clone());
context.add_download(download.clone());
let download_initiated =
if let Some(ref registry) = context.peer_registry {
@@ -354,6 +417,7 @@ impl Client {
if !download_initiated {
let _ = download.sender.send(DownloadStatus::Failed);
self.context.write().unwrap().update_download_with_status(token, DownloadStatus::Failed);
}
Ok(download_receiver)
@@ -365,25 +429,24 @@ impl Client {
filename: Option<&str>,
) {
let context = client_context.read().unwrap();
let download_tokens = &context.download_tokens;
let failed_downloads: Vec<_> = download_tokens
let failed_downloads: Vec<_> = context
.get_downloads()
.iter()
.filter(|(_, download)| {
.filter(|download| {
download.username == username
&& (filename.is_none_or(|f| download.filename == *f))
})
.map(|(token, download)| {
.map(|download| {
let _ = download.sender.send(DownloadStatus::Failed);
*token
download.token
})
.collect();
drop(context);
failed_downloads.iter().for_each(|token| {
client_context
.write()
.unwrap()
.download_tokens
.remove(token);
let mut context = client_context.write().unwrap();
context.update_download_with_status(*token, DownloadStatus::Failed);
context.remove_download(*token);
});
}
@@ -459,8 +522,7 @@ impl Client {
let client_context =
client_context.read().unwrap();
client_context
.download_tokens
.get(&token)
.get_download_by_token(token)
.cloned()
};
let own_username = own_username.clone();
@@ -489,16 +551,18 @@ impl Client {
match filename {
Some(filename) => {
match download_peer.download_file(
client_context_clone,
client_context_clone.clone(),
Some(download.clone()),
None
) {
Ok((download, filename)) => {
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);
}
Err(e) => {
let _ = download.sender.send(DownloadStatus::Failed);
client_context_clone.write().unwrap().update_download_with_status(download.token, DownloadStatus::Failed);
error!(
"Failed to download file '{}' from {}:{} (token: {}) - Error: {}",
filename, peer.host, peer.port, download.token, e
@@ -631,38 +695,36 @@ impl Client {
) => {
let mut context = client_context.write().unwrap();
let key_to_remove = context
.download_tokens
.iter()
.find_map(|(key, d)| {
let download_to_update =
context.get_downloads().iter().find_map(|d| {
if d.username == username
&& d.filename == transfer.filename
{
Some((*key, d.clone()))
Some((d.token, d.clone()))
} else {
None
}
});
if let Some((key, download)) = key_to_remove {
if let Some((old_token, download)) =
download_to_update
{
trace!(
"[client] UpdateDownloadTokens found {key}, transfer: {:?}",
"[client] UpdateDownloadTokens found {old_token}, transfer: {:?}",
transfer
);
context.download_tokens.insert(
transfer.token,
Download {
username: username.clone(),
filename: transfer.filename,
token: transfer.token,
size: transfer.size,
download_directory: download
.download_directory,
sender: download.sender.clone(),
},
);
context.download_tokens.remove(&key);
context.add_download(Download {
username: username.clone(),
filename: transfer.filename,
token: transfer.token,
size: transfer.size,
download_directory: download
.download_directory,
status: download.status.clone(),
sender: download.sender.clone(),
});
context.remove_download(old_token);
}
}
ClientOperation::UploadFailed(username, filename) => {
@@ -745,6 +807,7 @@ impl Client {
download.size
);
let _ = download.sender.send(DownloadStatus::Completed);
client_context.write().unwrap().update_download_with_status(download.token, DownloadStatus::Completed);
}
Err(e) => {
trace!("[client] failed to download: {}", e);
+2 -2
View File
@@ -13,7 +13,7 @@ pub mod utils;
pub mod prelude {
pub use crate::actor::server_actor::PeerAddress;
pub use crate::types::{
DownloadResult, DownloadStatus, File, Search, SearchResult, Transfer,
DownloadStatus, File, Search, SearchResult, Transfer,
};
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 error::{Result, SoulseekRs};
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
.read()
.map_err(|_| DownloadError::LockPoisoned)?;
let download_info =
client_guard.download_tokens.get(&token_u32).cloned();
let download_info = client_guard.get_download_by_token(token_u32).cloned();
drop(client_guard);
download_info.ok_or(DownloadError::TokenNotFound(token_u32))
@@ -318,13 +317,13 @@ impl DownloadPeer {
0.0
};
let _ =
dl.sender.send(DownloadStatus::InProgress {
bytes_downloaded: processor.total_bytes
as u64,
total_bytes: dl.size,
speed_bytes_per_sec: speed,
});
let status = DownloadStatus::InProgress {
bytes_downloaded: processor.total_bytes as u64,
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();
}
@@ -422,6 +421,7 @@ impl DownloadPeer {
if let Some(ref dl) = download {
let _ = dl.sender.send(DownloadStatus::Queued);
client_context.write().unwrap().update_download_with_status(dl.token, DownloadStatus::Queued);
}
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 download = context.download_tokens.get(&token).cloned();
let download = context.get_download_by_token(token).cloned();
if download.is_none() {
let download_tokens: Vec<u32> =
context.download_tokens.keys().cloned().collect();
let download_tokens = context.get_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)) => {
let _ = download.sender.send(DownloadStatus::Completed);
context.client_context.write().unwrap().update_download_with_status(download.token, DownloadStatus::Completed);
info!(
"Successfully downloaded {} bytes to {}",
download.size, filename
+1 -9
View File
@@ -94,16 +94,8 @@ pub struct Download {
pub token: u32,
pub size: u64,
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 elapsed_time: std::time::Duration,
pub sender: Sender<DownloadStatus>,
}
#[derive(Debug, Clone)]