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