//! [![crates.io version](https://img.shields.io/crates/v/safina-threadpool.svg)](https://crates.io/crates/safina-threadpool)
//! [![license: Apache 2.0](https://gitlab.com/leonhard-llc/safina-rs/-/raw/main/license-apache-2.0.svg)](http://www.apache.org/licenses/LICENSE-2.0)
//! [![unsafe forbidden](https://gitlab.com/leonhard-llc/safina-rs/-/raw/main/unsafe-forbidden-success.svg)](https://github.com/rust-secure-code/safety-dance/)
//! [![pipeline status](https://gitlab.com/leonhard-llc/safina-rs/badges/main/pipeline.svg)](https://gitlab.com/leonhard-llc/safina-rs/-/pipelines)
//!
//! This is a safe Rust thread pool library.
//!
//! It is part of [`safina`](https://crates.io/crates/safina), a safe async runtime.
//!
//! # Features
//! - `forbid(unsafe_code)`
//! - Depends only on `std`
//! - Good test coverage (100%)
//!
//! # Limitations
//! - Allocates memory
//! - Not optimized
//!
//! # Documentation
//! https://docs.rs/safina-threadpool
//!
//! # Examples
//! ```rust
//! # type ProcessResult = ();
//! # fn process_data(data: (), sender: std::sync::mpsc::Sender<ProcessResult>) -> ProcessResult {
//! #    sender.send(()).unwrap();
//! # }
//! # fn f() {
//! # let data_source = vec![(),()];
//! let pool =
//!     safina_threadpool::ThreadPool::new("worker", 2);
//! let receiver = {
//!     let (sender, receiver) =
//!         std::sync::mpsc::channel();
//!     for data in data_source {
//!         let sender_clone = sender.clone();
//!         pool.schedule(
//!             move || process_data(data, sender_clone));
//!     }
//!     receiver
//! };
//! let results: Vec<ProcessResult> =
//!     receiver.iter().collect();
//! // ...
//! # }
//! ```
//!
//! # Alternatives
//! - [blocking](https://crates.io/crates/blocking)
//!   - Popular
//!   - A little `unsafe` code
//! - [threadpool](https://crates.io/crates/threadpool)
//!   - Popular
//!   - Well maintained
//!   - Dependencies have `unsafe` code
//! - [futures-executor](https://crates.io/crates/futures-executor)
//!   - Very popular
//!   - Full of `unsafe`
//! - [scoped_threadpool](https://crates.io/crates/scoped_threadpool)
//!   - Popular
//!   - Contains `unsafe` code
//! - [scheduled-thread-pool](https://crates.io/crates/scheduled-thread-pool)
//!   - Used by a popular connection pool library
//!   - Dependencies have `unsafe` code
//! - [workerpool](https://crates.io/crates/workerpool)
//!   - Dependencies have `unsafe` code
//! - [threads_pool](https://crates.io/crates/threads_pool)
//!   - Full of `unsafe`
//! - [thread-pool](https://crates.io/crates/thread-pool)
//!   - Old
//!   - Dependencies have `unsafe` code
//! - [tasque](https://crates.io/crates/tasque)
//!   - Dependencies have `unsafe` code
//! - [fast-threadpool](https://crates.io/crates/fast-threadpool)
//!   - Dependencies have `unsafe` code
//! - [blocking-permit](https://crates.io/crates/blocking-permit)
//!   - Full of `unsafe`
//! - [rayon-core](https://crates.io/crates/rayon-core)
//!   - Full of `unsafe`
//!
//! # Changelog
//! - v0.1.0 - First release
//!
//! # TO DO
//! - DONE - Add `schedule` and `try_schedule`
//! - DONE - Add tests
//! - DONE - Add docs
//! - DONE - Publish on crates.io
//! - Add a stress test
//! - Add a benchmark.  See benchmarks in https://crates.io/crates/executors
//! - Add a way for a job to schedule another job on the same thread, with stealing.
//!
//! # Release Process
//! 1. Edit `Cargo.toml` and bump version number.
//! 1. Run `./release.sh`
#![forbid(unsafe_code)]

use core::fmt::{Display, Formatter};
use core::sync::atomic::{AtomicUsize, Ordering};
use core::time::Duration;
use std::error::Error;
use std::sync::mpsc::{Receiver, RecvTimeoutError, SyncSender, TrySendError};
use std::sync::{Arc, Mutex};

struct AtomicCounter {
    next_value: AtomicUsize,
}
impl AtomicCounter {
    pub fn new() -> Self {
        Self {
            next_value: AtomicUsize::new(0),
        }
    }
    pub fn next(&self) -> usize {
        self.next_value.fetch_add(1, Ordering::AcqRel)
    }
}

