Extract all event bases from an IO thread pool

Summary:
This functionality can be used to simplify use cases which require finer
control over the distribution of tasks across threads.

Reviewed By: iahs

Differential Revision: D34158426

fbshipit-source-id: fdbf18fec61f20b78961d11c23d9ed1a6d322c74

folly/executors/IOThreadPoolExecutor.cpp

@@ -178,6 +178,19 @@ EventBase* IOThreadPoolExecutor::getEventBase() {
   return pickThread()->eventBase;
 }
 
+std::vector<Executor::KeepAlive<EventBase>>
+IOThreadPoolExecutor::getAllEventBases() {
+  ensureMaxActiveThreads();
+  std::vector<Executor::KeepAlive<EventBase>> evbs;
+  SharedMutex::ReadHolder r{&threadListLock_};
+  const auto& threads = threadList_.get();
+  evbs.reserve(threads.size());
+  for (const auto& thr : threads) {
+    evbs.emplace_back(static_cast<IOThread&>(*thr).eventBase);
+  }
+  return evbs;
+}
+
 EventBase* IOThreadPoolExecutor::getEventBase(
     ThreadPoolExecutor::ThreadHandle* h) {
   auto thread = dynamic_cast<IOThread*>(h);

folly/executors/IOThreadPoolExecutor.h

@@ -98,6 +98,10 @@ class IOThreadPoolExecutor : public ThreadPoolExecutor, public IOExecutor {
 
   folly::EventBase* getEventBase() override;
 
+  // Ensures that the maximum number of active threads is running and returns
+  // the EventBase associated with each thread.
+  std::vector<folly::Executor::KeepAlive<folly::EventBase>> getAllEventBases();
+
   static folly::EventBase* getEventBase(ThreadPoolExecutor::ThreadHandle*);
 
   static std::mutex* getEventBaseShutdownMutex(

folly/executors/ThreadPoolExecutor.cpp

@@ -420,10 +420,7 @@ void ThreadPoolExecutor::addObserver(std::shared_ptr<Observer> o) {
       o->threadPreviouslyStarted(thread.get());
     }
   }
-  while (activeThreads_.load(std::memory_order_relaxed) <
-         maxThreads_.load(std::memory_order_relaxed)) {
-    ensureActiveThreads();
-  }
+  ensureMaxActiveThreads();
 }
 
 void ThreadPoolExecutor::removeObserver(std::shared_ptr<Observer> o) {
@@ -523,6 +520,13 @@ void ThreadPoolExecutor::ensureActiveThreads() {
   activeThreads_.store(active + 1, std::memory_order_relaxed);
 }
 
+void ThreadPoolExecutor::ensureMaxActiveThreads() {
+  while (activeThreads_.load(std::memory_order_relaxed) <
+         maxThreads_.load(std::memory_order_relaxed)) {
+    ensureActiveThreads();
+  }
+}
+
 // If an idle thread times out, only join it if there are at least
 // minThreads threads.
 bool ThreadPoolExecutor::minActive() {

folly/executors/ThreadPoolExecutor.h

@@ -325,6 +325,7 @@ class ThreadPoolExecutor : public DefaultKeepAliveExecutor {
   folly::ThreadPoolListHook threadPoolHook_;
 
   // Dynamic thread sizing functions and variables
+  void ensureMaxActiveThreads();
   void ensureActiveThreads();
   void ensureJoined();
   bool minActive();