use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use reqwest::Client;
use tokio::sync::mpsc;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use crate::engine::event::EngineEvent;
use crate::engine::models::{TestPhase, TransferStats};
use crate::engine::TransferConfig;
use crate::errors::EngineResult;
use super::transfer;
const RETRY_DELAY: Duration = Duration::from_millis(250);
pub async fn run_download(
client: &Client,
url: &str,
config: &TransferConfig,
events: &mpsc::Sender<EngineEvent>,
cancel: &CancellationToken,
) -> EngineResult<TransferStats> {
let counter = Arc::new(AtomicU64::new(0));
let workers_cancel = cancel.child_token();
let mut workers = JoinSet::new();
for _ in 0..config.connections.max(1) {
workers.spawn(worker(
client.clone(),
url.to_owned(),
Arc::clone(&counter),
workers_cancel.clone(),
));
}
transfer::monitor(
TestPhase::Download,
config,
counter,
&mut workers,
&workers_cancel,
events,
cancel,
)
.await
}
async fn worker(client: Client, url: String, counter: Arc<AtomicU64>, cancel: CancellationToken) {
while !cancel.is_cancelled() {
let response = tokio::select! {
() = cancel.cancelled() => return,
result = client.get(&url).send() => result,
};
let mut response = match response {
Ok(response) if response.status().is_success() => response,
Ok(_) | Err(_) => {
tokio::select! {
() = cancel.cancelled() => return,
() = tokio::time::sleep(RETRY_DELAY) => continue,
}
}
};
loop {
let chunk = tokio::select! {
() = cancel.cancelled() => return,
result = response.chunk() => result,
};
match chunk {
Ok(Some(bytes)) => {
counter.fetch_add(bytes.len() as u64, Ordering::Relaxed);
}
Ok(None) | Err(_) => break,
}
}
}
}