folly/executors: properly synchronize access to thread pool stats.

Summary:
ThreadPoolExecutor::getPoolStats depends on derived classes implementation.
Calling `getPoolStats` is specially useful from `ThreadPoolExecutor::withAll`.

This change moves registration of thread pool executor instances to derived classes
to avoid `ThreadPoolExecutor::withAll` racing with thread pool destruction while calling `getPoolStats`.

Reviewed By: yfeldblum

Differential Revision: D17965135

fbshipit-source-id: 637de55ed78b085dc86fffa0e0f4a66cd4cfcede

folly/executors/CPUThreadPoolExecutor.cpp

@@ -40,6 +40,7 @@ CPUThreadPoolExecutor::CPUThreadPoolExecutor(
           std::move(threadFactory)),
       taskQueue_(std::move(taskQueue)) {
   setNumThreads(numThreads);
+  registerThreadPoolExecutor(this);
 }
 
 CPUThreadPoolExecutor::CPUThreadPoolExecutor(
@@ -52,6 +53,7 @@ CPUThreadPoolExecutor::CPUThreadPoolExecutor(
           std::move(threadFactory)),
       taskQueue_(std::move(taskQueue)) {
   setNumThreads(numThreads.first);
+  registerThreadPoolExecutor(this);
 }
 
 CPUThreadPoolExecutor::CPUThreadPoolExecutor(
@@ -98,6 +100,7 @@ CPUThreadPoolExecutor::CPUThreadPoolExecutor(
           std::move(threadFactory)) {}
 
 CPUThreadPoolExecutor::~CPUThreadPoolExecutor() {
+  deregisterThreadPoolExecutor(this);
   stop();
   CHECK(threadsToStop_ == 0);
 }

folly/executors/CPUThreadPoolExecutor.h

@@ -141,7 +141,7 @@ class CPUThreadPoolExecutor : public ThreadPoolExecutor {
  private:
   void threadRun(ThreadPtr thread) override;
   void stopThreads(size_t n) override;
-  size_t getPendingTaskCountImpl() const override;
+  size_t getPendingTaskCountImpl() const override final;
 
   bool tryDecrToStop();
   bool taskShouldStop(folly::Optional<CPUTask>&);

folly/executors/EDFThreadPoolExecutor.cpp

@@ -248,9 +248,11 @@ EDFThreadPoolExecutor::EDFThreadPoolExecutor(
     : ThreadPoolExecutor(numThreads, numThreads, std::move(threadFactory)),
       taskQueue_(std::make_unique<TaskQueue>()) {
   setNumThreads(numThreads);
+  registerThreadPoolExecutor(this);
 }
 
 EDFThreadPoolExecutor::~EDFThreadPoolExecutor() {
+  deregisterThreadPoolExecutor(this);
   stop();
 }
 

folly/executors/EDFThreadPoolExecutor.h

@@ -60,7 +60,7 @@ class EDFThreadPoolExecutor : public SoftRealTimeExecutor,
  protected:
   void threadRun(ThreadPtr thread) override;
   void stopThreads(std::size_t numThreads) override;
-  std::size_t getPendingTaskCountImpl() const override;
+  std::size_t getPendingTaskCountImpl() const override final;
 
  private:
   bool shouldStop();

folly/executors/IOThreadPoolExecutor.cpp

@@ -80,9 +80,11 @@ IOThreadPoolExecutor::IOThreadPoolExecutor(
       nextThread_(0),
       eventBaseManager_(ebm) {
   setNumThreads(numThreads);
+  registerThreadPoolExecutor(this);
 }
 
 IOThreadPoolExecutor::~IOThreadPoolExecutor() {
+  deregisterThreadPoolExecutor(this);
   stop();
 }
 

folly/executors/IOThreadPoolExecutor.h

@@ -86,7 +86,7 @@ class IOThreadPoolExecutor : public ThreadPoolExecutor, public IOExecutor {
   std::shared_ptr<IOThread> pickThread();
   void threadRun(ThreadPtr thread) override;
   void stopThreads(size_t n) override;
-  size_t getPendingTaskCountImpl() const override;
+  size_t getPendingTaskCountImpl() const override final;
 
   std::atomic<size_t> nextThread_;
   folly::ThreadLocal<std::shared_ptr<IOThread>> thisThread_;

folly/executors/ThreadPoolExecutor.cpp

@@ -29,6 +29,16 @@ SyncVecThreadPoolExecutors& getSyncVecThreadPoolExecutors() {
   return *storage;
 }
 
+void ThreadPoolExecutor::registerThreadPoolExecutor(ThreadPoolExecutor* tpe) {
+  getSyncVecThreadPoolExecutors().wlock()->push_back(tpe);
+}
+
+void ThreadPoolExecutor::deregisterThreadPoolExecutor(ThreadPoolExecutor* tpe) {
+  getSyncVecThreadPoolExecutors().withWLock([tpe](auto& tpes) {
+    tpes.erase(std::remove(tpes.begin(), tpes.end(), tpe), tpes.end());
+  });
+}
+
 DEFINE_int64(
     threadtimeout_ms,
     60000,
@@ -44,16 +54,11 @@ ThreadPoolExecutor::ThreadPoolExecutor(
       taskStatsCallbacks_(std::make_shared<TaskStatsCallbackRegistry>()),
       threadPoolHook_("folly::ThreadPoolExecutor"),
       minThreads_(minThreads),
-      threadTimeout_(FLAGS_threadtimeout_ms) {
-  getSyncVecThreadPoolExecutors().wlock()->push_back(this);
-}
+      threadTimeout_(FLAGS_threadtimeout_ms) {}
 
 ThreadPoolExecutor::~ThreadPoolExecutor() {
   joinKeepAliveOnce();
   CHECK_EQ(0, threadList_.get().size());
-  getSyncVecThreadPoolExecutors().withWLock([this](auto& tpe) {
-    tpe.erase(std::remove(tpe.begin(), tpe.end(), this), tpe.end());
-  });
 }
 
 ThreadPoolExecutor::Task::Task(

folly/executors/ThreadPoolExecutor.h

@@ -53,6 +53,15 @@ namespace folly {
  * ensureJoined() is called on add(), such that we can join idle
  * threads that were destroyed (which can't be joined from
  * themselves).
+ *
+ * Thread pool stats accounting:
+ *
+ * Derived classes must register instances to keep stats on all thread
+ * pools by calling registerThreadPoolExecutor(this) on constructions
+ * and deregisterThreadPoolExecutor(this) on destruction.
+ *
+ * Registration must be done wherever getPendingTaskCountImpl is implemented
+ * and getPendingTaskCountImpl should be marked 'final' to avoid data races.
  */
 class ThreadPoolExecutor : public DefaultKeepAliveExecutor {
  public:
@@ -224,10 +233,11 @@ class ThreadPoolExecutor : public DefaultKeepAliveExecutor {
     return std::make_shared<Thread>(this);
   }
 
+  static void registerThreadPoolExecutor(ThreadPoolExecutor* tpe);
+  static void deregisterThreadPoolExecutor(ThreadPoolExecutor* tpe);
+
   // Prerequisite: threadListLock_ readlocked or writelocked
-  virtual size_t getPendingTaskCountImpl() const {
-    throw std::logic_error("getPendingTaskCountImpl not implemented");
-  }
+  virtual size_t getPendingTaskCountImpl() const = 0;
 
   class ThreadList {
    public: