648db22b |
1 | /* |
2 | * Copyright (c) Meta Platforms, Inc. and affiliates. |
3 | * All rights reserved. |
4 | * |
5 | * This source code is licensed under both the BSD-style license (found in the |
6 | * LICENSE file in the root directory of this source tree) and the GPLv2 (found |
7 | * in the COPYING file in the root directory of this source tree). |
8 | */ |
9 | #pragma once |
10 | |
11 | #include "utils/WorkQueue.h" |
12 | |
13 | #include <cstddef> |
14 | #include <functional> |
15 | #include <thread> |
16 | #include <vector> |
17 | |
18 | namespace pzstd { |
19 | /// A simple thread pool that pulls tasks off its queue in FIFO order. |
20 | class ThreadPool { |
21 | std::vector<std::thread> threads_; |
22 | |
23 | WorkQueue<std::function<void()>> tasks_; |
24 | |
25 | public: |
26 | /// Constructs a thread pool with `numThreads` threads. |
27 | explicit ThreadPool(std::size_t numThreads) { |
28 | threads_.reserve(numThreads); |
29 | for (std::size_t i = 0; i < numThreads; ++i) { |
30 | threads_.emplace_back([this] { |
31 | std::function<void()> task; |
32 | while (tasks_.pop(task)) { |
33 | task(); |
34 | } |
35 | }); |
36 | } |
37 | } |
38 | |
39 | /// Finishes all tasks currently in the queue. |
40 | ~ThreadPool() { |
41 | tasks_.finish(); |
42 | for (auto& thread : threads_) { |
43 | thread.join(); |
44 | } |
45 | } |
46 | |
47 | /** |
48 | * Adds `task` to the queue of tasks to execute. Since `task` is a |
49 | * `std::function<>`, it cannot be a move only type. So any lambda passed must |
50 | * not capture move only types (like `std::unique_ptr`). |
51 | * |
52 | * @param task The task to execute. |
53 | */ |
54 | void add(std::function<void()> task) { |
55 | tasks_.push(std::move(task)); |
56 | } |
57 | }; |
58 | } |