use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
use crate::bandwidth::BandwidthLimiter;
use crate::download_engine::{ConnectionTracker, ProgressUpdate, WorkerPool, build_job_submission};
use nzb_core::config::ServerConfig;
use nzb_core::models::NzbJob;
use parking_lot::Mutex;
#[allow(clippy::double_must_use)]
#[async_trait::async_trait]
pub trait DispatchEngine: Send + Sync {
fn start(&self);
fn submit_job(&self, job: &NzbJob, progress_tx: mpsc::Sender<ProgressUpdate>);
fn pause_job(&self, job_id: &str);
fn resume_job(&self, job_id: &str);
fn cancel_job(&self, job_id: &str);
fn abort_job(&self, job_id: &str, reason: String) -> bool;
fn has_job(&self, job_id: &str) -> bool;
fn job_active_secs(&self, _job_id: &str) -> Option<f64> {
None
}
fn release_completed_job(&self, job_id: &str);
fn update_servers(&self, servers: Vec<ServerConfig>);
fn connection_snapshot(&self) -> Vec<(String, usize, usize)>;
fn active_connection_snapshot(&self) -> Vec<(String, usize, usize)>;
fn connection_total(&self) -> usize;
fn reconcile_servers(&self);
fn set_max_worker_idle(&self, d: Duration);
fn eviction_count(&self) -> u64;
fn server_stats_snapshot(&self) -> Vec<(String, ServerAttemptStats)> {
Vec::new()
}
async fn shutdown(&self);
}
#[derive(Debug, Clone, Copy, Default)]
pub struct ServerAttemptStats {
pub attempted: u64,
pub succeeded: u64,
pub not_found: u64,
pub transient_failed: u64,
}
pub struct DispatchHandle(Arc<WorkerPool>);
impl DispatchHandle {
pub fn new(
servers: Arc<Mutex<Vec<ServerConfig>>>,
bandwidth: Arc<BandwidthLimiter>,
article_timeout_secs: u64,
) -> Self {
let tracker = Arc::new(ConnectionTracker::new());
for server in servers.lock().iter() {
tracker.set_limit(&server.id, &server.name, server.connections as usize);
}
Self(WorkerPool::new(
servers,
bandwidth,
tracker,
article_timeout_secs,
))
}
}
#[async_trait::async_trait]
impl DispatchEngine for DispatchHandle {
fn start(&self) {
self.0.start();
}
fn submit_job(&self, job: &NzbJob, progress_tx: mpsc::Sender<ProgressUpdate>) {
let enabled_server_ids = self
.0
.servers
.lock()
.iter()
.filter(|server| server.enabled)
.map(|server| server.id.clone())
.collect::<Vec<_>>();
let (ctx, items) = build_job_submission(job, progress_tx, &enabled_server_ids);
self.0.submit_job(ctx, items);
}
fn pause_job(&self, job_id: &str) {
self.0.pause_job(job_id);
}
fn resume_job(&self, job_id: &str) {
self.0.resume_job(job_id);
}
fn cancel_job(&self, job_id: &str) {
self.0.cancel_job(job_id);
}
fn abort_job(&self, job_id: &str, reason: String) -> bool {
self.0.abort_job(job_id, reason)
}
fn has_job(&self, job_id: &str) -> bool {
self.0.has_job(job_id)
}
fn job_active_secs(&self, job_id: &str) -> Option<f64> {
self.0.job_active_secs(job_id)
}
fn release_completed_job(&self, job_id: &str) {
self.0.release_completed_job(job_id);
}
fn update_servers(&self, servers: Vec<ServerConfig>) {
let new_ids: std::collections::HashSet<_> =
servers.iter().map(|server| server.id.clone()).collect();
for (old_id, _, _) in self.0.conn_tracker().snapshot() {
if !new_ids.contains(&old_id) {
self.0.conn_tracker().remove_server(&old_id);
}
}
for server in &servers {
self.0
.conn_tracker()
.set_limit(&server.id, &server.name, server.connections as usize);
}
*self.0.servers.lock() = servers;
self.0.reconcile_servers();
}
fn connection_snapshot(&self) -> Vec<(String, usize, usize)> {
self.0.conn_tracker().snapshot()
}
fn active_connection_snapshot(&self) -> Vec<(String, usize, usize)> {
self.0.conn_tracker().connected_snapshot()
}
fn connection_total(&self) -> usize {
self.0.conn_tracker().total()
}
fn reconcile_servers(&self) {
self.0.reconcile_servers();
}
fn set_max_worker_idle(&self, d: Duration) {
self.0.set_max_worker_idle(d);
}
fn eviction_count(&self) -> u64 {
self.0.eviction_count()
}
async fn shutdown(&self) {
self.0.shutdown().await;
}
}