panda-rpki/crates/panda-rpki-validator/src/parallel/repo_worker/pools.rs

// Worker-pool channels, lifecycle, and worker loops.

enum RepoWorkerMessage {
    Task(RepoSyncTask),
    Shutdown,
}

enum TransportWorkerMessage {
    Task(RepoTransportTask),
    Shutdown,
}

pub struct RepoWorkerPool<E: RepoSyncExecutor> {
    config: RepoWorkerPoolConfig,
    task_tx: Sender<RepoWorkerMessage>,
    result_rx: Receiver<RepoSyncResultEnvelope>,
    workers: Vec<JoinHandle<()>>,
    _executor: Arc<E>,
}

pub struct RepoTransportWorkerPool<E: RepoTransportExecutor> {
    config: RepoWorkerPoolConfig,
    task_tx: Sender<TransportWorkerMessage>,
    result_rx: Receiver<RepoTransportResultEnvelope>,
    workers: Vec<JoinHandle<()>>,
    _executor: Arc<E>,
}

impl<E: RepoSyncExecutor> RepoWorkerPool<E> {
    pub fn new(config: RepoWorkerPoolConfig, executor: E) -> Result<Self, String> {
        if config.max_workers == 0 {
            return Err("RepoWorkerPool requires at least one worker".to_string());
        }

        let executor = Arc::new(executor);
        let (task_tx, task_rx) = mpsc::channel::<RepoWorkerMessage>();
        let (result_tx, result_rx) = mpsc::channel::<RepoSyncResultEnvelope>();
        let shared_task_rx = Arc::new(Mutex::new(task_rx));

        let mut workers = Vec::with_capacity(config.max_workers);
        for idx in 0..config.max_workers {
            let task_rx = Arc::clone(&shared_task_rx);
            let result_tx = result_tx.clone();
            let executor = Arc::clone(&executor);
            workers.push(
                thread::Builder::new()
                    .name(format!("repo-sync-worker-{idx}"))
                    .spawn(move || worker_loop(task_rx, result_tx, executor))
                    .map_err(|e| format!("spawn repo worker failed: {e}"))?,
            );
        }

        Ok(Self {
            config,
            task_tx,
            result_rx,
            workers,
            _executor: executor,
        })
    }

    pub fn worker_count(&self) -> usize {
        self.config.max_workers
    }

    pub fn submit(&self, task: RepoSyncTask) -> Result<(), String> {
        self.task_tx
            .send(RepoWorkerMessage::Task(task))
            .map_err(|e| format!("submit repo task failed: {e}"))
    }

    pub fn recv_result_timeout(
        &self,
        timeout: Duration,
    ) -> Result<Option<RepoSyncResultEnvelope>, String> {
        match self.result_rx.recv_timeout(timeout) {
            Ok(msg) => Ok(Some(msg)),
            Err(RecvTimeoutError::Timeout) => Ok(None),
            Err(RecvTimeoutError::Disconnected) => {
                Err("repo worker result channel disconnected".to_string())
            }
        }
    }

    pub fn shutdown(mut self) -> Result<(), String> {
        self.shutdown_inner()
    }

    fn shutdown_inner(&mut self) -> Result<(), String> {
        if self.workers.is_empty() {
            return Ok(());
        }

        for _ in 0..self.workers.len() {
            self.task_tx
                .send(RepoWorkerMessage::Shutdown)
                .map_err(|e| format!("send shutdown to repo worker failed: {e}"))?;
        }

        let mut first_err: Option<String> = None;
        for handle in self.workers.drain(..) {
            if let Err(e) = handle.join() {
                if first_err.is_none() {
                    first_err = Some(format!("join repo worker failed: {e:?}"));
                }
            }
        }

        if let Some(err) = first_err {
            return Err(err);
        }
        Ok(())
    }
}

impl<E: RepoSyncExecutor> Drop for RepoWorkerPool<E> {
    fn drop(&mut self) {
        let _ = self.shutdown_inner();
    }
}

