Current thread executor

Summary: A thread-local query for the current executor. Grounding feature for logging operations.

Reviewed By: andriigrynenko

Differential Revision: D20070133

fbshipit-source-id: e320d2a3d97744b91c07a6060708e5cefc517cb7

folly/Executor.cpp

@@ -15,6 +15,7 @@
  */
 
 #include <folly/Executor.h>
+#include <folly/SingletonThreadLocal.h>
 
 #include <stdexcept>
 
@@ -35,4 +36,36 @@ void Executor::keepAliveRelease() {
   LOG(FATAL) << __func__ << "() should not be called for folly::Executor types "
              << "which do not override keepAliveAcquire()";
 }
+
+namespace {
+using BlockingContextSingletonT =
+    SingletonThreadLocal<folly::Optional<BlockingContext>>;
+} // namespace
+
+folly::Optional<BlockingContext> getBlockingContext() {
+  return BlockingContextSingletonT::get();
+}
+
+BlockingGuard::BlockingGuard(folly::StringPiece executorName)
+    : previousContext_{BlockingContextSingletonT::get()} {
+  BlockingContextSingletonT::get() = BlockingContext{executorName};
+}
+
+BlockingGuard::BlockingGuard()
+    : previousContext_{BlockingContextSingletonT::get()} {
+  BlockingContextSingletonT::get() = folly::none;
+}
+
+BlockingGuard::~BlockingGuard() {
+  BlockingContextSingletonT::get() = std::move(previousContext_);
+}
+
+BlockingGuard makeBlockingDisallowedGuard(folly::StringPiece executorName) {
+  return BlockingGuard{executorName};
+}
+
+BlockingGuard makeBlockingAllowedGuard() {
+  return BlockingGuard{};
+}
+
 } // namespace folly

folly/Executor.h

@@ -21,6 +21,8 @@
 #include <utility>
 
 #include <folly/Function.h>
+#include <folly/Optional.h>
+#include <folly/Range.h>
 #include <folly/Utility.h>
 
 namespace folly {
@@ -297,4 +299,32 @@ Executor::KeepAlive<ExecutorT> getKeepAliveToken(
   return ka.copy();
 }
 
+struct BlockingContext {
+  folly::StringPiece executorName;
+};
+
+class BlockingGuard;
+
+BlockingGuard makeBlockingDisallowedGuard(folly::StringPiece executorName);
+BlockingGuard makeBlockingAllowedGuard();
+
+class [[nodiscard]] BlockingGuard {
+ public:
+  ~BlockingGuard();
+
+ private:
+  // Disallow blocking
+  BlockingGuard(folly::StringPiece executorName);
+  // Empty guard treated as temporarily allowing blocking
+  BlockingGuard();
+
+  friend BlockingGuard makeBlockingDisallowedGuard(
+      folly::StringPiece executorName);
+  friend BlockingGuard makeBlockingAllowedGuard();
+
+  folly::Optional<BlockingContext> previousContext_;
+};
+
+folly::Optional<BlockingContext> getBlockingContext();
+
 } // namespace folly

folly/executors/CPUThreadPoolExecutor.cpp

@@ -37,6 +37,8 @@ namespace {
 using default_queue = UnboundedBlockingQueue<CPUThreadPoolExecutor::CPUTask>;
 using default_queue_alloc =
     AlignedSysAllocator<default_queue, FixedAlign<alignof(default_queue)>>;
+
+constexpr folly::StringPiece executorName = "CPUThreadPoolExecutor";
 } // namespace
 
 const size_t CPUThreadPoolExecutor::kDefaultMaxQueueSize = 1 << 14;
