feat: track uploads with progress and cancellation
serve_file streams in chunks, updating a shared byte counter and aborting with Interrupted when a cancel flag is set. Each served upload is registered (peer, virtual path, size, live progress, status) and exposed via Client::uploads(); Client::cancel_upload flags a matching in-progress transfer. e2e: the download test now also asserts the uploader recorded a completed upload with the right peer and byte count.
This commit is contained in:
@@ -59,6 +59,18 @@ fn next_upload_token() -> u32 {
|
|||||||
struct UploadJob {
|
struct UploadJob {
|
||||||
downloader: String,
|
downloader: String,
|
||||||
real_path: std::path::PathBuf,
|
real_path: std::path::PathBuf,
|
||||||
|
virtual_path: String,
|
||||||
|
size: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Live bookkeeping for an upload being served (or recently finished).
|
||||||
|
struct ActiveUpload {
|
||||||
|
username: String,
|
||||||
|
filename: String,
|
||||||
|
size: u64,
|
||||||
|
bytes_sent: Arc<std::sync::atomic::AtomicU64>,
|
||||||
|
cancel: Arc<std::sync::atomic::AtomicBool>,
|
||||||
|
status: crate::types::UploadStatus,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build a `FileSearchResponse` for `query` against `shares`, or `None` if
|
/// Build a `FileSearchResponse` for `query` against `shares`, or `None` if
|
||||||
@@ -210,6 +222,7 @@ pub struct ClientContext {
|
|||||||
pending_peer_messages: HashMap<String, Vec<crate::message::Message>>,
|
pending_peer_messages: HashMap<String, Vec<crate::message::Message>>,
|
||||||
/// Uploads we have offered, keyed by our transfer token.
|
/// Uploads we have offered, keyed by our transfer token.
|
||||||
uploads: HashMap<u32, UploadJob>,
|
uploads: HashMap<u32, UploadJob>,
|
||||||
|
active_uploads: HashMap<u32, ActiveUpload>,
|
||||||
/// Upload tokens waiting for the downloader's address to be resolved.
|
/// Upload tokens waiting for the downloader's address to be resolved.
|
||||||
pending_serves: HashMap<String, Vec<u32>>,
|
pending_serves: HashMap<String, Vec<u32>>,
|
||||||
/// Shared-file listings received from peers we browsed.
|
/// Shared-file listings received from peers we browsed.
|
||||||
@@ -507,6 +520,7 @@ impl ClientContext {
|
|||||||
peer_addresses: HashMap::new(),
|
peer_addresses: HashMap::new(),
|
||||||
pending_peer_messages: HashMap::new(),
|
pending_peer_messages: HashMap::new(),
|
||||||
uploads: HashMap::new(),
|
uploads: HashMap::new(),
|
||||||
|
active_uploads: HashMap::new(),
|
||||||
pending_serves: HashMap::new(),
|
pending_serves: HashMap::new(),
|
||||||
browse_results: HashMap::new(),
|
browse_results: HashMap::new(),
|
||||||
room_list: Vec::new(),
|
room_list: Vec::new(),
|
||||||
@@ -657,6 +671,54 @@ impl Client {
|
|||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Snapshot of the uploads served this session (active and finished),
|
||||||
|
/// most recent last.
|
||||||
|
#[must_use]
|
||||||
|
pub fn uploads(&self) -> Vec<crate::types::UploadInfo> {
|
||||||
|
self.context.read_safe().map_or_else(
|
||||||
|
|_| Vec::new(),
|
||||||
|
|ctx| {
|
||||||
|
let mut tokens: Vec<&u32> = ctx.active_uploads.keys().collect();
|
||||||
|
tokens.sort_unstable();
|
||||||
|
tokens
|
||||||
|
.into_iter()
|
||||||
|
.map(|token| {
|
||||||
|
let upload = &ctx.active_uploads[token];
|
||||||
|
crate::types::UploadInfo {
|
||||||
|
username: upload.username.clone(),
|
||||||
|
filename: upload.filename.clone(),
|
||||||
|
size: upload.size,
|
||||||
|
bytes_sent: upload
|
||||||
|
.bytes_sent
|
||||||
|
.load(std::sync::atomic::Ordering::Relaxed),
|
||||||
|
status: upload.status.clone(),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ask an in-progress upload to `username` of `filename` to stop.
|
||||||
|
/// Returns whether a matching in-progress upload was found.
|
||||||
|
pub fn cancel_upload(&self, username: &str, filename: &str) -> bool {
|
||||||
|
self.context.read_safe().is_ok_and(|ctx| {
|
||||||
|
let mut found = false;
|
||||||
|
for upload in ctx.active_uploads.values() {
|
||||||
|
if upload.username == username
|
||||||
|
&& upload.filename == filename
|
||||||
|
&& upload.status == crate::types::UploadStatus::InProgress
|
||||||
|
{
|
||||||
|
upload
|
||||||
|
.cancel
|
||||||
|
.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||||
|
found = true;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
found
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
/// Replace the shared directories at runtime: rescan into a fresh
|
/// Replace the shared directories at runtime: rescan into a fresh
|
||||||
/// index (served to peers from then on) and re-announce the new
|
/// index (served to peers from then on) and re-announce the new
|
||||||
/// folder/file counts to the server.
|
/// folder/file counts to the server.
|
||||||
|
|||||||
@@ -652,6 +652,8 @@ impl Client {
|
|||||||
UploadJob {
|
UploadJob {
|
||||||
downloader: downloader.clone(),
|
downloader: downloader.clone(),
|
||||||
real_path,
|
real_path,
|
||||||
|
virtual_path: filename.clone(),
|
||||||
|
size,
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
(ctx.peer_registry.clone(), size)
|
(ctx.peer_registry.clone(), size)
|
||||||
|
|||||||
@@ -1,7 +1,9 @@
|
|||||||
use super::{
|
use super::{
|
||||||
Arc, Client, ClientContext, DownloadStatus, RwLock, RwLockExt,
|
ActiveUpload, Arc, Client, ClientContext, DownloadStatus, RwLock,
|
||||||
collect_failed_tokens, error, thread,
|
RwLockExt, collect_failed_tokens, error, thread,
|
||||||
};
|
};
|
||||||
|
use crate::types::UploadStatus;
|
||||||
|
use std::sync::atomic::{AtomicBool, AtomicU64};
|
||||||
|
|
||||||
impl Client {
|
impl Client {
|
||||||
/// Consume the upload job for `token` and stream the file to `host:port`
|
/// Consume the upload job for `token` and stream the file to `host:port`
|
||||||
@@ -19,14 +21,47 @@ impl Client {
|
|||||||
let Some(job) = ctx.uploads.remove(&token) else {
|
let Some(job) = ctx.uploads.remove(&token) else {
|
||||||
return;
|
return;
|
||||||
};
|
};
|
||||||
|
let bytes_sent = Arc::new(AtomicU64::new(0));
|
||||||
|
let cancel = Arc::new(AtomicBool::new(false));
|
||||||
|
ctx.active_uploads.insert(
|
||||||
|
token,
|
||||||
|
ActiveUpload {
|
||||||
|
username: job.downloader.clone(),
|
||||||
|
filename: job.virtual_path.clone(),
|
||||||
|
size: job.size,
|
||||||
|
bytes_sent: bytes_sent.clone(),
|
||||||
|
cancel: cancel.clone(),
|
||||||
|
status: UploadStatus::InProgress,
|
||||||
|
},
|
||||||
|
);
|
||||||
drop(ctx);
|
drop(ctx);
|
||||||
let own = own_username.to_string();
|
let own = own_username.to_string();
|
||||||
let real_path = job.real_path;
|
let real_path = job.real_path;
|
||||||
|
let context = client_context.clone();
|
||||||
thread::spawn(move || {
|
thread::spawn(move || {
|
||||||
if let Err(e) = crate::peer::upload_peer::serve_file(
|
let result = crate::peer::upload_peer::serve_file(
|
||||||
&host, port, &own, token, &real_path,
|
&host,
|
||||||
) {
|
port,
|
||||||
error!("[client] serve {}: {}", real_path.display(), e);
|
&own,
|
||||||
|
token,
|
||||||
|
&real_path,
|
||||||
|
&bytes_sent,
|
||||||
|
&cancel,
|
||||||
|
);
|
||||||
|
let status = match &result {
|
||||||
|
Ok(()) => UploadStatus::Completed,
|
||||||
|
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
|
||||||
|
UploadStatus::Cancelled
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
error!("[client] serve {}: {}", real_path.display(), e);
|
||||||
|
UploadStatus::Failed(e.to_string())
|
||||||
|
}
|
||||||
|
};
|
||||||
|
if let Ok(mut ctx) = context.write_safe()
|
||||||
|
&& let Some(upload) = ctx.active_uploads.get_mut(&token)
|
||||||
|
{
|
||||||
|
upload.status = status;
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ use std::fs::File;
|
|||||||
use std::io::{self, Read, Write};
|
use std::io::{self, Read, Write};
|
||||||
use std::net::{TcpStream, ToSocketAddrs};
|
use std::net::{TcpStream, ToSocketAddrs};
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use crate::message::server::MessageFactory;
|
use crate::message::server::MessageFactory;
|
||||||
@@ -17,6 +18,9 @@ use crate::trace;
|
|||||||
/// download is matched by token), then the downloader sends an 8-byte
|
/// download is matched by token), then the downloader sends an 8-byte
|
||||||
/// START_DOWNLOAD offset before we stream the file.
|
/// START_DOWNLOAD offset before we stream the file.
|
||||||
///
|
///
|
||||||
|
/// `bytes_sent` is updated as the transfer progresses, and setting `cancel`
|
||||||
|
/// aborts the stream with an [`io::ErrorKind::Interrupted`] error.
|
||||||
|
///
|
||||||
/// # Errors
|
/// # Errors
|
||||||
/// Returns any I/O error opening the file or talking to the peer.
|
/// Returns any I/O error opening the file or talking to the peer.
|
||||||
pub fn serve_file(
|
pub fn serve_file(
|
||||||
@@ -25,6 +29,8 @@ pub fn serve_file(
|
|||||||
own_username: &str,
|
own_username: &str,
|
||||||
token: u32,
|
token: u32,
|
||||||
path: &Path,
|
path: &Path,
|
||||||
|
bytes_sent: &AtomicU64,
|
||||||
|
cancel: &AtomicBool,
|
||||||
) -> io::Result<()> {
|
) -> io::Result<()> {
|
||||||
let mut file = File::open(path)?;
|
let mut file = File::open(path)?;
|
||||||
|
|
||||||
@@ -53,7 +59,21 @@ pub fn serve_file(
|
|||||||
let mut offset = [0u8; 8];
|
let mut offset = [0u8; 8];
|
||||||
stream.read_exact(&mut offset)?;
|
stream.read_exact(&mut offset)?;
|
||||||
|
|
||||||
io::copy(&mut file, &mut stream)?;
|
let mut buffer = [0u8; 64 * 1024];
|
||||||
|
loop {
|
||||||
|
if cancel.load(Ordering::Relaxed) {
|
||||||
|
return Err(io::Error::new(
|
||||||
|
io::ErrorKind::Interrupted,
|
||||||
|
"upload cancelled",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
let read = file.read(&mut buffer)?;
|
||||||
|
if read == 0 {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
stream.write_all(&buffer[..read])?;
|
||||||
|
bytes_sent.fetch_add(read as u64, Ordering::Relaxed);
|
||||||
|
}
|
||||||
stream.flush()?;
|
stream.flush()?;
|
||||||
|
|
||||||
// Linger so the downloader drains everything before the socket closes.
|
// Linger so the downloader drains everything before the socket closes.
|
||||||
@@ -67,6 +87,8 @@ mod tests {
|
|||||||
use super::serve_file;
|
use super::serve_file;
|
||||||
use std::io::{Read, Write};
|
use std::io::{Read, Write};
|
||||||
use std::net::TcpListener;
|
use std::net::TcpListener;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn serve_file_streams_the_file_over_an_f_connection() {
|
fn serve_file_streams_the_file_over_an_f_connection() {
|
||||||
@@ -80,8 +102,18 @@ mod tests {
|
|||||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||||
let port = u32::from(listener.local_addr().unwrap().port());
|
let port = u32::from(listener.local_addr().unwrap().port());
|
||||||
|
|
||||||
|
let bytes_sent = Arc::new(AtomicU64::new(0));
|
||||||
|
let sent_counter = bytes_sent.clone();
|
||||||
let uploader = std::thread::spawn(move || {
|
let uploader = std::thread::spawn(move || {
|
||||||
serve_file("127.0.0.1", port, "me", 777, &path)
|
serve_file(
|
||||||
|
"127.0.0.1",
|
||||||
|
port,
|
||||||
|
"me",
|
||||||
|
777,
|
||||||
|
&path,
|
||||||
|
&sent_counter,
|
||||||
|
&AtomicBool::new(false),
|
||||||
|
)
|
||||||
});
|
});
|
||||||
|
|
||||||
// Downloader side: accept, read PeerInit(F) frame + raw token, send the
|
// Downloader side: accept, read PeerInit(F) frame + raw token, send the
|
||||||
@@ -104,6 +136,53 @@ mod tests {
|
|||||||
assert_eq!(received, content);
|
assert_eq!(received, content);
|
||||||
|
|
||||||
let _ = uploader.join();
|
let _ = uploader.join();
|
||||||
|
assert_eq!(bytes_sent.load(Ordering::Relaxed), 4096);
|
||||||
|
let _ = std::fs::remove_dir_all(dir);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn serve_file_stops_when_cancelled() {
|
||||||
|
let dir = std::env::temp_dir()
|
||||||
|
.join(format!("soulseek-upload-cancel-{}", std::process::id()));
|
||||||
|
std::fs::create_dir_all(&dir).unwrap();
|
||||||
|
let path = dir.join("big.bin");
|
||||||
|
std::fs::write(&path, vec![7u8; 1024 * 1024]).unwrap();
|
||||||
|
|
||||||
|
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||||
|
let port = u32::from(listener.local_addr().unwrap().port());
|
||||||
|
|
||||||
|
// Cancelled before the copy loop starts: the stream must abort with
|
||||||
|
// Interrupted instead of serving the whole file.
|
||||||
|
let cancel = Arc::new(AtomicBool::new(true));
|
||||||
|
let cancel_flag = cancel.clone();
|
||||||
|
let uploader = std::thread::spawn(move || {
|
||||||
|
serve_file(
|
||||||
|
"127.0.0.1",
|
||||||
|
port,
|
||||||
|
"me",
|
||||||
|
778,
|
||||||
|
&path,
|
||||||
|
&AtomicU64::new(0),
|
||||||
|
&cancel_flag,
|
||||||
|
)
|
||||||
|
});
|
||||||
|
|
||||||
|
let (mut stream, _) = listener.accept().unwrap();
|
||||||
|
let mut len_buf = [0u8; 4];
|
||||||
|
stream.read_exact(&mut len_buf).unwrap();
|
||||||
|
let mut payload = vec![0u8; u32::from_le_bytes(len_buf) as usize];
|
||||||
|
stream.read_exact(&mut payload).unwrap();
|
||||||
|
let mut token = [0u8; 4];
|
||||||
|
stream.read_exact(&mut token).unwrap();
|
||||||
|
stream.write_all(&[0u8; 8]).unwrap();
|
||||||
|
stream.flush().unwrap();
|
||||||
|
|
||||||
|
let mut received = Vec::new();
|
||||||
|
let _ = stream.read_to_end(&mut received);
|
||||||
|
assert!(received.is_empty(), "no file bytes after cancellation");
|
||||||
|
|
||||||
|
let err = uploader.join().unwrap().unwrap_err();
|
||||||
|
assert_eq!(err.kind(), std::io::ErrorKind::Interrupted);
|
||||||
let _ = std::fs::remove_dir_all(dir);
|
let _ = std::fs::remove_dir_all(dir);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -251,3 +251,24 @@ mod tests {
|
|||||||
assert_eq!(transfer.size, 0);
|
assert_eq!(transfer.size, 0);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Lifecycle of a file we are serving to a peer.
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
|
pub enum UploadStatus {
|
||||||
|
InProgress,
|
||||||
|
Completed,
|
||||||
|
Cancelled,
|
||||||
|
Failed(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A snapshot of one upload for display purposes.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct UploadInfo {
|
||||||
|
/// The peer receiving the file.
|
||||||
|
pub username: String,
|
||||||
|
/// The peer-facing virtual path being served.
|
||||||
|
pub filename: String,
|
||||||
|
pub size: u64,
|
||||||
|
pub bytes_sent: u64,
|
||||||
|
pub status: UploadStatus,
|
||||||
|
}
|
||||||
|
|||||||
@@ -1245,6 +1245,24 @@ fn two_real_clients_search_and_download() {
|
|||||||
.expect("downloaded file should exist");
|
.expect("downloaded file should exist");
|
||||||
assert_eq!(written, content, "downloaded bytes should match the source");
|
assert_eq!(written, content, "downloaded bytes should match the source");
|
||||||
|
|
||||||
|
// The uploader side tracked the transfer and saw it complete.
|
||||||
|
let deadline = Instant::now() + Duration::from_secs(5);
|
||||||
|
let mut uploads = sharer.uploads();
|
||||||
|
while Instant::now() < deadline
|
||||||
|
&& !uploads
|
||||||
|
.iter()
|
||||||
|
.any(|u| u.status == soulseek_rs::types::UploadStatus::Completed)
|
||||||
|
{
|
||||||
|
std::thread::sleep(Duration::from_millis(200));
|
||||||
|
uploads = sharer.uploads();
|
||||||
|
}
|
||||||
|
let upload = uploads
|
||||||
|
.iter()
|
||||||
|
.find(|u| u.status == soulseek_rs::types::UploadStatus::Completed)
|
||||||
|
.expect("uploader should record a completed upload");
|
||||||
|
assert_eq!(upload.username, "e2e_leecher");
|
||||||
|
assert_eq!(upload.bytes_sent, size);
|
||||||
|
|
||||||
let _ = std::fs::remove_dir_all(share_dir);
|
let _ = std::fs::remove_dir_all(share_dir);
|
||||||
let _ = std::fs::remove_dir_all(download_dir);
|
let _ = std::fs::remove_dir_all(download_dir);
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user