impl<E: RepoTransportExecutor> RepoTransportWorkerPool<E> {
    pub fn new(config: RepoWorkerPoolConfig, executor: E) -> Result<Self, String> {
        if config.max_workers == 0 {
            return Err("RepoTransportWorkerPool requires at least one worker".to_string());
        }
        let executor = Arc::new(executor);
        let (task_tx, task_rx) = mpsc::channel::<TransportWorkerMessage>();
        let (result_tx, result_rx) = mpsc::channel::<RepoTransportResultEnvelope>();
        let shared_task_rx = Arc::new(Mutex::new(task_rx));
        let mut workers = Vec::with_capacity(config.max_workers);
        for idx in 0..config.max_workers {
            let task_rx = Arc::clone(&shared_task_rx);
            let result_tx = result_tx.clone();
            let executor = Arc::clone(&executor);
            workers.push(
                thread::Builder::new()
                    .name(format!("repo-transport-worker-{idx}"))
                    .spawn(move || transport_worker_loop(task_rx, result_tx, executor))
                    .map_err(|e| format!("spawn repo transport worker failed: {e}"))?,
            );
        }
        Ok(Self {
            config,
            task_tx,
            result_rx,
            workers,
            _executor: executor,
        })
    }

    pub fn submit(&self, task: RepoTransportTask) -> Result<(), String> {
        self.task_tx
            .send(TransportWorkerMessage::Task(task))
            .map_err(|e| format!("submit repo transport task failed: {e}"))
    }

    pub fn recv_result_timeout(
        &self,
        timeout: Duration,
    ) -> Result<Option<RepoTransportResultEnvelope>, String> {
        match self.result_rx.recv_timeout(timeout) {
            Ok(msg) => Ok(Some(msg)),
            Err(RecvTimeoutError::Timeout) => Ok(None),
            Err(RecvTimeoutError::Disconnected) => {
                Err("repo transport worker result channel disconnected".to_string())
            }
        }
    }

    pub fn worker_count(&self) -> usize {
        self.config.max_workers
    }

    pub fn shutdown(mut self) -> Result<(), String> {
        self.shutdown_inner()
    }

    fn shutdown_inner(&mut self) -> Result<(), String> {
        if self.workers.is_empty() {
            return Ok(());
        }
        for _ in 0..self.workers.len() {
            self.task_tx
                .send(TransportWorkerMessage::Shutdown)
                .map_err(|e| format!("send shutdown to repo transport worker failed: {e}"))?;
        }
        let mut first_err: Option<String> = None;
        for handle in self.workers.drain(..) {
            if let Err(e) = handle.join() {
                if first_err.is_none() {
                    first_err = Some(format!("join repo transport worker failed: {e:?}"));
                }
            }
        }
        if let Some(err) = first_err {
            return Err(err);
        }
        Ok(())
    }
}

impl<E: RepoTransportExecutor> Drop for RepoTransportWorkerPool<E> {
    fn drop(&mut self) {
        let _ = self.shutdown_inner();
    }
}

fn worker_loop<E: RepoSyncExecutor>(
    task_rx: Arc<Mutex<Receiver<RepoWorkerMessage>>>,
    result_tx: Sender<RepoSyncResultEnvelope>,
    executor: Arc<E>,
) {
    loop {
        let message = {
            let rx = task_rx.lock().expect("repo worker receiver lock poisoned");
            rx.recv()
        };

        match message {
            Ok(RepoWorkerMessage::Task(task)) => {
                let result = executor.execute(task);
                if result_tx.send(result).is_err() {
                    break;
                }
            }
            Ok(RepoWorkerMessage::Shutdown) | Err(_) => break,
        }
    }
}

fn transport_worker_loop<E: RepoTransportExecutor>(
    task_rx: Arc<Mutex<Receiver<TransportWorkerMessage>>>,
    result_tx: Sender<RepoTransportResultEnvelope>,
    executor: Arc<E>,
) {
    loop {
        let message = {
            let rx = task_rx
                .lock()
                .expect("repo transport worker receiver lock poisoned");
            rx.recv()
        };

        match message {
            Ok(TransportWorkerMessage::Task(task)) => {
                let result = executor.execute_transport(task);
                if result_tx.send(result).is_err() {
                    break;
                }
            }
            Ok(TransportWorkerMessage::Shutdown) | Err(_) => break,
        }
    }
}