Add ctor to CPUThreadPoolExecutor to enable custom queue sizes
[folly.git] / folly / wangle / concurrent / CPUThreadPoolExecutor.h
1 /*
2  * Copyright 2015 Facebook, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *   http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16
17 #pragma once
18
19 #include <folly/wangle/concurrent/ThreadPoolExecutor.h>
20
21 namespace folly { namespace wangle {
22
23 class CPUThreadPoolExecutor : public ThreadPoolExecutor {
24  public:
25   struct CPUTask;
26
27   explicit CPUThreadPoolExecutor(
28       size_t numThreads,
29       std::unique_ptr<BlockingQueue<CPUTask>> taskQueue,
30       std::shared_ptr<ThreadFactory> threadFactory =
31           std::make_shared<NamedThreadFactory>("CPUThreadPool"));
32
33   explicit CPUThreadPoolExecutor(size_t numThreads);
34
35   explicit CPUThreadPoolExecutor(
36       size_t numThreads,
37       std::shared_ptr<ThreadFactory> threadFactory);
38
39   explicit CPUThreadPoolExecutor(
40       size_t numThreads,
41       uint32_t numPriorities,
42       std::shared_ptr<ThreadFactory> threadFactory =
43           std::make_shared<NamedThreadFactory>("CPUThreadPool"));
44
45   explicit CPUThreadPoolExecutor(
46       size_t numThreads,
47       uint32_t numPriorities,
48       size_t maxQueueSize,
49       std::shared_ptr<ThreadFactory> threadFactory =
50           std::make_shared<NamedThreadFactory>("CPUThreadPool"));
51
52   ~CPUThreadPoolExecutor();
53
54   void add(Func func) override;
55   void add(
56       Func func,
57       std::chrono::milliseconds expiration,
58       Func expireCallback = nullptr) override;
59
60   void add(Func func, uint32_t priority);
61   void add(
62       Func func,
63       uint32_t priority,
64       std::chrono::milliseconds expiration,
65       Func expireCallback = nullptr);
66
67   uint32_t getNumPriorities() const;
68
69   struct CPUTask : public ThreadPoolExecutor::Task {
70     // Must be noexcept move constructible so it can be used in MPMCQueue
71     explicit CPUTask(
72         Func&& f,
73         std::chrono::milliseconds expiration,
74         Func&& expireCallback)
75       : Task(std::move(f), expiration, std::move(expireCallback)),
76         poison(false) {}
77     CPUTask()
78       : Task(nullptr, std::chrono::milliseconds(0), nullptr),
79         poison(true) {}
80     CPUTask(CPUTask&& o) noexcept : Task(std::move(o)), poison(o.poison) {}
81     CPUTask(const CPUTask&) = default;
82     CPUTask& operator=(const CPUTask&) = default;
83     bool poison;
84   };
85
86   static const size_t kDefaultMaxQueueSize;
87   static const size_t kDefaultNumPriorities;
88
89  protected:
90   BlockingQueue<CPUTask>* getTaskQueue();
91
92  private:
93   void threadRun(ThreadPtr thread) override;
94   void stopThreads(size_t n) override;
95   uint64_t getPendingTaskCount() override;
96
97   std::unique_ptr<BlockingQueue<CPUTask>> taskQueue_;
98   std::atomic<ssize_t> threadsToStop_{0};
99 };
100
101 }} // folly::wangle