chore: refactor downloads
This commit is contained in:
@@ -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);
|
||||||
|
|||||||
@@ -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,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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)]
|
||||||
|
|||||||
Reference in New Issue
Block a user