use std::collections::HashMap;
use std::thread;
use std::time::Duration;
use crate::api::ApiClient;
use crate::error::BioLibError;
const MIN_CHUNK_SIZE: usize = 10_000_000; const MAX_CHUNK_COUNT: usize = 9_000;
const MAX_CHUNK_RETRIES: u32 = 20;
pub fn upload_module_input(
api_client: &mut ApiClient,
job_uuid: &str,
job_auth_token: &str,
module_input: &[u8],
) -> crate::Result<()> {
let mut headers = HashMap::new();
headers.insert("job-auth-token".to_string(), job_auth_token.to_string());
api_client.start_multipart_upload(
&format!("/jobs/{job_uuid}/storage/input/start_upload/"),
Some(&headers),
)?;
let chunk_size = calculate_chunk_size(module_input.len());
let mut parts = Vec::new();
let mut offset = 0usize;
let mut part_number = 1u32;
while offset < module_input.len() {
let end = std::cmp::min(offset + chunk_size, module_input.len());
let chunk = &module_input[offset..end];
let (etag, _chunk_len) = upload_chunk(
api_client,
&format!("/jobs/{job_uuid}/storage/input/presigned_upload_url/"),
part_number,
chunk,
&headers,
)?;
parts.push(serde_json::json!({
"ETag": etag,
"PartNumber": part_number,
}));
offset = end;
part_number += 1;
}
api_client.complete_multipart_upload(
&format!("/jobs/{job_uuid}/storage/input/complete_upload/"),
&parts,
Some(&headers),
)?;
Ok(())
}
fn upload_chunk(
api_client: &mut ApiClient,
presigned_url_path: &str,
part_number: u32,
chunk: &[u8],
headers: &HashMap<String, String>,
) -> crate::Result<(String, usize)> {
let mut last_error: Option<BioLibError> = None;
for attempt in 0..MAX_CHUNK_RETRIES {
if attempt > 0 {
let delay_secs = (attempt as u64) * (attempt as u64) + 2;
crate::logging::warning(&format!(
"Retrying upload of part {part_number} (attempt {attempt}, waiting {delay_secs}s)..."
));
thread::sleep(Duration::from_secs(delay_secs));
}
let presigned_url = match api_client.get_presigned_upload_url(
presigned_url_path,
part_number,
Some(headers),
) {
Ok(url) => url,
Err(err) => {
crate::logging::warning(&format!(
"Error getting URL for part {part_number}: {err}"
));
last_error = Some(err);
continue;
}
};
match api_client.upload_to_presigned_url(&presigned_url, chunk) {
Ok(etag) => return Ok((etag, chunk.len())),
Err(err) => {
crate::logging::warning(&format!("Error uploading part {part_number}: {err}"));
last_error = Some(err);
}
}
}
Err(last_error.unwrap_or_else(|| {
BioLibError::General(format!("Max retries hit when uploading part {part_number}"))
}))
}
fn calculate_chunk_size(total_size: usize) -> usize {
std::cmp::max(MIN_CHUNK_SIZE, total_size / MAX_CHUNK_COUNT)
}