/// Returned by [`try_schedule`](struct.ThreadPool.html#method.try_schedule)
/// when the queue is full.
/// This can happen when the program schedules many closures at one time.
/// It can also happen when closures panic their threads.  The pool's
/// throughput goes down when it must create new threads.
#[derive(Debug)]
pub struct QueueFull {}
impl Display for QueueFull {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
        std::fmt::Debug::fmt(self, f)
    }
}
impl Error for QueueFull {}

struct Inner {
    name: &'static str,
    next_name_num: AtomicCounter,
    size: usize,
    receiver: Mutex<Receiver<Box<dyn FnOnce() + Send>>>,
}
impl Inner {
    pub fn num_live_threads(self: &Arc<Inner>) -> usize {
        Arc::strong_count(self) - 1
    }

    fn work(self: &Arc<Inner>) {
        loop {
            let recv_result = self
                .receiver
                .lock()
                .unwrap()
                .recv_timeout(Duration::from_millis(100));
            self.start_threads();
            match recv_result {
                Ok(f) => f(),
                Err(RecvTimeoutError::Timeout) => {}
                // ThreadPool was dropped.
                Err(RecvTimeoutError::Disconnected) => return,
            };
            self.start_threads();
        }
    }

    fn start_thread(self: &Arc<Inner>) {
        let self_clone = self.clone();
        if self.num_live_threads() <= self.size {
            std::thread::Builder::new()
                .name(format!("{}{}", self.name, self.next_name_num.next()))
                .spawn(move || self_clone.work())
                .unwrap();
        }
    }

    fn start_threads(self: &Arc<Inner>) {
        while self.num_live_threads() < self.size {
            self.start_thread();
        }
    }

    fn start_thread_if_none(self: &Arc<Inner>) {
        if self.num_live_threads() < 1 {
            self.start_thread();
        }
    }
}

/// A collection of threads and a queue for jobs (`FnOnce` structs) they execute.
///
/// Threads die when they execute a job that panics.
/// If one thread survives, it will recreate all the threads.
/// The next call to [`schedule`](#method.schedule) or [`try_schedule`](#method.try_schedule)
/// also recreates threads.
///
/// TODO(mleonhard) Show how to use safina_timer to recover when all threads panicked.
///
/// # Example
/// ```rust
/// # type ProcessResult = ();
/// # fn process_data(data: (), sender: std::sync::mpsc::Sender<ProcessResult>) -> ProcessResult {
/// #    sender.send(()).unwrap();
/// # }
/// # fn f() {
/// # let data_source = vec![(),()];
/// let pool =
///     safina_threadpool::ThreadPool::new("worker", 2);
/// let receiver = {
///     let (sender, receiver) =
///         std::sync::mpsc::channel();
///     for data in data_source {
///         let sender_clone = sender.clone();
///         pool.schedule(
///             move || process_data(data, sender_clone));
///     }
///     receiver
/// };
/// let results: Vec<ProcessResult> =
///     receiver.iter().collect();
/// // ...
/// # }
/// ```
pub struct ThreadPool {
    inner: Arc<Inner>,
    sender: SyncSender<Box<dyn FnOnce() + Send>>,
}

impl ThreadPool {
    /// Creates a new thread pool containing `size` threads.
    /// The threads all start immediately.
    ///
    /// Threads are named with `name` with a number.
    /// For example, `ThreadPool::new("worker", 2)`
    /// creates threads named "worker-1" and "worker-2".
    /// If one of those threads panics, the pool creates "worker-3".
    ///
    /// Panics if `name` is empty or `size` is zero.
    pub fn new(name: &'static str, size: usize) -> Self {
        if name.is_empty() {
            panic!("ThreadPool::new called with empty name")
        }
        if size < 1 {
            panic!("ThreadPool::new called with invalid size value: {:?}", size)
        }
        // Use a channel with bounded size.
        // If the channel was unbounded, the process could OOM when throughput goes down.
        let (sender, receiver) = std::sync::mpsc::sync_channel(size * 200);
        let pool = ThreadPool {
            inner: Arc::new(Inner {
                name,
                next_name_num: AtomicCounter::new(),
                size,
                receiver: Mutex::new(receiver),
            }),
            sender,
        };
        pool.inner.start_threads();
        pool
    }

    /// Returns the number of threads in the pool.
    pub fn size(&self) -> usize {
        self.inner.size
    }

    /// Returns the number of threads currently alive.
    pub fn num_live_threads(&self) -> usize {
        self.inner.num_live_threads()
    }

