chore: initial sanitized public snapshot
This commit is contained in:
Aria2 Rust Pro Contributors
@@ -0,0 +1,409 @@
|
||||
#![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)
|
||||
}
|
||||
Reference in New Issue
Block a user