Add ThreadIdCollector to IOThreadPoolDeadlockDetectorObserver

Summary:
Add an implementation of `WorkerProvider` to `IOThreadPoolDeadlockDetectorObserver`.

- On thread start, the observer creates a `TaskStallDetector` and provides it a `ThreadIdCollector` containing the TID of the thread that started
- The thread will be blocked from exiting until any existing `WorkerProvider::KeepAlive`s are destroyed

Differential Revision: D31963855

fbshipit-source-id: cd749d76a433684a8280d95a4dbff441e32bfda9

folly/concurrency/DeadlockDetector.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <folly/Executor.h>
+#include <folly/executors/QueueObserver.h>
 
 namespace folly {
 class DeadlockDetector {
@@ -28,7 +29,9 @@ class DeadlockDetectorFactory {
  public:
   virtual ~DeadlockDetectorFactory() {}
   virtual std::unique_ptr<DeadlockDetector> create(
-      Executor* executor, const std::string& name) = 0;
+      Executor* executor,
+      const std::string& name,
+      std::unique_ptr<WorkerProvider> threadIdCollector) = 0;
   static DeadlockDetectorFactory* instance();
 };
 

folly/executors/IOThreadPoolDeadlockDetectorObserver.cpp

@@ -21,6 +21,35 @@
 
 namespace folly {
 
+namespace {
+class ThreadIdCollector : public WorkerProvider {
+ public:
+  ThreadIdCollector(pid_t tid, std::mutex& eventBaseShutdownMutex)
+      : tid_(tid), eventBaseShutdownMutex_(eventBaseShutdownMutex) {}
+
+  IdsWithKeepAlive collectThreadIds() override final {
+    std::unique_lock<std::mutex> guard(eventBaseShutdownMutex_);
+    return {
+        std::make_unique<WorkerKeepAlive>(std::move(guard)),
+        std::vector<pid_t>{tid_}};
+  }
+
+ private:
+  class WorkerKeepAlive : public WorkerProvider::KeepAlive {
+   public:
+    explicit WorkerKeepAlive(std::unique_lock<std::mutex> lock)
+        : lock_(std::move(lock)) {}
+    ~WorkerKeepAlive() override {}
+
+   private:
+    std::unique_lock<std::mutex> lock_;
+  };
+
+  pid_t tid_;
+  std::mutex& eventBaseShutdownMutex_;
+};
+} // namespace
+
 IOThreadPoolDeadlockDetectorObserver::IOThreadPoolDeadlockDetectorObserver(
     DeadlockDetectorFactory* deadlockDetectorFactory, const std::string& name)
     : name_(name), deadlockDetectorFactory_(deadlockDetectorFactory) {}
@@ -32,13 +61,19 @@ void IOThreadPoolDeadlockDetectorObserver::threadStarted(
   }
 
   auto eventBase = folly::IOThreadPoolExecutor::getEventBase(h);
+  auto eventBaseShutdownMutex =
+      folly::IOThreadPoolExecutor::getEventBaseShutdownMutex(h);
   // This Observer only works with IOThreadPoolExecutor class.
   CHECK_NOTNULL(eventBase);
+  CHECK_NOTNULL(eventBaseShutdownMutex);
   eventBase->runInEventBaseThread([=] {
-    auto thid = folly::to<std::string>(folly::getOSThreadID());
-    auto name = name_ + ":" + thid;
-    detectors_.wlock()->insert_or_assign(
-        h, deadlockDetectorFactory_->create(eventBase, name));
+    auto tid = folly::getOSThreadID();
+    auto name = name_ + ":" + folly::to<std::string>(tid);
+    auto deadlockDetector = deadlockDetectorFactory_->create(
+        eventBase,
+        name,
+        std::make_unique<ThreadIdCollector>(tid, *eventBaseShutdownMutex));
+    detectors_.wlock()->insert_or_assign(h, std::move(deadlockDetector));
   });
 }
 

folly/executors/IOThreadPoolExecutor.cpp

@@ -183,6 +183,17 @@ EventBase* IOThreadPoolExecutor::getEventBase(
   return nullptr;
 }
 
+std::mutex* IOThreadPoolExecutor::getEventBaseShutdownMutex(
+    ThreadPoolExecutor::ThreadHandle* h) {
+  auto thread = dynamic_cast<IOThread*>(h);
+
+  if (thread) {
+    return &thread->eventBaseShutdownMutex_;
+  }
+
+  return nullptr;
+}
+
 EventBaseManager* IOThreadPoolExecutor::getEventBaseManager() {
   return eventBaseManager_;
 }

folly/executors/IOThreadPoolExecutor.h

@@ -83,6 +83,9 @@ class IOThreadPoolExecutor : public ThreadPoolExecutor, public IOExecutor {
 
   static folly::EventBase* getEventBase(ThreadPoolExecutor::ThreadHandle*);
 
+  static std::mutex* getEventBaseShutdownMutex(
+      ThreadPoolExecutor::ThreadHandle*);
+
   folly::EventBaseManager* getEventBaseManager();
 
  private:

folly/executors/test/IOThreadPoolDeadlockDetectorObserverTest.cpp

@@ -40,11 +40,18 @@ class DeadlockDetectorMock : public DeadlockDetector {
 class DeadlockDetectorFactoryMock : public DeadlockDetectorFactory {
  public:
   std::unique_ptr<DeadlockDetector> create(
-      folly::Executor* executor, const std::string& name) override {
+      folly::Executor* executor,
+      const std::string& name,
+      std::unique_ptr<WorkerProvider> tidCollector) override {
     EXPECT_TRUE(executor != nullptr);
-    std::string expectedName =
-        fmt::format("TestPool:{}", folly::getOSThreadID());
+    auto tid = folly::getOSThreadID();
+    std::string expectedName = fmt::format("TestPool:{}", tid);
     EXPECT_EQ(expectedName, name);
+    {
+      auto tids = tidCollector->collectThreadIds();
+      EXPECT_EQ(1, tids.threadIds.size());
+      EXPECT_EQ(tid, tids.threadIds[0]);
+    }
     name_ = name;
     auto retval = std::make_unique<DeadlockDetectorMock>(counter_);
     baton.post();