folly/experimental/wangle/concurrent/ThreadPoolExecutor.h

   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 #pragma once
  18 #include <folly/wangle/Executor.h>
  19 #include <folly/experimental/wangle/concurrent/LifoSemMPMCQueue.h>
  20 #include <folly/experimental/wangle/concurrent/NamedThreadFactory.h>
  21 #include <folly/experimental/wangle/rx/Observable.h>
  22 #include <folly/Baton.h>
  23 #include <folly/Memory.h>
  24 #include <folly/RWSpinLock.h>
  25
  26 #include <algorithm>
  27 #include <mutex>
  28 #include <queue>
  29
  30 #include <glog/logging.h>
  31
  32 namespace folly { namespace wangle {
  33
  34 class ThreadPoolExecutor : public Executor {
  35  public:
  36   explicit ThreadPoolExecutor(
  37       size_t numThreads,
  38       std::shared_ptr<ThreadFactory> threadFactory);
  39
  40   ~ThreadPoolExecutor();
  41
  42   virtual void add(Func func) override = 0;
  43   virtual void add(
  44       Func func,
  45       std::chrono::milliseconds expiration,
  46       Func expireCallback) = 0;
  47
  48   void setThreadFactory(std::shared_ptr<ThreadFactory> threadFactory) {
  49     CHECK(numThreads() == 0);
  50     threadFactory_ = std::move(threadFactory);
  51   }
  52
  53   size_t numThreads();
  54   void setNumThreads(size_t numThreads);
  55   void stop();
  56   void join();
  57
  58   struct PoolStats {
  59     PoolStats() : threadCount(0), idleThreadCount(0), activeThreadCount(0),
  60                   pendingTaskCount(0), totalTaskCount(0) {}
  61     size_t threadCount, idleThreadCount, activeThreadCount;
  62     uint64_t pendingTaskCount, totalTaskCount;
  63   };
  64
  65   PoolStats getPoolStats();
  66
  67   struct TaskStats {
  68     TaskStats() : expired(false), waitTime(0), runTime(0) {}
  69     bool expired;
  70     std::chrono::nanoseconds waitTime;
  71     std::chrono::nanoseconds runTime;
  72   };
  73
  74   Subscription<TaskStats> subscribeToTaskStats(
  75       const ObserverPtr<TaskStats>& observer) {
  76     return taskStatsSubject_.subscribe(observer);
  77   }
  78
  79  protected:
  80   // Prerequisite: threadListLock_ writelocked
  81   void addThreads(size_t n);
  82   // Prerequisite: threadListLock_ writelocked
  83   void removeThreads(size_t n, bool isJoin);
  84
  85   struct FOLLY_ALIGN_TO_AVOID_FALSE_SHARING Thread {
  86     virtual ~Thread() {}
  87     Thread() : id(nextId++), handle(), idle(true) {};
  88     static std::atomic<uint64_t> nextId;
  89     uint64_t id;
  90     std::thread handle;
  91     bool idle;
  92     Baton<> startupBaton;
  93   };
  94
  95   typedef std::shared_ptr<Thread> ThreadPtr;
  96
  97   struct Task {
  98     explicit Task(
  99         Func&& func,
 100         std::chrono::milliseconds expiration,
 101         Func&& expireCallback);
 102     Func func_;
 103     TaskStats stats_;
 104     std::chrono::steady_clock::time_point enqueueTime_;
 105     std::chrono::milliseconds expiration_;
 106     Func expireCallback_;
 107   };
 108
 109   void runTask(const ThreadPtr& thread, Task&& task);
 110
 111   // The function that will be bound to pool threads. It must call
 112   // thread->startupBaton.post() when it's ready to consume work.
 113   virtual void threadRun(ThreadPtr thread) = 0;
 114
 115   // Stop n threads and put their ThreadPtrs in the threadsStopped_ queue
 116   // Prerequisite: threadListLock_ writelocked
 117   virtual void stopThreads(size_t n) = 0;
 118
 119   // Create a suitable Thread struct
 120   virtual ThreadPtr makeThread() {
 121     return std::make_shared<Thread>();
 122   }
 123
 124   // Prerequisite: threadListLock_ readlocked
 125   virtual uint64_t getPendingTaskCount() = 0;
 126
 127   class ThreadList {
 128    public:
 129     void add(const ThreadPtr& state) {
 130       auto it = std::lower_bound(vec_.begin(), vec_.end(), state, compare);
 131       vec_.insert(it, state);
 132     }
 133
 134     void remove(const ThreadPtr& state) {
 135       auto itPair = std::equal_range(vec_.begin(), vec_.end(), state, compare);
 136       CHECK(itPair.first != vec_.end());
 137       CHECK(std::next(itPair.first) == itPair.second);
 138       vec_.erase(itPair.first);
 139     }
 140
 141     const std::vector<ThreadPtr>& get() const {
 142       return vec_;
 143     }
 144
 145    private:
 146     static bool compare(const ThreadPtr& ts1, const ThreadPtr& ts2) {
 147       return ts1->id < ts2->id;
 148     }
 149
 150     std::vector<ThreadPtr> vec_;
 151   };
 152
 153   class StoppedThreadQueue : public BlockingQueue<ThreadPtr> {
 154    public:
 155     void add(ThreadPtr item) override;
 156     ThreadPtr take() override;
 157     size_t size() override;
 158
 159    private:
 160     LifoSem sem_;
 161     std::mutex mutex_;
 162     std::queue<ThreadPtr> queue_;
 163   };
 164
 165   std::shared_ptr<ThreadFactory> threadFactory_;
 166   ThreadList threadList_;
 167   RWSpinLock threadListLock_;
 168   StoppedThreadQueue stoppedThreads_;
 169   std::atomic<bool> isJoin_; // whether the current downsizing is a join
 170
 171   Subject<TaskStats> taskStatsSubject_;
 172 };
 173
 174 }} // folly::wangle