fix: made uploads and downloads more consistent
This commit is contained in:
@@ -380,16 +380,16 @@ where
|
||||
};
|
||||
|
||||
if let Some(handler) = lock.get(&id) {
|
||||
let _ = handler.data_sender.send((number, bytes));
|
||||
let _ = handler.msg_sender.send(upload_file::UploadFileMsg::Data((number, bytes)));
|
||||
}
|
||||
}
|
||||
Command::GetUploadStatus(id, up_to) => {
|
||||
Command::GetUploadStatus(id, bid, up_to) => {
|
||||
let Ok(lock) = uploaded_files.lock() else {
|
||||
continue;
|
||||
};
|
||||
|
||||
if let Some(handler) = lock.get(&id) {
|
||||
let _ = handler.request_status.send(up_to);
|
||||
let _ = handler.msg_sender.send(upload_file::UploadFileMsg::ReqStatus(bid, up_to));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -414,13 +414,13 @@ where
|
||||
|
||||
lock.insert(handler.id, handler);
|
||||
}
|
||||
Command::DownloadFileStatus(id, needed) => {
|
||||
Command::DownloadFileStatus(id, bid, needed) => {
|
||||
let Ok(lock) = download_files.lock() else {
|
||||
continue;
|
||||
};
|
||||
|
||||
if let Some(handler) = lock.get(&id) {
|
||||
let _ = handler.download_status.send(needed);
|
||||
let _ = handler.download_status.send((bid, needed));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ use std::{
|
||||
mpsc::{channel, Sender},
|
||||
Arc, Mutex,
|
||||
},
|
||||
thread::Scope,
|
||||
thread::Scope, time::Duration,
|
||||
};
|
||||
|
||||
use sparse_05_common::messages::{Response, FILE_BUFFER_BUFFER_SIZE, FILE_TRANSFER_PACKET_SIZE};
|
||||
@@ -19,7 +19,7 @@ static CURRENT_FILE_DOWNLOAD_ID: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
pub(super) struct DownloadFileHandler {
|
||||
pub id: u64,
|
||||
pub download_status: Sender<Vec<u64>>,
|
||||
pub download_status: Sender<(u64, Vec<u64>)>,
|
||||
}
|
||||
|
||||
pub(super) fn start_file_download<'a, 'b: 'a>(
|
||||
@@ -41,6 +41,8 @@ pub(super) fn start_file_download<'a, 'b: 'a>(
|
||||
conninfo
|
||||
.send(conninfo.encrypt_and_sign_resp(Response::StartDownloadFile(id, packet_count))?)?;
|
||||
|
||||
let mut current_buffer = 0;
|
||||
|
||||
loop {
|
||||
let mut file_data: Vec<Vec<u8>> = Vec::with_capacity(FILE_BUFFER_BUFFER_SIZE);
|
||||
let mut buffer = [0u8; FILE_TRANSFER_PACKET_SIZE];
|
||||
@@ -68,13 +70,24 @@ pub(super) fn start_file_download<'a, 'b: 'a>(
|
||||
}
|
||||
|
||||
loop {
|
||||
std::thread::sleep(std::time::Duration::from_millis(1000));
|
||||
|
||||
conninfo.send(conninfo.encrypt_and_sign_resp(
|
||||
Response::GetDownloadFileStatus(id, file_data.len() as u64),
|
||||
Response::GetDownloadFileStatus(id, current_buffer, file_data.len() as u64),
|
||||
)?)?;
|
||||
|
||||
let buffers_needed: Vec<u64> = receive_download_status.recv()?;
|
||||
let buffers_needed: Vec<u64> = loop {
|
||||
let buffers = receive_download_status.recv_timeout(Duration::from_millis(125));
|
||||
|
||||
match buffers {
|
||||
Ok((bid, buf)) if bid == current_buffer => {
|
||||
break buf;
|
||||
}
|
||||
_ => {
|
||||
conninfo.send(conninfo.encrypt_and_sign_resp(
|
||||
Response::GetDownloadFileStatus(id, current_buffer, file_data.len() as u64),
|
||||
)?)?;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if buffers_needed.is_empty() {
|
||||
break;
|
||||
@@ -90,6 +103,8 @@ pub(super) fn start_file_download<'a, 'b: 'a>(
|
||||
}
|
||||
}
|
||||
|
||||
current_buffer += 1;
|
||||
|
||||
if done {
|
||||
break;
|
||||
}
|
||||
|
||||
@@ -18,10 +18,14 @@ use super::ConnectionInformation;
|
||||
|
||||
static CURRENT_FILE_UPLOAD_ID: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
pub(super) enum UploadFileMsg {
|
||||
Data((u64, Vec<u8>)),
|
||||
ReqStatus(u64, u64)
|
||||
}
|
||||
|
||||
pub(super) struct UploadFileHandler {
|
||||
pub id: u64,
|
||||
pub data_sender: Sender<(u64, Vec<u8>)>,
|
||||
pub request_status: Sender<u64>,
|
||||
pub msg_sender: Sender<UploadFileMsg>,
|
||||
}
|
||||
|
||||
pub(super) fn start_file_upload<'a, 'b: 'a>(
|
||||
@@ -31,8 +35,7 @@ pub(super) fn start_file_upload<'a, 'b: 'a>(
|
||||
conninfo: ConnectionInformation,
|
||||
upload_file_map: Arc<Mutex<HashMap<u64, UploadFileHandler>>>,
|
||||
) -> anyhow::Result<UploadFileHandler> {
|
||||
let (data_sender, data_receiver) = channel();
|
||||
let (request_status, receive_request_status) = channel();
|
||||
let (msg_sender, msg_receiver) = channel();
|
||||
|
||||
let id = CURRENT_FILE_UPLOAD_ID.fetch_add(1, Ordering::Relaxed);
|
||||
let id_2 = id;
|
||||
@@ -46,37 +49,47 @@ pub(super) fn start_file_upload<'a, 'b: 'a>(
|
||||
.truncate(true)
|
||||
.open(&file_path)?;
|
||||
let mut current_packet_count = 0;
|
||||
let mut current_buffer = 0;
|
||||
|
||||
while current_packet_count < packet_count {
|
||||
let mut buffers: Vec<Option<Vec<u8>>> = vec![None; FILE_BUFFER_BUFFER_SIZE];
|
||||
|
||||
loop {
|
||||
let Ok((i, buffer)) = data_receiver.recv_timeout(Duration::from_millis(250)) else {
|
||||
let up_to = receive_request_status.recv()?;
|
||||
|
||||
let needed = buffers[..up_to as usize]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.flat_map(|(i, b)| match b {
|
||||
Some(..) => None,
|
||||
None => Some(i as u64),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let is_empty = needed.is_empty();
|
||||
|
||||
conninfo.send(
|
||||
conninfo.encrypt_and_sign_resp(Response::UploadFileStatus(id, needed))?,
|
||||
)?;
|
||||
|
||||
if is_empty {
|
||||
current_packet_count += up_to;
|
||||
break;
|
||||
} else {
|
||||
continue;
|
||||
match msg_receiver.recv() {
|
||||
Ok(UploadFileMsg::Data((i, buffer))) => {
|
||||
buffers[i as usize] = Some(buffer);
|
||||
}
|
||||
};
|
||||
buffers[i as usize] = Some(buffer);
|
||||
Ok(UploadFileMsg::ReqStatus(bid, up_to)) if bid == current_buffer => {
|
||||
let needed = buffers[..up_to as usize]
|
||||
.iter()
|
||||
.enumerate()
|
||||
.flat_map(|(i, b)| match b {
|
||||
Some(..) => None,
|
||||
None => Some(i as u64),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let is_empty = needed.is_empty();
|
||||
|
||||
println!("Status requested {current_buffer}; packets needed: {:?}", needed);
|
||||
|
||||
conninfo.send(
|
||||
conninfo.encrypt_and_sign_resp(Response::UploadFileStatus(id, current_buffer, needed.clone()))?,
|
||||
)?;
|
||||
|
||||
if is_empty {
|
||||
current_packet_count += up_to;
|
||||
current_buffer += 1;
|
||||
break;
|
||||
} else {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
Ok(..) => {}
|
||||
Err(e) => {
|
||||
eprintln!("Error receiving packet: {e:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for buffer in buffers {
|
||||
@@ -95,7 +108,6 @@ pub(super) fn start_file_upload<'a, 'b: 'a>(
|
||||
|
||||
Ok(UploadFileHandler {
|
||||
id: id_2,
|
||||
data_sender,
|
||||
request_status,
|
||||
msg_sender
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user