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. ThreadPool(std::size_t numThreads)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. ~ThreadPool()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 */ add(std::function<void ()> task)54 void add(std::function<void()> task) { 55 tasks_.push(std::move(task)); 56 } 57 }; 58 } 59