    /// Adds a job to the queue.  The next idle thread will execute it.
    /// Jobs are started in FIFO order.
    ///
    /// Blocks when the queue is full.
    /// See [`try_schedule`](#method.try_schedule).
    ///
    /// Recreates any threads that panicked.
    ///
    /// Puts `f` in a [`Box`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html) before
    /// adding it to the queue.
    pub fn schedule(&self, f: impl FnOnce() + Send + 'static) {
        // TODO(mleonhard) Log a warning before blocking.
        // If all workers panicked and the queue is full, adding to the queue will
        // block forever.  So we start at least one thread first.
        self.inner.start_thread_if_none();
        self.sender.send(Box::new(f)).unwrap();
        self.inner.start_threads();
    }

    /// Adds a job to the queue.  The next idle thread will execute it.
    /// Jobs are started in FIFO order.
    ///
    /// Returns `Err(QueueFull)` when the queue is full.
    ///
    /// Recreates any threads that panicked.
    ///
    /// Puts `f` in a [`Box`](https://doc.rust-lang.org/stable/std/boxed/struct.Box.html) before
    /// adding it to the queue.
    pub fn try_schedule(&self, f: impl FnOnce() + Send + 'static) -> Result<(), QueueFull> {
        self.inner.start_thread_if_none();
        match self.sender.try_send(Box::new(f)) {
            Ok(_) => {
                self.inner.start_threads();
                Ok(())
            }
            Err(TrySendError::Disconnected(_)) => unreachable!(),
            Err(TrySendError::Full(_)) => Err(QueueFull {}),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use core::fmt::Debug;
    use core::ops::Range;
    use core::time::Duration;
    use std::time::Instant;

    fn assert_in_range<T: PartialOrd + Debug>(range: Range<T>, value: T) {
        if range.is_empty() {
            panic!("invalid range {:?}", range)
        }
        println!(
            "measured concurrency value {:?}, expected range {:?}",
            value, range,
        );
        if !range.contains(&value) {
            panic!(
                "measured concurrency value {:?} out of range {:?}",
                value, range,
            );
        }
    }

    pub fn assert_elapsed(before: Instant, range_ms: Range<u64>) {
        if range_ms.is_empty() {
            panic!("invalid range {:?}", range_ms)
        }
        let elapsed = before.elapsed();
        let duration_range =
            Duration::from_millis(range_ms.start)..Duration::from_millis(range_ms.end);
        if !duration_range.contains(&elapsed) {
            panic!("{:?} elapsed, out of range {:?}", elapsed, duration_range);
        }
    }

    fn measure_concurrency(pool: &ThreadPool, num_jobs: usize) -> f32 {
        const WAIT_DURATION: Duration = Duration::from_millis(100);
        let before = Instant::now();
        let receiver = {
            let (sender, receiver) = std::sync::mpsc::channel();
            for _ in 0..num_jobs {
                let sender_clone = sender.clone();
                pool.schedule(move || {
                    std::thread::sleep(WAIT_DURATION);
                    sender_clone.send(()).unwrap();
                });
            }
            receiver
        };
        for _ in 0..num_jobs {
            receiver.recv_timeout(Duration::from_millis(500)).unwrap();
        }
        let elapsed = before.elapsed();
        elapsed.as_secs_f32() / WAIT_DURATION.as_secs_f32()
    }

    fn sleep(ms: u64) {
        std::thread::sleep(Duration::from_millis(ms));
    }

    #[test]
    fn atomic_counter() {
        let counter = Arc::new(AtomicCounter::new());
        assert_eq!(0, counter.next());
        assert_eq!(1, counter.next());
        assert_eq!(2, counter.next());
    }

    #[test]
    fn atomic_counter_many_readers() {
        let receiver = {
            let counter = Arc::new(AtomicCounter::new());
            let (sender, receiver) = std::sync::mpsc::channel();
            for _ in 0..10 {
                let counter_clone = counter.clone();
                let sender_clone = sender.clone();
                std::thread::spawn(move || {
                    for _ in 0..10 {
                        sender_clone.send(counter_clone.next()).unwrap();
                    }
                });
            }
            receiver
        };
        let mut values: Vec<usize> = receiver.iter().collect();
        values.sort();
        assert_eq!((0usize..100).collect::<Vec<usize>>(), values);
    }

    #[test]
    fn queue_full_display() {
        assert_eq!("QueueFull", format!("{}", QueueFull {}));
    }

    #[test]
    fn empty_name() {
        match std::panic::catch_unwind(|| ThreadPool::new("", 1)) {
            Ok(_) => panic!("expected panic"),
            Err(_) => {}
        }
    }

    #[test]
    fn zero_size() {
        match std::panic::catch_unwind(|| ThreadPool::new("pool1", 0)) {
            Ok(_) => panic!("expected panic"),
            Err(_) => {}
        }
    }

    #[test]
    fn test_size() {
        let pool = ThreadPool::new("pool1", 3);
        assert_eq!(3, pool.size());
    }

    #[test]
    fn test_num_live_threads() {
        let pool = ThreadPool::new("pool1", 3);
        sleep(100);
        assert_eq!(3, pool.num_live_threads());
        pool.schedule(move || {
            sleep(100);
            panic!("ignore this panic")
        });
        pool.schedule(move || {
            sleep(100);
            panic!("ignore this panic")
        });
        pool.schedule(move || {
            sleep(100);
            panic!("ignore this panic")
        });
        sleep(200);
        assert_eq!(0, pool.num_live_threads());
        pool.schedule(move || {});
        assert_eq!(3, pool.num_live_threads());
    }

    #[test]
    fn schedule_should_run_the_fn() {
        let pool = ThreadPool::new("pool1", 1);
        let before = Instant::now();
        let (sender, receiver) = std::sync::mpsc::channel();
        pool.schedule(move || {
            sender.send(()).unwrap();
        });
        receiver.recv_timeout(Duration::from_millis(500)).unwrap();
        assert_elapsed(before, 0..100);
    }

    #[test]
    fn schedule_should_start_a_thread_if_none() {
        let pool = ThreadPool::new("pool1", 3);
        sleep(100);
        pool.schedule(move || {
            sleep(100);
            panic!("ignore this panic")
        });
        pool.schedule(move || {
            sleep(100);
            panic!("ignore this panic")
        });
        pool.schedule(move || {
            sleep(100);
            panic!("ignore this panic")
        });
        sleep(200);
        assert_eq!(0, pool.num_live_threads());
        pool.schedule(|| {});
        assert_eq!(3, pool.num_live_threads());
    }

    #[test]
    fn try_schedule_should_run_the_fn() {
        let pool = ThreadPool::new("pool1", 1);
        let before = Instant::now();
        let (sender, receiver) = std::sync::mpsc::channel();
        pool.try_schedule(move || {
            sender.send(()).unwrap();
        })
        .unwrap();
        receiver.recv_timeout(Duration::from_millis(500)).unwrap();
        assert_elapsed(before, 0..100);
    }

    #[test]
    fn try_schedule_queue_full() {
        let pool = ThreadPool::new("pool1", 1);
        let before = Instant::now();
        while Instant::now() - before < Duration::from_millis(500) {
            match pool.try_schedule(move || panic!("ignore this panic")) {
                Ok(_) => {}
                Err(e) => {
                    println!("try_schedule got {:?}", e);
                    // Sometimes a thread's panic message is interspersed with the
                    // test runner result, confusing the IDE.  We sleep here to
                    // let the threads run and finish before we return.
                    sleep(100);
                    return;
                }
            }
        }
        panic!("timeout");
    }

    #[test]
    fn check_concurrency1() {
        let pool = ThreadPool::new("pool1", 1);
        assert_in_range(1.0..1.99, measure_concurrency(&pool, 1));
        assert_in_range(2.0..2.99, measure_concurrency(&pool, 2));
    }

    #[test]
    fn check_concurrency2() {
        let pool = ThreadPool::new("pool1", 2);
        assert_in_range(1.0..1.99, measure_concurrency(&pool, 1));
        assert_in_range(1.0..1.99, measure_concurrency(&pool, 2));
        assert_in_range(2.0..2.99, measure_concurrency(&pool, 3));
        assert_in_range(2.0..2.99, measure_concurrency(&pool, 4));
    }

    #[test]
    fn check_concurrency5() {
        let pool = ThreadPool::new("pool1", 5);
        assert_in_range(1.0..1.99, measure_concurrency(&pool, 5));
        assert_in_range(2.0..2.99, measure_concurrency(&pool, 6));
    }

    #[test]
    fn should_quickly_respawn_panicked_threads() {
        let pool = ThreadPool::new("pool1", 2);
        sleep(50);
        pool.schedule(move || panic!("ignore this panic"));
        sleep(100);
        assert_eq!(1, pool.num_live_threads());
        sleep(200);
        assert_eq!(2, pool.num_live_threads());
        assert_in_range(1.0..1.99, measure_concurrency(&pool, 2));
    }

    #[test]
    fn should_respawn_after_recv() {
        let pool = ThreadPool::new("pool1", 2);
        sleep(50);
        pool.schedule(move || panic!("ignore this panic"));
        sleep(100);
        assert_eq!(1, pool.num_live_threads());
        pool.schedule(move || sleep(200));
        sleep(100);
        assert_eq!(2, pool.num_live_threads());
    }

    #[test]
    fn should_respawn_after_executing_job() {
        let pool = ThreadPool::new("pool1", 2);
        pool.schedule(move || sleep(200));
        pool.schedule(move || panic!("ignore this panic"));
        sleep(100);
        assert_eq!(1, pool.num_live_threads());
        sleep(200);
        assert_eq!(2, pool.num_live_threads());
    }
}