410 lines
15 KiB
Rust
410 lines
15 KiB
Rust
#![doc(hidden)]
|
|
|
|
use super::{
|
|
Arc, Cow, Downloader, HttpResponseModel, HttpTransferTaskModel, Mutex, RangeSpec, RangeUnit,
|
|
RetryPolicy, RuntimeConfig, VecDeque, planned_segment_span, saturating_u32_from_usize,
|
|
saturating_usize_from_u64,
|
|
};
|
|
|
|
const fn should_retry_response(response: &HttpResponseModel, policy: &RetryPolicy) -> bool {
|
|
match response.status {
|
|
300..=399 => policy.retry_on_3xx,
|
|
400..=499 => policy.retry_on_4xx,
|
|
500..=599 => policy.retry_on_5xx,
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
fn response_total_length(response: &HttpResponseModel) -> Option<u64> {
|
|
response.total_length()
|
|
}
|
|
|
|
fn response_completed_length(response: &HttpResponseModel) -> u64 {
|
|
response.completed_length()
|
|
}
|
|
|
|
pub(crate) struct HttpTransferExecution {
|
|
pub(crate) response: Option<HttpResponseModel>,
|
|
pub(crate) retry_count: u32,
|
|
pub(crate) retry_attempts: Vec<aria2_rust_pro_protocol::RetryAttempt>,
|
|
pub(crate) planned_ranges: Vec<Option<RangeSpec>>,
|
|
pub(crate) checksum_observed: bool,
|
|
pub(crate) checksum_complete: bool,
|
|
}
|
|
|
|
#[expect(
|
|
clippy::single_match_else,
|
|
clippy::too_many_lines,
|
|
reason = "retry, resume, checksum, and terminal-status branches are kept together to preserve transfer semantics"
|
|
)]
|
|
pub(crate) fn execute_http_transfer_with_retry<D: Downloader>(
|
|
downloader: &D,
|
|
base_task: &HttpTransferTaskModel,
|
|
runtime: &RuntimeConfig,
|
|
) -> HttpTransferExecution {
|
|
let max_attempts = base_task.retry.policy.max_attempts.max(1);
|
|
let requested_start = base_task
|
|
.request
|
|
.range
|
|
.as_ref()
|
|
.map_or(0, |range| range.start);
|
|
let requested_end_exclusive = base_task
|
|
.request
|
|
.range
|
|
.as_ref()
|
|
.and_then(|range| range.end_inclusive.map(|end| end.saturating_add(1)));
|
|
let mut completed_length = requested_start;
|
|
let mut retry_attempts = Vec::new();
|
|
let mut current_total_length = 0_u64;
|
|
let mut planned_ranges = Vec::new();
|
|
let mut checksum_observed = false;
|
|
let mut checksum_complete = false;
|
|
|
|
for attempt in 0..max_attempts {
|
|
let next_attempt = attempt.saturating_add(1);
|
|
let task = if completed_length > requested_start || !retry_attempts.is_empty() {
|
|
let mut owned_task = base_task.clone();
|
|
if completed_length > requested_start {
|
|
let end_inclusive = requested_end_exclusive
|
|
.map(|end| end.saturating_sub(1))
|
|
.or_else(|| {
|
|
planned_segment_span(current_total_length, runtime).and_then(|span| {
|
|
let next_end = completed_length.saturating_add(span).saturating_sub(1);
|
|
(current_total_length > 0)
|
|
.then_some(next_end.min(current_total_length.saturating_sub(1)))
|
|
})
|
|
});
|
|
owned_task.request.range = Some(RangeSpec {
|
|
start: completed_length,
|
|
end_inclusive,
|
|
unit: RangeUnit::Bytes,
|
|
});
|
|
owned_task.resume_state = Some(aria2_rust_pro_protocol::ResumeState {
|
|
requested_offset: completed_length,
|
|
accepted_offset: None,
|
|
resumed: true,
|
|
});
|
|
}
|
|
if !retry_attempts.is_empty() {
|
|
owned_task.retry_attempts.clone_from(&retry_attempts);
|
|
}
|
|
Cow::Owned(owned_task)
|
|
} else {
|
|
Cow::Borrowed(base_task)
|
|
};
|
|
planned_ranges.push(task.request.range);
|
|
|
|
match downloader.start_http_transfer(task.as_ref()) {
|
|
Ok(response) => {
|
|
let success = (200..=299).contains(&response.status);
|
|
let total_length = response_total_length(&response).unwrap_or(0);
|
|
if response.checksum.is_some() {
|
|
checksum_observed = true;
|
|
}
|
|
if total_length > 0 {
|
|
current_total_length = total_length;
|
|
}
|
|
let completed_after = response_completed_length(&response);
|
|
if success {
|
|
completed_length = completed_length.max(completed_after);
|
|
if response.checksum.is_some()
|
|
&& total_length > 0
|
|
&& completed_length >= total_length
|
|
&& response.completion_model().checksum_verified
|
|
{
|
|
checksum_complete = true;
|
|
}
|
|
}
|
|
let terminal_success = if let Some(segment_end) = requested_end_exclusive {
|
|
success && completed_length >= segment_end
|
|
} else {
|
|
success
|
|
&& (!response.partial_content
|
|
|| (total_length > 0 && completed_length >= total_length))
|
|
};
|
|
if terminal_success || next_attempt >= max_attempts {
|
|
return HttpTransferExecution {
|
|
response: Some(response),
|
|
retry_count: saturating_u32_from_usize(retry_attempts.len()),
|
|
retry_attempts,
|
|
planned_ranges,
|
|
checksum_observed,
|
|
checksum_complete,
|
|
};
|
|
}
|
|
if success
|
|
&& response.partial_content
|
|
&& requested_end_exclusive
|
|
.is_some_and(|segment_end| completed_length < segment_end)
|
|
{
|
|
retry_attempts.push(aria2_rust_pro_protocol::RetryAttempt {
|
|
attempt: next_attempt,
|
|
reason: aria2_rust_pro_protocol::RetryReason::Other,
|
|
status: Some(response.status),
|
|
backoff_ms: Some(0),
|
|
});
|
|
continue;
|
|
}
|
|
if !should_retry_response(&response, &task.retry.policy) {
|
|
return HttpTransferExecution {
|
|
response: Some(response),
|
|
retry_count: saturating_u32_from_usize(retry_attempts.len()),
|
|
retry_attempts,
|
|
planned_ranges,
|
|
checksum_observed,
|
|
checksum_complete,
|
|
};
|
|
}
|
|
let retry_reason = match response.status {
|
|
300..=399 => aria2_rust_pro_protocol::RetryReason::Http3xx,
|
|
400..=499 => aria2_rust_pro_protocol::RetryReason::Http4xx,
|
|
500..=599 => aria2_rust_pro_protocol::RetryReason::Http5xx,
|
|
_ => aria2_rust_pro_protocol::RetryReason::Other,
|
|
};
|
|
retry_attempts.push(aria2_rust_pro_protocol::RetryAttempt {
|
|
attempt: next_attempt,
|
|
reason: retry_reason,
|
|
status: Some(response.status),
|
|
backoff_ms: None,
|
|
});
|
|
}
|
|
Err(_) => {
|
|
if next_attempt >= max_attempts || !task.retry.policy.retry_on_network_error {
|
|
return HttpTransferExecution {
|
|
response: None,
|
|
retry_count: saturating_u32_from_usize(retry_attempts.len()),
|
|
retry_attempts,
|
|
planned_ranges,
|
|
checksum_observed,
|
|
checksum_complete,
|
|
};
|
|
}
|
|
retry_attempts.push(aria2_rust_pro_protocol::RetryAttempt {
|
|
attempt: next_attempt,
|
|
reason: aria2_rust_pro_protocol::RetryReason::NetworkError,
|
|
status: None,
|
|
backoff_ms: None,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
HttpTransferExecution {
|
|
response: None,
|
|
retry_count: saturating_u32_from_usize(retry_attempts.len()),
|
|
retry_attempts,
|
|
planned_ranges,
|
|
checksum_observed,
|
|
checksum_complete,
|
|
}
|
|
}
|
|
|
|
pub(crate) fn execute_tagged_segment_transfers<D: Downloader + Sync, T: Send>(
|
|
downloader: &D,
|
|
planned_tasks: Vec<(T, HttpTransferTaskModel)>,
|
|
runtime: &RuntimeConfig,
|
|
) -> Vec<(T, HttpTransferTaskModel, HttpTransferExecution)> {
|
|
let parallelism = effective_segment_transfer_parallelism(runtime, planned_tasks.len());
|
|
execute_tagged_segment_transfers_with_parallelism(
|
|
downloader,
|
|
planned_tasks,
|
|
runtime,
|
|
parallelism,
|
|
)
|
|
}
|
|
|
|
pub(crate) fn execute_tagged_segment_transfers_with_parallelism<D: Downloader + Sync, T: Send>(
|
|
downloader: &D,
|
|
planned_tasks: Vec<(T, HttpTransferTaskModel)>,
|
|
runtime: &RuntimeConfig,
|
|
parallelism: usize,
|
|
) -> Vec<(T, HttpTransferTaskModel, HttpTransferExecution)> {
|
|
if planned_tasks.len() <= 1 || parallelism <= 1 {
|
|
return planned_tasks
|
|
.into_iter()
|
|
.map(|(tag, task)| {
|
|
let execution = execute_http_transfer_with_retry(downloader, &task, runtime);
|
|
(tag, task, execution)
|
|
})
|
|
.collect();
|
|
}
|
|
|
|
if planned_tasks.len() <= parallelism.saturating_mul(4) {
|
|
return execute_tagged_segment_transfers_static_partitioned(
|
|
downloader,
|
|
planned_tasks,
|
|
runtime,
|
|
parallelism,
|
|
);
|
|
}
|
|
|
|
let task_count = planned_tasks.len();
|
|
let chunk_size = task_count
|
|
.div_ceil(parallelism.saturating_mul(4).max(1))
|
|
.max(1);
|
|
let mut chunk_queue = VecDeque::new();
|
|
let mut current_chunk = Vec::with_capacity(chunk_size);
|
|
for item in planned_tasks.into_iter().enumerate() {
|
|
current_chunk.push(item);
|
|
if current_chunk.len() >= chunk_size {
|
|
chunk_queue.push_back(std::mem::take(&mut current_chunk));
|
|
current_chunk = Vec::with_capacity(chunk_size);
|
|
}
|
|
}
|
|
if !current_chunk.is_empty() {
|
|
chunk_queue.push_back(current_chunk);
|
|
}
|
|
let work_chunks = Arc::new(Mutex::new(chunk_queue));
|
|
|
|
let mut indexed_results = std::thread::scope(|scope| {
|
|
let mut handles = Vec::with_capacity(parallelism);
|
|
for _ in 0..parallelism {
|
|
let work_chunks = Arc::clone(&work_chunks);
|
|
handles.push(scope.spawn(move || {
|
|
let mut local_results = Vec::new();
|
|
loop {
|
|
let Some(chunk) = work_chunks
|
|
.lock()
|
|
.expect("segment work chunk mutex should not be poisoned")
|
|
.pop_front()
|
|
else {
|
|
break;
|
|
};
|
|
local_results.reserve(chunk.len());
|
|
for (index, (tag, task)) in chunk {
|
|
let execution =
|
|
execute_http_transfer_with_retry(downloader, &task, runtime);
|
|
local_results.push((index, (tag, task, execution)));
|
|
}
|
|
}
|
|
local_results
|
|
}));
|
|
}
|
|
handles
|
|
.into_iter()
|
|
.flat_map(|handle| {
|
|
handle
|
|
.join()
|
|
.expect("dynamic segment transfer worker should not panic")
|
|
})
|
|
.collect::<Vec<_>>()
|
|
});
|
|
debug_assert_eq!(indexed_results.len(), task_count);
|
|
indexed_results.sort_by_key(|(index, _)| *index);
|
|
indexed_results
|
|
.into_iter()
|
|
.map(|(_, result)| result)
|
|
.collect()
|
|
}
|
|
|
|
fn execute_tagged_segment_transfers_static_partitioned<D: Downloader + Sync, T: Send>(
|
|
downloader: &D,
|
|
planned_tasks: Vec<(T, HttpTransferTaskModel)>,
|
|
runtime: &RuntimeConfig,
|
|
parallelism: usize,
|
|
) -> Vec<(T, HttpTransferTaskModel, HttpTransferExecution)> {
|
|
let task_count = planned_tasks.len();
|
|
let indexed_chunks = partition_indexed_work_evenly(
|
|
planned_tasks.into_iter().enumerate(),
|
|
task_count,
|
|
parallelism,
|
|
);
|
|
|
|
std::thread::scope(|scope| {
|
|
let mut handles = Vec::with_capacity(indexed_chunks.len());
|
|
for chunk in indexed_chunks {
|
|
handles.push(scope.spawn(move || {
|
|
chunk
|
|
.into_iter()
|
|
.map(|(index, (tag, task))| {
|
|
let execution =
|
|
execute_http_transfer_with_retry(downloader, &task, runtime);
|
|
(index, (tag, task, execution))
|
|
})
|
|
.collect::<Vec<_>>()
|
|
}));
|
|
}
|
|
|
|
let mut ordered_results = std::iter::repeat_with(|| None)
|
|
.take(task_count)
|
|
.collect::<Vec<_>>();
|
|
for handle in handles {
|
|
for (index, result) in handle
|
|
.join()
|
|
.expect("static segment transfer worker should not panic")
|
|
{
|
|
*ordered_results
|
|
.get_mut(index)
|
|
.expect("static segment transfer worker should return an in-bounds index") =
|
|
Some(result);
|
|
}
|
|
}
|
|
|
|
ordered_results
|
|
.into_iter()
|
|
.map(|result| result.expect("static segment transfer worker should fill every slot"))
|
|
.collect()
|
|
})
|
|
}
|
|
|
|
pub(crate) fn execute_segment_transfers<D: Downloader + Sync>(
|
|
downloader: &D,
|
|
planned_tasks: Vec<HttpTransferTaskModel>,
|
|
runtime: &RuntimeConfig,
|
|
) -> Vec<(HttpTransferTaskModel, HttpTransferExecution)> {
|
|
execute_tagged_segment_transfers(
|
|
downloader,
|
|
planned_tasks.into_iter().map(|task| ((), task)).collect(),
|
|
runtime,
|
|
)
|
|
.into_iter()
|
|
.map(|((), task, execution)| (task, execution))
|
|
.collect()
|
|
}
|
|
|
|
pub(crate) fn partition_indexed_work_evenly<T>(
|
|
indexed_work: impl IntoIterator<Item = (usize, T)>,
|
|
task_count: usize,
|
|
parallelism: usize,
|
|
) -> Vec<Vec<(usize, T)>> {
|
|
if task_count == 0 {
|
|
return Vec::new();
|
|
}
|
|
|
|
let worker_count = parallelism.max(1).min(task_count);
|
|
let mut chunks = std::iter::repeat_with(Vec::new)
|
|
.take(worker_count)
|
|
.collect::<Vec<_>>();
|
|
for (ordinal, item) in indexed_work.into_iter().enumerate() {
|
|
let chunk_index = ordinal
|
|
.checked_rem(worker_count)
|
|
.expect("worker count is nonzero when partitioning indexed work");
|
|
chunks
|
|
.get_mut(chunk_index)
|
|
.expect("round-robin chunk index should be in bounds")
|
|
.push(item);
|
|
}
|
|
chunks.retain(|chunk| !chunk.is_empty());
|
|
chunks
|
|
}
|
|
|
|
pub(crate) fn effective_segment_transfer_parallelism(
|
|
runtime: &RuntimeConfig,
|
|
planned_tasks: usize,
|
|
) -> usize {
|
|
if planned_tasks <= 1 {
|
|
return planned_tasks;
|
|
}
|
|
|
|
let segment_unit = runtime.min_split_size.max(runtime.piece_length).max(1);
|
|
let mut parallelism = planned_tasks;
|
|
if let Some(limit) = runtime
|
|
.max_download_limit
|
|
.map(|bytes_per_second| saturating_usize_from_u64(bytes_per_second.div_ceil(segment_unit)))
|
|
{
|
|
parallelism = parallelism.min(limit.max(1));
|
|
}
|
|
|
|
parallelism.max(1).min(planned_tasks)
|
|
}
|