@@ -194,10 +196,12 @@ bool CPUThreadPoolExecutor::taskShouldStop(folly::Optional<CPUTask>& task) {
 
 void CPUThreadPoolExecutor::threadRun(ThreadPtr thread) {
   this->threadPoolHook_.registerThread();
+  auto guard = folly::makeBlockingDisallowedGuard(executorName);
 
   thread->startupBaton.post();
   while (true) {
     auto task = taskQueue_->try_take_for(threadTimeout_);
+
     // Handle thread stopping, either by task timeout, or
     // by 'poison' task added in join() or stop().
     if (UNLIKELY(!task || task.value().poison)) {

folly/executors/EDFThreadPoolExecutor.cpp

@@ -30,6 +30,9 @@
 #include <folly/executors/EDFThreadPoolExecutor.h>
 
 namespace folly {
+namespace {
+constexpr folly::StringPiece executorName = "EDFThreadPoolExecutor";
+}
 
 class EDFThreadPoolExecutor::Task {
  public:
@@ -340,6 +343,7 @@ folly::Executor::KeepAlive<> EDFThreadPoolExecutor::deadlineExecutor(
 
 void EDFThreadPoolExecutor::threadRun(ThreadPtr thread) {
   this->threadPoolHook_.registerThread();
+  auto guard = folly::makeBlockingDisallowedGuard(executorName);
 
   thread->startupBaton.post();
   for (;;) {

folly/executors/test/ThreadPoolExecutorTest.cpp

@@ -1002,3 +1002,47 @@ TEST(ThreadPoolExecutorTest, VirtualExecutorTestCPU) {
 TEST(ThreadPoolExecutorTest, VirtualExecutorTestEDF) {
   virtualExecutorTest<EDFThreadPoolExecutor>();
 }
+
+// Test use of guard inside executors
+template <class TPE>
+static void currentThreadTest(folly::StringPiece executorName) {
+  folly::Optional<BlockingContext> blockingContext{};
+  TPE tpe(1);
+  tpe.add([&blockingContext]() { blockingContext = getBlockingContext(); });
+  tpe.join();
+  EXPECT_EQ(blockingContext->executorName, executorName);
+}
+
+// Test the nesting of the allowing guard
+template <class TPE>
+static void currentThreadTestDisabled(folly::StringPiece executorName) {
+  folly::Optional<BlockingContext> blockingContextAllowed{};
+  folly::Optional<BlockingContext> blockingContextDisallowed{};
+  TPE tpe(1);
+  tpe.add([&]() {
+    {
+      // Nest the guard that reallows blocking
+      auto guard = folly::makeBlockingAllowedGuard();
+      blockingContextAllowed = getBlockingContext();
+    }
+    blockingContextDisallowed = getBlockingContext();
+  });
+  tpe.join();
+  EXPECT_TRUE(!blockingContextAllowed.has_value());
+  EXPECT_EQ(blockingContextDisallowed->executorName, executorName);
+}
+
+TEST(ThreadPoolExecutorTest, CPUCurrentThreadExecutor) {
+  currentThreadTest<CPUThreadPoolExecutor>("CPUThreadPoolExecutor");
+  currentThreadTestDisabled<CPUThreadPoolExecutor>("CPUThreadPoolExecutor");
+}
+
+TEST(ThreadPoolExecutorTest, IOCurrentThreadExecutor) {
+  currentThreadTest<IOThreadPoolExecutor>("EventBase");
+  currentThreadTestDisabled<IOThreadPoolExecutor>("EventBase");
+}
+
+TEST(ThreadPoolExecutorTest, EDFCurrentThreadExecutor) {
+  currentThreadTest<EDFThreadPoolExecutor>("EDFThreadPoolExecutor");
+  currentThreadTestDisabled<EDFThreadPoolExecutor>("EDFThreadPoolExecutor");
+}

folly/io/async/EventBase.cpp

@@ -36,6 +36,8 @@
 #include <folly/system/ThreadName.h>
 
 namespace {
+constexpr folly::StringPiece executorName = "EventBase";
+
 class EventBaseBackend : public folly::EventBaseBackendBase {
  public:
   EventBaseBackend();
@@ -309,6 +311,7 @@ void EventBase::waitUntilRunning() {
 
 // enters the event_base loop -- will only exit when forced to
 bool EventBase::loop() {
+  auto guard = folly::makeBlockingDisallowedGuard(executorName);
   return loopBody();
 }