| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| #pragma once |
|
|
| #include <cassert> |
| #include <condition_variable> |
| #include <functional> |
| #include <future> |
| #include <memory> |
| #include <mutex> |
| #include <queue> |
| #include <thread> |
| #include <unordered_map> |
|
|
| namespace PaddlePool { |
|
|
| class ThreadPool { |
| public: |
| using MutexGuard = std::lock_guard<std::mutex>; |
| using UniqueLock = std::unique_lock<std::mutex>; |
| using Thread = std::thread; |
| using ThreadID = std::thread::id; |
| using Task = std::function<void()>; |
|
|
| ThreadPool(); |
| explicit ThreadPool(size_t maxThreads); |
|
|
| ThreadPool(const ThreadPool &) = delete; |
| ThreadPool &operator=(const ThreadPool &) = delete; |
|
|
| ~ThreadPool(); |
|
|
| template <typename Func, typename... Ts> |
| auto submit(Func &&func, Ts &&...params) |
| -> std::future<typename std::result_of<Func(Ts...)>::type>; |
|
|
| size_t threadsNum() const; |
|
|
| private: |
| static constexpr size_t WAIT_SECONDS = 2; |
| void worker(); |
| void joinFinishedThreads(); |
|
|
| bool quit_; |
| size_t currentThreads_; |
| size_t idleThreads_; |
| size_t maxThreads_; |
|
|
| mutable std::mutex mutex_; |
| std::condition_variable cv_; |
| std::queue<Task> tasks_; |
| std::queue<ThreadID> finishedThreadIDs_; |
| std::unordered_map<ThreadID, Thread> threads_; |
| }; |
|
|
| } |
|
|
| namespace PaddlePool { |
|
|
| template <typename Func, typename... Ts> |
| auto ThreadPool::submit(Func &&func, Ts &&...params) |
| -> std::future<typename std::result_of<Func(Ts...)>::type> { |
| auto execute = |
| std::bind(std::forward<Func>(func), std::forward<Ts>(params)...); |
|
|
| using ReturnType = typename std::result_of<Func(Ts...)>::type; |
| using PackagedTask = std::packaged_task<ReturnType()>; |
|
|
| auto task = std::make_shared<PackagedTask>(std::move(execute)); |
| auto result = task->get_future(); |
|
|
| MutexGuard guard(mutex_); |
| assert(!quit_); |
|
|
| tasks_.emplace([task]() { (*task)(); }); |
| if (idleThreads_ > 0) { |
| cv_.notify_one(); |
| } else if (currentThreads_ < maxThreads_) { |
| Thread t(&ThreadPool::worker, this); |
| assert(threads_.find(t.get_id()) == threads_.end()); |
| threads_[t.get_id()] = std::move(t); |
| ++currentThreads_; |
| } |
|
|
| return result; |
| } |
|
|
| } |
|
|