Fix ExecutorLoopController

Summary: D37231232 introduced the assumption that `runLoop` is always called from the same thread, but this may not be true for arbitrary executors. So we need to reset `loopThread_` at the end of the loop.

Reviewed By: andriigrynenko

Differential Revision: D52566640

fbshipit-source-id: 4c3849f19a262072ee7f87335527aada277a6249

folly/fibers/ExecutorLoopController-inl.h

@@ -16,17 +16,12 @@
 
 #pragma once
 #include <thread>
+#include <folly/ScopeGuard.h>
 #include <folly/futures/Future.h>
 
 namespace folly {
 namespace fibers {
 
-inline void ExecutorLoopController::assumeCalledFromLoopThread() {
-  DCHECK(
-      loopThread_ == std::thread::id{} ||
-      loopThread_ == std::this_thread::get_id());
-}
-
 inline ExecutorLoopController::ExecutorLoopController(folly::Executor* executor)
     : executor_(executor),
       timeoutManager_(executor_),
@@ -39,7 +34,6 @@ inline void ExecutorLoopController::setFiberManager(fibers::FiberManager* fm) {
 }
 
 inline void ExecutorLoopController::schedule() {
-  assumeCalledFromLoopThread();
   // add() is thread-safe, so this isn't properly optimized for addTask()
   if (!executorKeepAlive_) {
     executorKeepAlive_ = getKeepAliveToken(executor_);
@@ -57,9 +51,12 @@ inline void ExecutorLoopController::schedule() {
 }
 
 inline void ExecutorLoopController::runLoop() {
-  assumeCalledFromLoopThread();
-  // We are assuming that Executor will be always pinned to a thread.
-  loopThread_ = std::this_thread::get_id();
+  auto oldLoopThread = loopThread_.exchange(std::this_thread::get_id());
+  DCHECK(
+      oldLoopThread == std::thread::id{} ||
+      oldLoopThread == std::this_thread::get_id());
+  SCOPE_EXIT { loopThread_ = oldLoopThread; };
+
   if (!executorKeepAlive_) {
     if (!fm_->hasTasks()) {
       return;
@@ -73,9 +70,7 @@ inline void ExecutorLoopController::runLoop() {
 }
 
 inline void ExecutorLoopController::runEagerFiber(Fiber* fiber) {
-  assumeCalledFromLoopThread();
-  // We are assuming that Executor will be always pinned to a thread.
-  loopThread_ = std::this_thread::get_id();
+  DCHECK(loopThread_ == std::this_thread::get_id());
   if (!executorKeepAlive_) {
     executorKeepAlive_ = getKeepAliveToken(executor_);
   }

folly/fibers/ExecutorLoopController.h

@@ -142,8 +142,6 @@ class ExecutorLoopController : public fibers::ExecutorBasedLoopController {
     return loopThread_ == std::this_thread::get_id();
   }
 
-  void assumeCalledFromLoopThread();
-
   friend class fibers::FiberManager;
 };