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 {
|
||||
downloader: String,
|
||||
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
|
||||
@@ -210,6 +222,7 @@ pub struct ClientContext {
|
||||
pending_peer_messages: HashMap<String, Vec<crate::message::Message>>,
|
||||
/// Uploads we have offered, keyed by our transfer token.
|
||||
uploads: HashMap<u32, UploadJob>,
|
||||
active_uploads: HashMap<u32, ActiveUpload>,
|
||||
/// Upload tokens waiting for the downloader's address to be resolved.
|
||||
pending_serves: HashMap<String, Vec<u32>>,
|
||||
/// Shared-file listings received from peers we browsed.
|
||||
@@ -507,6 +520,7 @@ impl ClientContext {
|
||||
peer_addresses: HashMap::new(),
|
||||
pending_peer_messages: HashMap::new(),
|
||||
uploads: HashMap::new(),
|
||||
active_uploads: HashMap::new(),
|
||||
pending_serves: HashMap::new(),
|
||||
browse_results: HashMap::new(),
|
||||
room_list: Vec::new(),
|
||||
@@ -657,6 +671,54 @@ impl Client {
|
||||
.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
|
||||
/// index (served to peers from then on) and re-announce the new
|
||||
/// folder/file counts to the server.
|
||||
|
||||
@@ -652,6 +652,8 @@ impl Client {
|
||||
UploadJob {
|
||||
downloader: downloader.clone(),
|
||||
real_path,
|
||||
virtual_path: filename.clone(),
|
||||
size,
|
||||
},
|
||||
);
|
||||
(ctx.peer_registry.clone(), size)
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use super::{
|
||||
Arc, Client, ClientContext, DownloadStatus, RwLock, RwLockExt,
|
||||
collect_failed_tokens, error, thread,
|
||||
ActiveUpload, Arc, Client, ClientContext, DownloadStatus, RwLock,
|
||||
RwLockExt, collect_failed_tokens, error, thread,
|
||||
};
|
||||
use crate::types::UploadStatus;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64};
|
||||
|
||||
impl Client {
|
||||
/// 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 {
|
||||
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);
|
||||
let own = own_username.to_string();
|
||||
let real_path = job.real_path;
|
||||
let context = client_context.clone();
|
||||
thread::spawn(move || {
|
||||
if let Err(e) = crate::peer::upload_peer::serve_file(
|
||||
&host, port, &own, token, &real_path,
|
||||
) {
|
||||
error!("[client] serve {}: {}", real_path.display(), e);
|
||||
let result = crate::peer::upload_peer::serve_file(
|
||||
&host,
|
||||
port,
|
||||
&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;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user