mirror of
https://github.com/Yuyi-Oak/BlueArchiveToolkit.git
synced 2026-09-18 06:34:54 +08:00
270 lines
9.3 KiB
Rust
270 lines
9.3 KiB
Rust
//! Downloader backend and bounded scheduling contracts.
|
|
//!
|
|
//! The scheduler is deliberately independent from curl, manifests, and
|
|
//! official URL rules. Those concerns belong to a backend and the caller,
|
|
//! which keeps retry, proxy, and verification policy composable.
|
|
|
|
use std::sync::{mpsc, Arc, Mutex};
|
|
|
|
/// Lowest supported download concurrency.
|
|
pub const MIN_DOWNLOAD_CONCURRENCY: usize = 1;
|
|
/// Highest supported download concurrency.
|
|
pub const MAX_DOWNLOAD_CONCURRENCY: usize = 256;
|
|
/// Default official download concurrency.
|
|
pub const DEFAULT_DOWNLOAD_CONCURRENCY: usize = 8;
|
|
|
|
/// A backend that executes one already-planned download task.
|
|
pub trait DownloaderBackend<T>: Send + Sync {
|
|
/// Successful result returned for one task.
|
|
type Output: Send;
|
|
/// Failure returned for one task.
|
|
type Error: Send;
|
|
|
|
/// Executes one task. The scheduler owns ordering and concurrency only.
|
|
fn download(&self, task: T) -> Result<Self::Output, Self::Error>;
|
|
}
|
|
|
|
/// Results returned by a scheduler for a task type and backend.
|
|
pub type DownloadResults<T, B> =
|
|
Vec<Result<<B as DownloaderBackend<T>>::Output, <B as DownloaderBackend<T>>::Error>>;
|
|
|
|
/// A bounded worker scheduler.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct DownloadScheduler {
|
|
max_concurrency: usize,
|
|
}
|
|
|
|
impl DownloadScheduler {
|
|
/// Creates a scheduler with the supported bounded range.
|
|
///
|
|
/// The official CLI validates input and reports out-of-range values.
|
|
/// This lower-level constructor remains total for library callers and
|
|
/// clamps values to the same safety bounds.
|
|
pub fn new(max_concurrency: usize) -> Self {
|
|
Self {
|
|
max_concurrency: max_concurrency
|
|
.clamp(MIN_DOWNLOAD_CONCURRENCY, MAX_DOWNLOAD_CONCURRENCY),
|
|
}
|
|
}
|
|
|
|
/// Returns the configured upper bound.
|
|
pub fn max_concurrency(self) -> usize {
|
|
self.max_concurrency
|
|
}
|
|
|
|
/// Executes tasks with a bounded number of workers.
|
|
///
|
|
/// Results are returned in input order even when workers finish out of
|
|
/// order. A failed task does not cause additional tasks to be scheduled
|
|
/// after it, because already-started bounded work must be joined cleanly;
|
|
/// callers decide whether a failed result invalidates the whole release.
|
|
pub fn execute<T, B>(self, backend: &B, tasks: Vec<T>) -> DownloadResults<T, B>
|
|
where
|
|
T: Send + 'static,
|
|
B: DownloaderBackend<T>,
|
|
{
|
|
self.execute_with_observer(
|
|
backend,
|
|
tasks,
|
|
|_, _| Ok::<(), std::convert::Infallible>(()),
|
|
)
|
|
.expect("infallible download observer cannot fail")
|
|
}
|
|
|
|
/// Executes tasks and observes each result as soon as a worker returns it.
|
|
///
|
|
/// The observer runs on the coordinator thread, while worker threads
|
|
/// immediately take another pending task after sending their result. An
|
|
/// observer error stops further observation but still drains and joins all
|
|
/// workers before returning, so no background transfer is left detached.
|
|
pub fn execute_with_observer<T, B, F, E>(
|
|
self,
|
|
backend: &B,
|
|
tasks: Vec<T>,
|
|
mut observer: F,
|
|
) -> Result<DownloadResults<T, B>, E>
|
|
where
|
|
T: Send + 'static,
|
|
B: DownloaderBackend<T>,
|
|
F: FnMut(usize, &Result<B::Output, B::Error>) -> Result<(), E>,
|
|
{
|
|
if tasks.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
if self.max_concurrency == 1 {
|
|
let mut results = Vec::with_capacity(tasks.len());
|
|
for (index, task) in tasks.into_iter().enumerate() {
|
|
let result = backend.download(task);
|
|
observer(index, &result)?;
|
|
results.push(result);
|
|
}
|
|
return Ok(results);
|
|
}
|
|
|
|
let total = tasks.len();
|
|
let worker_count = self.max_concurrency.min(total);
|
|
let pending = Arc::new(Mutex::new(tasks.into_iter().enumerate()));
|
|
let (result_sender, result_receiver) = mpsc::channel();
|
|
|
|
std::thread::scope(|scope| {
|
|
for _ in 0..worker_count {
|
|
let pending = Arc::clone(&pending);
|
|
let result_sender = result_sender.clone();
|
|
scope.spawn(move || loop {
|
|
let task = pending
|
|
.lock()
|
|
.expect("download scheduler task queue poisoned")
|
|
.next();
|
|
let Some((index, task)) = task else {
|
|
break;
|
|
};
|
|
let result = backend.download(task);
|
|
if result_sender.send((index, result)).is_err() {
|
|
break;
|
|
}
|
|
});
|
|
}
|
|
drop(result_sender);
|
|
|
|
let mut results = std::iter::repeat_with(|| None)
|
|
.take(total)
|
|
.collect::<Vec<_>>();
|
|
let mut observer_error = None;
|
|
for (index, result) in result_receiver {
|
|
if observer_error.is_none() {
|
|
if let Err(error) = observer(index, &result) {
|
|
observer_error = Some(error);
|
|
}
|
|
}
|
|
results[index] = Some(result);
|
|
}
|
|
let results = results
|
|
.into_iter()
|
|
.map(|result| result.expect("download scheduler lost a task result"))
|
|
.collect();
|
|
match observer_error {
|
|
Some(error) => Err(error),
|
|
None => Ok(results),
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::thread;
|
|
use std::time::Duration;
|
|
|
|
struct TestBackend {
|
|
active: AtomicUsize,
|
|
max_active: AtomicUsize,
|
|
}
|
|
|
|
impl DownloaderBackend<usize> for TestBackend {
|
|
type Output = usize;
|
|
type Error = String;
|
|
|
|
fn download(&self, task: usize) -> Result<Self::Output, Self::Error> {
|
|
let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
|
|
self.max_active.fetch_max(active, Ordering::SeqCst);
|
|
thread::sleep(Duration::from_millis(2));
|
|
self.active.fetch_sub(1, Ordering::SeqCst);
|
|
Ok(task * 2)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn scheduler_preserves_result_order_and_respects_bound() {
|
|
let backend = TestBackend {
|
|
active: AtomicUsize::new(0),
|
|
max_active: AtomicUsize::new(0),
|
|
};
|
|
let results = DownloadScheduler::new(2).execute(&backend, (0..8).collect());
|
|
|
|
assert_eq!(
|
|
results.into_iter().map(Result::unwrap).collect::<Vec<_>>(),
|
|
(0..8).map(|value| value * 2).collect::<Vec<_>>()
|
|
);
|
|
assert!(backend.max_active.load(Ordering::SeqCst) <= 2);
|
|
assert!(backend.max_active.load(Ordering::SeqCst) >= 2);
|
|
}
|
|
|
|
#[test]
|
|
fn zero_concurrency_is_conservative() {
|
|
assert_eq!(
|
|
DownloadScheduler::new(0).max_concurrency(),
|
|
MIN_DOWNLOAD_CONCURRENCY
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn scheduler_caps_untrusted_upper_bound() {
|
|
assert_eq!(
|
|
DownloadScheduler::new(usize::MAX).max_concurrency(),
|
|
MAX_DOWNLOAD_CONCURRENCY
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn observer_receives_completion_without_a_global_barrier() {
|
|
struct UnevenBackend {
|
|
active: AtomicUsize,
|
|
task_two_started_while_task_zero_active: AtomicUsize,
|
|
}
|
|
|
|
impl DownloaderBackend<usize> for UnevenBackend {
|
|
type Output = usize;
|
|
type Error = String;
|
|
|
|
fn download(&self, task: usize) -> Result<Self::Output, Self::Error> {
|
|
if task == 0 {
|
|
self.active.fetch_add(1, Ordering::SeqCst);
|
|
thread::sleep(Duration::from_millis(20));
|
|
self.active.fetch_sub(1, Ordering::SeqCst);
|
|
} else {
|
|
if task == 1 {
|
|
while self.active.load(Ordering::SeqCst) == 0 {
|
|
thread::yield_now();
|
|
}
|
|
}
|
|
if task == 2 && self.active.load(Ordering::SeqCst) > 0 {
|
|
self.task_two_started_while_task_zero_active
|
|
.fetch_add(1, Ordering::SeqCst);
|
|
}
|
|
thread::sleep(Duration::from_millis(if task == 1 { 1 } else { 5 }));
|
|
}
|
|
Ok(task)
|
|
}
|
|
}
|
|
|
|
let backend = UnevenBackend {
|
|
active: AtomicUsize::new(0),
|
|
task_two_started_while_task_zero_active: AtomicUsize::new(0),
|
|
};
|
|
let mut completed = Vec::new();
|
|
let results = DownloadScheduler::new(2)
|
|
.execute_with_observer(&backend, vec![0, 1, 2], |index, _| {
|
|
completed.push(index);
|
|
Ok::<(), ()>(())
|
|
})
|
|
.unwrap();
|
|
|
|
assert_eq!(
|
|
results.into_iter().map(Result::unwrap).collect::<Vec<_>>(),
|
|
vec![0, 1, 2]
|
|
);
|
|
assert_eq!(completed.len(), 3);
|
|
assert!(completed[0] == 1, "短任务应在长任务之前回传：{completed:?}");
|
|
assert_eq!(
|
|
backend
|
|
.task_two_started_while_task_zero_active
|
|
.load(Ordering::SeqCst),
|
|
1,
|
|
"worker 完成 task 1 后应立即领取 task 2"
|
|
);
|
|
}
|
|
}
|