Sunshine/src/thread_pool.h

/**
 * @file src/thread_pool.h
 * @brief Declarations for the thread pool system.
 */
#pragma once

// standard includes
#include <thread>

// local includes
#include "platform/common.h"
#include "task_pool.h"

namespace thread_pool_util {
  /**
   * Allow threads to execute unhindered while keeping full control over the threads.
   */
  class ThreadPool: public task_pool_util::TaskPool {
  public:
    /**
     * @brief Callable unit executed by a worker thread.
     */
    typedef TaskPool::__task __task;

  private:
    std::vector<std::thread> _thread;

    std::condition_variable _cv;
    std::mutex _lock;

    bool _continue;

  public:
    ThreadPool():
        _continue {false} {
    }

    /**
     * @brief Start a pool with the requested number of worker threads.
     *
     * @param threads Number of worker threads to start.
     */
    explicit ThreadPool(int threads):
        _thread(threads),
        _continue {true} {
      for (auto &t : _thread) {
        t = std::thread(&ThreadPool::_main, this);
      }
    }

    ~ThreadPool() noexcept {
      if (!_continue) {
        return;
      }

      stop();
      join();
    }

    /**
     * @brief Queue work for asynchronous execution.
     *
     * @param newTask New task.
     * @param args Arguments forwarded to the callable or parser.
     * @return Identifier assigned to the queued task.
     */
    template<class Function, class... Args>
    auto push(Function &&newTask, Args &&...args) {
      std::lock_guard lg(_lock);
      auto future = TaskPool::push(std::forward<Function>(newTask), std::forward<Args>(args)...);

      _cv.notify_one();
      return future;
    }

    /**
     * @brief Queue a task that becomes runnable after a delay.
     *
     * @param task Task object to enqueue or execute.
     */
    void pushDelayed(std::pair<__time_point, __task> &&task) {
      std::lock_guard lg(_lock);

      TaskPool::pushDelayed(std::move(task));
    }

    /**
     * @brief Queue a task that becomes runnable after a delay.
     *
     * @param newTask New task.
     * @param duration Delay before the timed task should run.
     * @param args Arguments forwarded to the callable or parser.
     * @return Future that becomes ready after the delayed task completes.
     */
    template<class Function, class X, class Y, class... Args>
    auto pushDelayed(Function &&newTask, std::chrono::duration<X, Y> duration, Args &&...args) {
      std::lock_guard lg(_lock);
      auto future = TaskPool::pushDelayed(std::forward<Function>(newTask), duration, std::forward<Args>(args)...);

      // Update all timers for wait_until
      _cv.notify_all();
      return future;
    }

    /**
     * @brief Start worker threads for queued thread-pool tasks.
     *
     * @param threads Number of worker threads to start.
     */
    void start(int threads) {
      _continue = true;

      _thread.resize(threads);

      for (auto &t : _thread) {
        t = std::thread(&ThreadPool::_main, this);
      }
    }

    /**
     * @brief Stop worker threads and prevent additional task execution.
     */
    void stop() {
      std::lock_guard lg(_lock);

      _continue = false;
      _cv.notify_all();
    }

    /**
     * @brief Wait for worker threads owned by the session to exit.
     */
    void join() {
      for (auto &t : _thread) {
        t.join();
      }
    }

  public:
    /**
     * @brief Run the main application or worker loop.
     */
    void _main() {
      platf::set_thread_name("TaskPool::worker");
      while (_continue) {
        if (auto task = this->pop()) {
          (*task)->run();
        } else {
          std::unique_lock uniq_lock(_lock);

          if (ready()) {
            continue;
          }

          if (!_continue) {
            break;
          }

          if (auto tp = next()) {
            _cv.wait_until(uniq_lock, *tp);
          } else {
            _cv.wait(uniq_lock);
          }
        }
      }

      // Execute remaining tasks
      while (auto task = this->pop()) {
        (*task)->run();
      }
    }
  };
}  // namespace thread_pool_util