File size: 2,729 Bytes
8207382 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | // Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#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
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;
}
} // namespace PaddlePool
|