From b8520e0d5b02e58f0a23d024c7a439b269fcecf1 Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Fri, 8 Jan 2016 17:15:21 -0800
Subject: [PATCH] Fix Python FiberTask to properly init/reset Python thread
 state

Reviewed By: alikhtarov

Differential Revision: D2813735

fb-gh-sync-id: ab3c3e18618ed8bf15f478bcfca008786834e65c
---
 folly/experimental/fibers/Fiber.cpp        | 30 ++++++++++++++--------
 folly/experimental/fibers/FiberManager.cpp |  4 +++
 folly/experimental/fibers/FiberManager.h   | 20 +++++++++++++++
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/folly/experimental/fibers/Fiber.cpp b/folly/experimental/fibers/Fiber.cpp
index d5a8b1e9..17502153 100644
--- a/folly/experimental/fibers/Fiber.cpp
+++ b/folly/experimental/fibers/Fiber.cpp
@@ -182,20 +182,30 @@ void Fiber::fiberFunc() {
 }
 
 intptr_t Fiber::preempt(State state) {
-  DCHECK_EQ(fiberManager_.activeFiber_, this);
-  DCHECK_EQ(state_, RUNNING);
-  DCHECK_NE(state, RUNNING);
+  intptr_t ret;
 
-  fiberManager_.activeFiber_ = nullptr;
-  state_ = state;
+  auto preemptImpl = [&]() mutable {
+    DCHECK_EQ(fiberManager_.activeFiber_, this);
+    DCHECK_EQ(state_, RUNNING);
+    DCHECK_NE(state, RUNNING);
 
-  recordStackPosition();
+    fiberManager_.activeFiber_ = nullptr;
+    state_ = state;
+
+    recordStackPosition();
+
+    ret = jumpContext(&fcontext_, &fiberManager_.mainContext_, 0);
 
-  auto ret = jumpContext(&fcontext_, &fiberManager_.mainContext_, 0);
+    DCHECK_EQ(fiberManager_.activeFiber_, this);
+    DCHECK_EQ(state_, READY_TO_RUN);
+    state_ = RUNNING;
+  };
 
-  DCHECK_EQ(fiberManager_.activeFiber_, this);
-  DCHECK_EQ(state_, READY_TO_RUN);
-  state_ = RUNNING;
+  if (fiberManager_.preemptRunner_) {
+    fiberManager_.preemptRunner_->run(std::ref(preemptImpl));
+  } else {
+    preemptImpl();
+  }
 
   return ret;
 }
diff --git a/folly/experimental/fibers/FiberManager.cpp b/folly/experimental/fibers/FiberManager.cpp
index e405630d..2a9a2ea1 100644
--- a/folly/experimental/fibers/FiberManager.cpp
+++ b/folly/experimental/fibers/FiberManager.cpp
@@ -121,6 +121,10 @@ void FiberManager::setObserver(ExecutionObserver* observer) {
   observer_ = observer;
 }
 
+void FiberManager::setPreemptRunner(InlineFunctionRunner* preemptRunner) {
+  preemptRunner_ = preemptRunner;
+}
+
 void FiberManager::doFibersPoolResizing() {
   while (fibersAllocated_ > maxFibersActiveLastPeriod_ &&
          fibersPoolSize_ > options_.maxFibersPoolSize) {
diff --git a/folly/experimental/fibers/FiberManager.h b/folly/experimental/fibers/FiberManager.h
index 0a74500f..549d3f32 100644
--- a/folly/experimental/fibers/FiberManager.h
+++ b/folly/experimental/fibers/FiberManager.h
@@ -48,6 +48,16 @@ template <typename T>
 class LocalType {
 };
 
+class InlineFunctionRunner {
+ public:
+  virtual ~InlineFunctionRunner() {}
+
+  /**
+   * func must be executed inline and only once.
+   */
+  virtual void run(std::function<void()> func) = 0;
+};
+
 /**
  * @class FiberManager
  * @brief Single-threaded task execution engine.
@@ -260,6 +270,11 @@ class FiberManager : public ::folly::Executor {
    */
   void setObserver(ExecutionObserver* observer);
 
+  /**
+   * Setup fibers preempt runner.
+   */
+  void setPreemptRunner(InlineFunctionRunner* preemptRunner);
+
   /**
    * Returns an estimate of the number of fibers which are waiting to run (does
    * not include fibers or tasks scheduled remotely).
@@ -387,6 +402,11 @@ class FiberManager : public ::folly::Executor {
    */
   std::function<void()> immediateFunc_;
 
+  /**
+   * Preempt runner.
+   */
+  InlineFunctionRunner* preemptRunner_{nullptr};
+
   /**
    * Fiber's execution observer.
    */
-- 
2.34.1