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:
@@ -4,6 +4,7 @@ use std::fs::File;
|
||||
use std::io::{self, Read, Write};
|
||||
use std::net::{TcpStream, ToSocketAddrs};
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::message::server::MessageFactory;
|
||||
@@ -17,6 +18,9 @@ use crate::trace;
|
||||
/// download is matched by token), then the downloader sends an 8-byte
|
||||
/// 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
|
||||
/// Returns any I/O error opening the file or talking to the peer.
|
||||
pub fn serve_file(
|
||||
@@ -25,6 +29,8 @@ pub fn serve_file(
|
||||
own_username: &str,
|
||||
token: u32,
|
||||
path: &Path,
|
||||
bytes_sent: &AtomicU64,
|
||||
cancel: &AtomicBool,
|
||||
) -> io::Result<()> {
|
||||
let mut file = File::open(path)?;
|
||||
|
||||
@@ -53,7 +59,21 @@ pub fn serve_file(
|
||||
let mut offset = [0u8; 8];
|
||||
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()?;
|
||||
|
||||
// Linger so the downloader drains everything before the socket closes.
|
||||
@@ -67,6 +87,8 @@ mod tests {
|
||||
use super::serve_file;
|
||||
use std::io::{Read, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
|
||||
#[test]
|
||||
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 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 || {
|
||||
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
|
||||
@@ -104,6 +136,53 @@ mod tests {
|
||||
assert_eq!(received, content);
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user