Fix -Wsign-compare
[folly.git] / folly / experimental / wangle / concurrent / IOThreadPoolExecutor.cpp
1 /*
2  * Copyright 2014 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 #include <folly/experimental/wangle/concurrent/IOThreadPoolExecutor.h>
18
19 #include <folly/MoveWrapper.h>
20 #include <glog/logging.h>
21 #include <folly/io/async/EventBaseManager.h>
22
23 #include <folly/detail/MemoryIdler.h>
24
25 namespace folly { namespace wangle {
26
27 using folly::detail::MemoryIdler;
28
29 /* Class that will free jemalloc caches and madvise the stack away
30  * if the event loop is unused for some period of time
31  */
32 class MemoryIdlerTimeout
33     : public AsyncTimeout , public EventBase::LoopCallback {
34  public:
35   explicit MemoryIdlerTimeout(EventBase* b) : AsyncTimeout(b), base_(b) {}
36
37   virtual void timeoutExpired() noexcept {
38     idled = true;
39   }
40
41   virtual void runLoopCallback() noexcept {
42     if (idled) {
43       MemoryIdler::flushLocalMallocCaches();
44       MemoryIdler::unmapUnusedStack(MemoryIdler::kDefaultStackToRetain);
45
46       idled = false;
47     } else {
48       std::chrono::steady_clock::duration idleTimeout =
49         MemoryIdler::defaultIdleTimeout.load(
50           std::memory_order_acquire);
51
52       idleTimeout = MemoryIdler::getVariationTimeout(idleTimeout);
53
54       scheduleTimeout(std::chrono::duration_cast<std::chrono::milliseconds>(
55                         idleTimeout).count());
56     }
57
58     // reschedule this callback for the next event loop.
59     base_->runBeforeLoop(this);
60   }
61  private:
62   EventBase* base_;
63   bool idled{false};
64 } ;
65
66 IOThreadPoolExecutor::IOThreadPoolExecutor(
67     size_t numThreads,
68     std::shared_ptr<ThreadFactory> threadFactory)
69   : ThreadPoolExecutor(numThreads, std::move(threadFactory)),
70     nextThread_(0) {
71   addThreads(numThreads);
72   CHECK(threadList_.get().size() == numThreads);
73 }
74
75 IOThreadPoolExecutor::~IOThreadPoolExecutor() {
76   stop();
77 }
78
79 void IOThreadPoolExecutor::add(Func func) {
80   add(std::move(func), std::chrono::milliseconds(0));
81 }
82
83 void IOThreadPoolExecutor::add(
84     Func func,
85     std::chrono::milliseconds expiration,
86     Func expireCallback) {
87   RWSpinLock::ReadHolder{&threadListLock_};
88   if (threadList_.get().empty()) {
89     throw std::runtime_error("No threads available");
90   }
91   auto thread = threadList_.get()[nextThread_++ % threadList_.get().size()];
92   auto ioThread = std::static_pointer_cast<IOThread>(thread);
93
94   auto moveTask = folly::makeMoveWrapper(
95       Task(std::move(func), expiration, std::move(expireCallback)));
96   auto wrappedFunc = [this, ioThread, moveTask] () mutable {
97     runTask(ioThread, std::move(*moveTask));
98     ioThread->pendingTasks--;
99   };
100
101   ioThread->pendingTasks++;
102   if (!ioThread->eventBase->runInEventBaseThread(std::move(wrappedFunc))) {
103     ioThread->pendingTasks--;
104     throw std::runtime_error("Unable to run func in event base thread");
105   }
106 }
107
108 std::shared_ptr<ThreadPoolExecutor::Thread>
109 IOThreadPoolExecutor::makeThread() {
110   return std::make_shared<IOThread>();
111 }
112
113 void IOThreadPoolExecutor::threadRun(ThreadPtr thread) {
114   const auto ioThread = std::static_pointer_cast<IOThread>(thread);
115   ioThread->eventBase =
116     folly::EventBaseManager::get()->getEventBase();
117
118   auto idler = new MemoryIdlerTimeout(ioThread->eventBase);
119   ioThread->eventBase->runBeforeLoop(idler);
120
121   thread->startupBaton.post();
122   while (ioThread->shouldRun) {
123     ioThread->eventBase->loopForever();
124   }
125   if (isJoin_) {
126     while (ioThread->pendingTasks > 0) {
127       ioThread->eventBase->loopOnce();
128     }
129   }
130   stoppedThreads_.add(ioThread);
131 }
132
133 // threadListLock_ is writelocked
134 void IOThreadPoolExecutor::stopThreads(size_t n) {
135   for (size_t i = 0; i < n; i++) {
136     const auto ioThread = std::static_pointer_cast<IOThread>(
137         threadList_.get()[i]);
138     ioThread->shouldRun = false;
139     ioThread->eventBase->terminateLoopSoon();
140   }
141 }
142
143 // threadListLock_ is readlocked
144 uint64_t IOThreadPoolExecutor::getPendingTaskCount() {
145   uint64_t count = 0;
146   for (const auto& thread : threadList_.get()) {
147     auto ioThread = std::static_pointer_cast<IOThread>(thread);
148     size_t pendingTasks = ioThread->pendingTasks;
149     if (pendingTasks > 0 && !ioThread->idle) {
150       pendingTasks--;
151     }
152     count += pendingTasks;
153   }
154   return count;
155 }
156
157 }} // folly::wangle