jtp 0.1.0

use crate::{
    task::{Task, TaskFn, TaskListeners},
    worker::Worker,
    TPError, TPResult,
};

use crossbeam_channel::{Receiver, Sender, TrySendError};

use std::{panic::UnwindSafe, sync::Arc, time::Duration};

/// If a task is rejected, the task will be handled by this.
pub enum RejectedTaskHandler {
    /// Returns a [`TPError::AbortError`].
    Abort,

    /// Nothing to do, just return [`Ok`].
    Discard,

    /// Immediately run the rejected task in the caller thread.
    CallerRuns,
}

/// A `ThreadPool` consists of a collection of threads and a bounded
/// task queue.
///
/// # Bounded Task Queue(channel)
/// A bounded task queue is a method of delivering messages across
/// multiple threads. In this thread pool, we use it to pass tasks to
/// an available worker thread.
///
/// # Worker Thread
/// We use a special struct, `Worker`, to represent a thread that is
/// always accepting tasks and executing them. In this library,
/// there are two kinds of the worker:
/// 1. Core worker: A thread in the worker never be terminated except
/// the associated thread pool is disconnected.
/// 2. Non-core worker: A thread in this worker can be idle if no task
/// is received for a certain period of time(`keep_alive_time`).
///
/// This thread pool will store core workers and non-core workers in
/// two vectors. When you execute a task, it creates a core worker to
/// process the task if the core worker vector is not full, otherwise
/// the task will be sent to the task channel queue. If the channel
/// buffer is full, it attempts to find an idle worker or creates a
/// non-core worker to process the task.
pub struct ThreadPool {
    pub(crate) sender: Option<Sender<Task>>,
    pub(crate) reciver: Receiver<Task>,
    pub(crate) core_workers: Vec<Worker>,
    pub(crate) workers: Vec<Worker>,
    pub(crate) core_pool_size: usize,
    pub(crate) max_pool_size: usize,
    pub(crate) keep_alive_time: Duration,
    pub(crate) rejected_task_handler: RejectedTaskHandler,
    pub(crate) next_task_id: usize,
    pub(crate) task_lisenters: Arc<TaskListeners>,
}

impl ThreadPool {
    /// Executes the given task in the future.
    ///
    /// If the task queue is full and no thread can be allocated to
    /// execute the task, the task is handled by the setted
    /// [`RejectedTaskHandler`].
    ///
    /// # Errors
    /// Returns an `AbortError` if the task is handled by the
    /// [`RejectedTaskHandler::Abort`], or returns a `DisconnectedError`
    /// if the channel was disconnected.
    pub fn execute<F>(&mut self, task_fn: F) -> Result<(), TPError>
    where
        F: FnOnce() + Send + UnwindSafe + 'static,
    {
        if self.sender.is_none() {
            return Err(TPError::DisconnectedError);
        }
        let task = self.create_task(Box::new(task_fn));
        if self.core_workers.len() < self.core_pool_size {
            // Create a core thread to process the task.
            self.add_worker(task, true);
            Ok(())
        } else {
            self.try_send_task(task)
        }
    }

    /// Close the current task channel.
    ///
    /// If the task channel is disconnected, all worker threads will
    /// be terminated and the thread pool will not accept any tasks.
    ///
    /// # Examples
    ///
    /// ```
    /// use jtp::ThreadPoolBuilder;
    /// let mut thread_pool = ThreadPoolBuilder::default()
    ///     .build()
    ///     .unwrap();
    ///
    /// thread_pool.close_channel();
    /// assert!(thread_pool.execute(|| {
    ///     println!("Hello");
    /// }).is_err());
    /// ```
    pub fn close_channel(&mut self) {
        self.sender.take();
    }

    /// Waits for all worker threads to finish. Note that is worker
    /// threads instead of tasks.
    ///
    /// # Errors
    /// An error is returned if a thread panics.
    ///
    /// # Examples
    ///
    /// ```
    /// use jtp::ThreadPoolBuilder;
    /// let mut thread_pool = ThreadPoolBuilder::default()
    ///     .build()
    ///     .unwrap();
    /// for _ in 0..10 {
    ///     thread_pool.execute(|| {
    ///         std::thread::sleep(std::time::Duration::from_millis(20));
    ///     });
    /// }
    ///
    /// thread_pool.close_channel();
    ///
    /// // Block current thread until all worker threads are finished.
    /// thread_pool.wait().unwrap();
    /// ```
    ///
    /// Panics if the task channel is not disconnected.
    /// ``` should_panic
    /// use jtp::ThreadPoolBuilder;
    /// let mut thread_pool = ThreadPoolBuilder::default()
    ///     .build()
    ///     .unwrap();
    /// thread_pool.wait().unwrap();
    /// ```
    pub fn wait(self) -> std::thread::Result<()> {
        assert!(self.sender.is_none());
        Self::wait_workers(self.core_workers)?;
        Self::wait_workers(self.workers)
    }

    fn wait_workers(workers: Vec<Worker>) -> std::thread::Result<()> {
        for worker in workers {
            worker.join()?
        }
        Ok(())
    }

    fn create_task(&mut self, task_fn: TaskFn) -> Task {
        let id = self.next_task_id;
        self.next_task_id += 1;
        Task::create(id, task_fn, self.task_lisenters.clone())
    }

    fn add_worker(&mut self, task: Task, is_core: bool) {
        let new_worker = Worker::new(is_core, self.keep_alive_time, self.reciver.clone(), task);
        if is_core {
            self.core_workers.push(new_worker);
        } else {
            self.workers.push(new_worker);
        }
    }

    fn try_send_task(&mut self, task: Task) -> TPResult<()> {
        let sender = self.sender.as_ref();
        if sender.is_none() {
            return Err(TPError::DisconnectedError);
        }
        if let Err(err) = sender.unwrap().try_send(task) {
            return match err {
                TrySendError::Full(task) => self.process_task_if_channel_full(task),
                TrySendError::Disconnected(_) => Err(TPError::DisconnectedError),
            };
        }
        Ok(())
    }

    fn process_task_if_channel_full(&mut self, task: Task) -> TPResult<()> {
        // Attempt to find an idle worker.
        let idle_worker = self.workers.iter_mut().find(|worker| worker.is_finished());

        if let Some(idle_worker) = idle_worker {
            idle_worker.restart(task);
            return Ok(());
        }

        // If the vector of non-core worker is not full.
        if self.workers.len() < self.max_pool_size - self.core_pool_size {
            // Add a new non-core worker to process the task.
            self.add_worker(task, false);
            Ok(())
        } else {
            // Handle the task by RejectedTaskHandler.
            self.reject(task)
        }
    }

    fn reject(&self, task: Task) -> Result<(), TPError> {
        match &self.rejected_task_handler {
            RejectedTaskHandler::Abort => Err(TPError::AbortError),
            RejectedTaskHandler::CallerRuns => {
                task.run();
                Ok(())
            }
            RejectedTaskHandler::Discard => Ok(()),
        }
    }
}