use std::sync::Arc;
use dig_dht::ContentId;
use tokio::sync::{mpsc, oneshot, Mutex};
use crate::error::DownloadError;
use crate::orchestrator::{DownloadOptions, Downloader};
use crate::progress::DownloadEvent;
use crate::sink::Sink;
pub const DEFAULT_MAX_ACTIVE_DOWNLOADS: usize = 3;
struct QueuedJob {
content: ContentId,
sink: Arc<dyn Sink>,
opts: DownloadOptions,
events: mpsc::Sender<DownloadEvent>,
result: oneshot::Sender<Result<u64, DownloadError>>,
}
pub struct DownloadQueue {
submit_tx: mpsc::UnboundedSender<QueuedJob>,
max_active: usize,
}
impl DownloadQueue {
pub fn new(downloader: Arc<Downloader>, max_active: usize) -> Arc<Self> {
let max_active = max_active.max(1);
let (submit_tx, submit_rx) = mpsc::unbounded_channel::<QueuedJob>();
let submit_rx = Arc::new(Mutex::new(submit_rx));
for _ in 0..max_active {
let downloader = downloader.clone();
let submit_rx = submit_rx.clone();
tokio::spawn(async move {
loop {
let job = {
let mut rx = submit_rx.lock().await;
rx.recv().await
};
let Some(job) = job else {
return; };
run_job(&downloader, job).await;
}
});
}
Arc::new(DownloadQueue {
submit_tx,
max_active,
})
}
pub fn with_defaults(downloader: Arc<Downloader>) -> Arc<Self> {
Self::new(downloader, DEFAULT_MAX_ACTIVE_DOWNLOADS)
}
pub fn max_active(&self) -> usize {
self.max_active
}
pub fn submit(
&self,
content: ContentId,
sink: Arc<dyn Sink>,
opts: DownloadOptions,
) -> QueuedHandle {
let (events_tx, events_rx) = mpsc::channel(256);
let (result_tx, result_rx) = oneshot::channel();
let job = QueuedJob {
content,
sink,
opts,
events: events_tx,
result: result_tx,
};
if self.submit_tx.send(job).is_err() {
}
QueuedHandle {
events: events_rx,
result: Some(result_rx),
}
}
}
async fn run_job(downloader: &Downloader, job: QueuedJob) {
let mut handle = downloader.download(job.content, job.sink, job.opts);
while let Some(event) = handle.next_event().await {
if job.events.send(event).await.is_err() {
break;
}
}
let _ = job.result.send(handle.join().await);
}
pub struct QueuedHandle {
events: mpsc::Receiver<DownloadEvent>,
result: Option<oneshot::Receiver<Result<u64, DownloadError>>>,
}
impl QueuedHandle {
pub async fn next_event(&mut self) -> Option<DownloadEvent> {
self.events.recv().await
}
pub async fn join(mut self) -> Result<u64, DownloadError> {
match self.result.take() {
Some(rx) => rx.await.unwrap_or(Err(DownloadError::TaskEnded)),
None => Err(DownloadError::TaskEnded),
}
}
}