RequestContext support

Summary: Integrating RequestContext into fibers library. RequestContext is saved for every task, when that task is created. If RequestContext is overriden when a task is being run (within fiber, in runInMainContext, within function passed to await call) the new context will continue to be used for the task.

Reviewed By: @alikhtarov

Differential Revision: D2240152

folly/experimental/fibers/Fiber.h

@@ -24,6 +24,7 @@
 #include <folly/CPortability.h>
 #include <folly/IntrusiveList.h>
 #include <folly/experimental/fibers/BoostContextCompatibility.h>
+#include <folly/io/async/Request.h>
 #include <folly/Portability.h>
 
 namespace folly { namespace fibers {
@@ -102,6 +103,7 @@ class Fiber {
   FiberManager& fiberManager_;  /**< Associated FiberManager */
   FContext fcontext_;           /**< current task execution context */
   intptr_t data_;               /**< Used to keep some data with the Fiber */
+  std::shared_ptr<RequestContext> rcontext_; /**< current RequestContext */
   std::function<void()> func_;  /**< task function */
   bool recordStackUsed_{false};
   bool stackFilledWithMagic_{false};

folly/experimental/fibers/FiberManager-inl.h

@@ -66,6 +66,7 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
   assert(fiber->state_ == Fiber::NOT_STARTED ||
          fiber->state_ == Fiber::READY_TO_RUN);
   currentFiber_ = fiber;
+  fiber->rcontext_ = RequestContext::setContext(std::move(fiber->rcontext_));
   if (observer_) {
     observer_->starting(reinterpret_cast<uintptr_t>(fiber));
   }
@@ -92,6 +93,7 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
       observer_->stopped(reinterpret_cast<uintptr_t>(fiber));
     }
     currentFiber_ = nullptr;
+    fiber->rcontext_ = RequestContext::setContext(std::move(fiber->rcontext_));
   } else if (fiber->state_ == Fiber::INVALID) {
     assert(fibersActive_ > 0);
     --fibersActive_;
@@ -113,7 +115,9 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
       observer_->stopped(reinterpret_cast<uintptr_t>(fiber));
     }
     currentFiber_ = nullptr;
+    fiber->rcontext_ = RequestContext::setContext(std::move(fiber->rcontext_));
     fiber->localData_.reset();
+    fiber->rcontext_.reset();
 
     if (fibersPoolSize_ < options_.maxFibersPoolSize) {
       fibersPool_.push_front(*fiber);
@@ -128,6 +132,7 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
       observer_->stopped(reinterpret_cast<uintptr_t>(fiber));
     }
     currentFiber_ = nullptr;
+    fiber->rcontext_ = RequestContext::setContext(std::move(fiber->rcontext_));
     fiber->state_ = Fiber::READY_TO_RUN;
     yieldedFibers_.push_back(*fiber);
   }
