ncd 0.1.3

use std::{io, mem, panic, sync::Arc, thread};

use super::TaskQueue;

/// General purpose threadpool. Unwinding panics in tasks are caught.
///
/// Dropping a pool drains queued tasks and waits for its workers to exit. This
/// can block indefinitely if a running or queued task does not return. If a
/// worker drops the final pool owner, joining is deferred until that task
/// returns to avoid deadlocking the pool. Non-string custom panic payloads are
/// leaked because running their arbitrary destructors could terminate a worker.
pub struct ThreadPool {
    task_queue: Arc<TaskQueue>,
    workers: Vec<thread::JoinHandle<()>>,
}

impl ThreadPool {
    /// Creates a new threadpool.
    ///
    /// # Notes
    /// * Thread names will be "name_pattern-X", where X is the number
    ///   in creation order.
    ///
    /// # Panics
    /// * Panics if `num_threads` is zero.
    /// * Panics if `std::thread::Builder` fails.
    pub fn new(name_pattern: &str, num_threads: usize) -> Self {
        assert!(
            num_threads > 0,
            "a threadpool must have at least one worker"
        );

        Self::try_new_with(name_pattern, num_threads, |name, task_queue| {
            thread::Builder::new()
                .name(name)
                .spawn(move || worker_loop(task_queue))
        })
        .unwrap()
    }

    fn try_new_with<F>(name_pattern: &str, num_threads: usize, mut spawn: F) -> io::Result<Self>
    where
        F: FnMut(String, Arc<TaskQueue>) -> io::Result<thread::JoinHandle<()>>,
    {
        let task_queue = Arc::new(TaskQueue::new());
        let mut pool = ThreadPool {
            task_queue,
            workers: Vec::with_capacity(num_threads),
        };

        for i in 0..num_threads {
            let worker = spawn(format!("{name_pattern}-{i}"), Arc::clone(&pool.task_queue))?;
            pool.workers.push(worker);
        }

        Ok(pool)
    }

    /// Submit a task to the threadpool.
    pub fn submit<F>(&self, task: F)
    where
        F: FnOnce() + Send + 'static,
    {
        if let Err(task) = self.task_queue.push(Box::new(task)) {
            // Task destructors are user code and must run outside the queue lock.
            drop(task);
        }
    }
}

impl Drop for ThreadPool {
    fn drop(&mut self) {
        self.task_queue.close();

        let workers = mem::take(&mut self.workers);
        let current_thread = thread::current().id();
        let called_from_worker = workers
            .iter()
            .any(|worker| worker.thread().id() == current_thread);

        if called_from_worker {
            // The current worker cannot finish until this destructor returns. A
            // reaper can safely join it and any workers that depend on it.
            let _ = thread::Builder::new()
                .name("ncd-reaper".to_owned())
                .spawn(move || join_workers(workers));
        } else {
            join_workers(workers);
        }
    }
}

fn worker_loop(task_queue: Arc<TaskQueue>) {
    while let Some(task) = task_queue.pop() {
        if let Err(payload) = panic::catch_unwind(panic::AssertUnwindSafe(task)) {
            discard_panic_payload(payload);
        }
    }
}

fn join_workers(workers: Vec<thread::JoinHandle<()>>) {
    let current_thread = thread::current().id();

    for worker in workers {
        if worker.thread().id() == current_thread {
            continue;
        }

        if let Err(payload) = worker.join() {
            discard_panic_payload(payload);
        }
    }
}

fn discard_panic_payload(payload: Box<dyn std::any::Any + Send>) {
    if payload.is::<&'static str>() || payload.is::<String>() {
        drop(payload);
    } else {
        // Dropping an arbitrary payload can panic or abort through nested
        // destructor panics. Preserve worker capacity instead.
        mem::forget(payload);
    }
}

#[cfg(test)]
mod tests {
    use std::{
        sync::{
            atomic::{AtomicUsize, Ordering},
            mpsc, Arc, Barrier,
        },
        time::Duration,
    };

    use super::*;

    #[test]
    fn drop_drains_pending_tasks() {
        let completed = Arc::new(AtomicUsize::new(0));

        {
            let pool = ThreadPool::new("drain", 2);
            for _ in 0..100 {
                let completed = Arc::clone(&completed);
                pool.submit(move || {
                    completed.fetch_add(1, Ordering::Relaxed);
                });
            }
        }

        assert_eq!(completed.load(Ordering::Relaxed), 100);
    }

    #[test]
    #[should_panic(expected = "a threadpool must have at least one worker")]
    fn zero_workers_are_rejected() {
        let _ = ThreadPool::new("zero", 0);
    }

    #[test]
    fn ordinary_panic_does_not_stop_worker() {
        let pool = ThreadPool::new("panic", 1);
        let (sender, receiver) = mpsc::channel();

        pool.submit(|| panic!("expected task panic"));
        pool.submit(move || sender.send(()).unwrap());

        receiver.recv_timeout(Duration::from_secs(1)).unwrap();
    }

    #[test]
    fn custom_panic_payload_does_not_stop_worker() {
        struct PanicsOnDrop;

        impl Drop for PanicsOnDrop {
            fn drop(&mut self) {
                panic!("expected payload drop panic");
            }
        }

        let pool = ThreadPool::new("panic-payload", 1);
        let (sender, receiver) = mpsc::channel();

        pool.submit(|| panic::panic_any(PanicsOnDrop));
        pool.submit(move || sender.send(()).unwrap());

        receiver.recv_timeout(Duration::from_secs(1)).unwrap();
    }

    #[test]
    fn final_drop_from_worker_does_not_join_dependent_workers() {
        let pool = Arc::new(ThreadPool::new("worker-drop", 2));
        let worker_pool = Arc::clone(&pool);
        let barrier = Arc::new(Barrier::new(3));
        let dropping_worker_barrier = Arc::clone(&barrier);
        let dependent_worker_barrier = Arc::clone(&barrier);
        let (release_sender, release_receiver) = mpsc::channel();
        let (completed_sender, completed_receiver) = mpsc::channel();

        pool.submit(move || {
            dropping_worker_barrier.wait();
            drop(worker_pool);
            release_sender.send(()).unwrap();
        });
        pool.submit(move || {
            dependent_worker_barrier.wait();
            release_receiver.recv().unwrap();
            completed_sender.send(()).unwrap();
        });

        drop(pool);
        barrier.wait();

        completed_receiver
            .recv_timeout(Duration::from_secs(1))
            .unwrap();
    }

    #[test]
    fn construction_failure_joins_started_workers() {
        let (sender, receiver) = mpsc::channel();
        let mut spawn_count = 0;

        let result = ThreadPool::try_new_with("partial", 2, move |_, task_queue| {
            spawn_count += 1;
            if spawn_count == 2 {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "expected spawn failure",
                ));
            }

            let sender = sender.clone();
            thread::Builder::new().spawn(move || {
                worker_loop(task_queue);
                sender.send(()).unwrap();
            })
        });

        assert!(result.is_err());
        receiver.recv_timeout(Duration::from_secs(1)).unwrap();
    }
}