deps/libchdr/deps/zstd-1.5.5/contrib/pzstd/utils/ThreadPool.h

   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 }