@@ -168,6 +173,7 @@ inline bool FiberManager::loopUntilNoReady() {
         if (task->localData) {
           fiber->localData_ = *task->localData;
         }
+        fiber->rcontext_ = std::move(task->rcontext);
 
         fiber->setFunction(std::move(task->func));
         fiber->data_ = reinterpret_cast<intptr_t>(fiber);
@@ -471,6 +477,7 @@ inline void FiberManager::initLocalData(Fiber& fiber) {
   if (fm && fm->currentFiber_ && fm->localType_ == localType_) {
     fiber.localData_ = fm->currentFiber_->localData_;
   }
+  fiber.rcontext_ = RequestContext::saveContext();
 }
 
 template <typename LocalT>

folly/experimental/fibers/FiberManager.h

@@ -27,6 +27,7 @@
 #include <folly/Executor.h>
 #include <folly/Likely.h>
 #include <folly/IntrusiveList.h>
+#include <folly/io/async/Request.h>
 #include <folly/futures/Try.h>
 
 #include <folly/experimental/ExecutionObserver.h>
@@ -252,13 +253,17 @@ class FiberManager : public ::folly::Executor {
 
   struct RemoteTask {
     template <typename F>
-    explicit RemoteTask(F&& f) : func(std::forward<F>(f)) {}
+    explicit RemoteTask(F&& f) :
+        func(std::forward<F>(f)),
+        rcontext(RequestContext::saveContext()) {}
     template <typename F>
     RemoteTask(F&& f, const Fiber::LocalData& localData_) :
         func(std::forward<F>(f)),
-        localData(folly::make_unique<Fiber::LocalData>(localData_)) {}
+        localData(folly::make_unique<Fiber::LocalData>(localData_)),
+        rcontext(RequestContext::saveContext()) {}
     std::function<void()> func;
     std::unique_ptr<Fiber::LocalData> localData;
+    std::shared_ptr<RequestContext> rcontext;
     AtomicLinkedListHook<RemoteTask> nextRemoteTask;
   };
 

folly/experimental/fibers/test/FibersTest.cpp

@@ -1332,6 +1332,81 @@ TEST(FiberManager, yieldTest) {
   EXPECT_TRUE(checkRan);
 }
 
+TEST(FiberManager, RequestContext) {
+  FiberManager fm(folly::make_unique<SimpleLoopController>());
+  auto& loopController =
+    dynamic_cast<SimpleLoopController&>(fm.loopController());
+
+  bool checkRun1 = false;
+  bool checkRun2 = false;
+  bool checkRun3 = false;
+
+  folly::fibers::Baton baton1;
+  folly::fibers::Baton baton2;
+  folly::fibers::Baton baton3;
+
+  folly::RequestContext::create();
+  auto rcontext1 = folly::RequestContext::get();
+  fm.addTask([&]() {
+      EXPECT_EQ(rcontext1, folly::RequestContext::get());
+      baton1.wait([&]() {
+          EXPECT_EQ(rcontext1, folly::RequestContext::get());
+        });
+      EXPECT_EQ(rcontext1, folly::RequestContext::get());
+      runInMainContext([&]() {
+          EXPECT_EQ(rcontext1, folly::RequestContext::get());
+        });
+      checkRun1 = true;
+    });
+
+  folly::RequestContext::create();
+  auto rcontext2 = folly::RequestContext::get();
+  fm.addTaskRemote([&]() {
+      EXPECT_EQ(rcontext2, folly::RequestContext::get());
+      baton2.wait();
+      EXPECT_EQ(rcontext2, folly::RequestContext::get());
+      checkRun2 = true;
+    });
+
+  folly::RequestContext::create();
+  auto rcontext3 = folly::RequestContext::get();
+  fm.addTaskFinally([&]() {
+      EXPECT_EQ(rcontext3, folly::RequestContext::get());
+      baton3.wait();
+      EXPECT_EQ(rcontext3, folly::RequestContext::get());
+
+      return folly::Unit();
+    },
+    [&](Try<folly::Unit>&& t) {
+      EXPECT_EQ(rcontext3, folly::RequestContext::get());
+      checkRun3 = true;
+    });
+
+  folly::RequestContext::create();
+  auto rcontext = folly::RequestContext::get();
+
+  fm.loopUntilNoReady();
+  EXPECT_EQ(rcontext, folly::RequestContext::get());
+
+  baton1.post();
+  EXPECT_EQ(rcontext, folly::RequestContext::get());
+  fm.loopUntilNoReady();
+  EXPECT_TRUE(checkRun1);
+  EXPECT_EQ(rcontext, folly::RequestContext::get());
+
+  baton2.post();
+  EXPECT_EQ(rcontext, folly::RequestContext::get());
+  fm.loopUntilNoReady();
+  EXPECT_TRUE(checkRun2);
+  EXPECT_EQ(rcontext, folly::RequestContext::get());
+
+  baton3.post();
+  EXPECT_EQ(rcontext, folly::RequestContext::get());
+  fm.loopUntilNoReady();
+  EXPECT_TRUE(checkRun3);
+  EXPECT_EQ(rcontext, folly::RequestContext::get());
+}
+
 static size_t sNumAwaits;
 
 void runBenchmark(size_t numAwaits, size_t toSend) {