Simplify ThreadedExecutor

Summary: The implementation doesn't need explicit locking or condition variables, everything can be done with a single concurrent queue.

Reviewed By: yfeldblum

Differential Revision: D25919429

fbshipit-source-id: a2b1e6332305174ba41d02fd79d9525a12b82218

folly/executors/ThreadedExecutor.cpp

@@ -17,65 +17,43 @@
 #include <folly/executors/ThreadedExecutor.h>
 
 #include <chrono>
+#include <utility>
 
 #include <glog/logging.h>
 
+#include <folly/ScopeGuard.h>
 #include <folly/executors/thread_factory/NamedThreadFactory.h>
 #include <folly/system/ThreadName.h>
 
 namespace folly {
 
-template <typename F>
-static auto with_unique_lock(std::mutex& m, F&& f) -> decltype(f()) {
-  std::unique_lock<std::mutex> lock(m);
-  return f();
-}
-
 ThreadedExecutor::ThreadedExecutor(std::shared_ptr<ThreadFactory> threadFactory)
-    : threadFactory_(std::move(threadFactory)) {
-  controlt_ = std::thread([this] { control(); });
-}
+    : threadFactory_(std::move(threadFactory)),
+      controlThread_([this] { control(); }) {}
 
 ThreadedExecutor::~ThreadedExecutor() {
   stopping_.store(true, std::memory_order_release);
-  notify();
-  controlt_.join();
+  controlMessages_.enqueue({Message::Type::StopControl, {}, {}});
+  controlThread_.join();
   CHECK(running_.empty());
-  CHECK(finished_.empty());
+  CHECK(controlMessages_.empty());
 }
 
 void ThreadedExecutor::add(Func func) {
   CHECK(!stopping_.load(std::memory_order_acquire));
-  with_unique_lock(enqueuedm_, [&] { enqueued_.push_back(std::move(func)); });
-  notify();
+  controlMessages_.enqueue({Message::Type::Start, std::move(func), {}});
 }
 
 std::shared_ptr<ThreadFactory> ThreadedExecutor::newDefaultThreadFactory() {
   return std::make_shared<NamedThreadFactory>("Threaded");
 }
 
-void ThreadedExecutor::notify() {
-  with_unique_lock(controlm_, [&] { controls_ = true; });
-  controlc_.notify_one();
-}
-
-void ThreadedExecutor::control() {
-  folly::setThreadName("ThreadedCtrl");
-  auto looping = true;
-  while (looping) {
-    controlWait();
-    looping = controlPerformAll();
-  }
-}
-
-void ThreadedExecutor::controlWait() {
-  constexpr auto kMaxWait = std::chrono::seconds(10);
-  std::unique_lock<std::mutex> lock(controlm_);
-  controlc_.wait_for(lock, kMaxWait, [&] { return controls_; });
-  controls_ = false;
-}
-
 void ThreadedExecutor::work(Func& func) {
+  SCOPE_EXIT {
+    controlMessages_.enqueue(
+        {Message::Type::Join, {}, std::this_thread::get_id()});
+  };
+
   try {
     func();
   } catch (const std::exception& e) {
@@ -84,35 +62,34 @@ void ThreadedExecutor::work(Func& func) {
   } catch (...) {
     LOG(ERROR) << "ThreadedExecutor: func threw unhandled non-exception object";
   }
-  auto id = std::this_thread::get_id();
-  with_unique_lock(finishedm_, [&] { finished_.push_back(id); });
-  notify();
 }
 
-void ThreadedExecutor::controlJoinFinishedThreads() {
-  std::deque<std::thread::id> finishedt;
-  with_unique_lock(finishedm_, [&] { std::swap(finishedt, finished_); });
-  for (auto id : finishedt) {
-    running_[id].join();
-    running_.erase(id);
-  }
-}
-
-void ThreadedExecutor::controlLaunchEnqueuedTasks() {
-  std::deque<Func> enqueuedt;
-  with_unique_lock(enqueuedm_, [&] { std::swap(enqueuedt, enqueued_); });
-  for (auto& f : enqueuedt) {
-    auto th = threadFactory_->newThread(
-        [this, f = std::move(f)]() mutable { work(f); });
-    auto id = th.get_id();
-    running_[id] = std::move(th);
+void ThreadedExecutor::control() {
+  folly::setThreadName("ThreadedCtrl");
+  bool controlStopping = false;
+  while (!(controlStopping && running_.empty())) {
+    auto msg = controlMessages_.dequeue();
+    switch (msg.type) {
+      case Message::Type::Start: {
+        auto th = threadFactory_->newThread(
+            [this, func = std::move(msg.startFunc)]() mutable { work(func); });
+        auto id = th.get_id();
+        running_[id] = std::move(th);
+        break;
+      }
+      case Message::Type::Join: {
+        auto it = running_.find(msg.joinTid);
+        CHECK(it != running_.end());
+        it->second.join();
+        running_.erase(it);
+        break;
+      }
+      case Message::Type::StopControl: {
+        CHECK(!std::exchange(controlStopping, true));
+        break;
+      }
+    }
   }
 }
 
-bool ThreadedExecutor::controlPerformAll() {
-  auto stopping = stopping_.load(std::memory_order_acquire);
-  controlJoinFinishedThreads();
-  controlLaunchEnqueuedTasks();
-  return !stopping || !running_.empty();
-}
 } // namespace folly

folly/executors/ThreadedExecutor.h

@@ -17,14 +17,12 @@
 #pragma once
 
 #include <atomic>
-#include <condition_variable>
-#include <deque>
-#include <map>
 #include <memory>
-#include <mutex>
 #include <thread>
 
 #include <folly/Executor.h>
+#include <folly/concurrency/UnboundedQueue.h>
+#include <folly/container/F14Map.h>
 #include <folly/executors/thread_factory/ThreadFactory.h>
 
 namespace folly {
@@ -67,33 +65,28 @@ class ThreadedExecutor : public virtual folly::Executor {
   void add(Func func) override;
 
  private:
-  static std::shared_ptr<ThreadFactory> newDefaultThreadFactory();
+  // TODO(ott): Switch to std::variant when available everywhere.
+  struct Message {
+    enum class Type { Start, Join, StopControl };
+    Type type;
+    Func startFunc;
+    std::thread::id joinTid;
+  };
 
-  void notify();
-  void control();
-  void controlWait();
-  bool controlPerformAll();
-  void controlJoinFinishedThreads();
-  void controlLaunchEnqueuedTasks();
+  static std::shared_ptr<ThreadFactory> newDefaultThreadFactory();
 
   void work(Func& func);
+  void control();
 
   std::shared_ptr<ThreadFactory> threadFactory_;
 
   std::atomic<bool> stopping_{false};
 
-  std::mutex controlm_;
-  std::condition_variable controlc_;
-  bool controls_ = false;
-  std::thread controlt_;
+  UMPSCQueue<Message, /* MayBlock */ true> controlMessages_;
+  std::thread controlThread_;
 
-  std::mutex enqueuedm_;
-  std::deque<Func> enqueued_;
-
-  //  Accessed only by the control thread, so no synchronization.
-  std::map<std::thread::id, std::thread> running_;
-
-  std::mutex finishedm_;
-  std::deque<std::thread::id> finished_;
+  // Accessed only by the control thread, so no synchronization.
+  F14FastMap<std::thread::id, std::thread> running_;
 };
+
 } // namespace folly