Make RequestContext provider overridable in order to save cost of set…

…Context() on each fiber context switch

Summary:
Each fiber context switch currently involves the cost of saving/restoring `RequestContext`.  Certain
fibers-based applications such as mcrouter don't use `RequestContext` at all, so updating the current
thread-global `RequestContext` on each fiber context switch is unnecessary overhead.  Avoid this cost
by allowing `FiberManager` to override the `RequestContext` provider at the start and end of each fiber drain
loop.

Reviewed By: andriigrynenko

Differential Revision: D5787837

fbshipit-source-id: ea9041ce228063c8701165366fd1e34132868d22

folly/fibers/FiberManagerInternal-inl.h

@@ -113,7 +113,6 @@ inline void FiberManager::runReadyFiber(Fiber* fiber) {
       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));
   }
@@ -139,7 +138,6 @@ 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_;
@@ -161,7 +159,6 @@ 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();
 
@@ -179,7 +176,6 @@ 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);
   }
@@ -200,8 +196,20 @@ inline void FiberManager::loopUntilNoReadyImpl() {
   auto originalFiberManager = this;
   std::swap(currentFiberManager_, originalFiberManager);
 
+  RequestContext::Provider oldRequestContextProvider;
+  auto newRequestContextProvider =
+      [this, &oldRequestContextProvider]() -> std::shared_ptr<RequestContext>& {
+    return currentFiber_ ? currentFiber_->rcontext_
+                         : oldRequestContextProvider();
+  };
+  oldRequestContextProvider = RequestContext::setRequestContextProvider(
+      std::ref(newRequestContextProvider));
+
   SCOPE_EXIT {
     isLoopScheduled_ = false;
+    // Restore RequestContext provider before call to ensureLoopScheduled()
+    RequestContext::setRequestContextProvider(
+        std::move(oldRequestContextProvider));
     if (!readyFibers_.empty()) {
       ensureLoopScheduled();
     }

folly/fibers/test/FibersTest.cpp

@@ -33,6 +33,7 @@
 #include <folly/fibers/SimpleLoopController.h>
 #include <folly/fibers/TimedMutex.h>
 #include <folly/fibers/WhenN.h>
+#include <folly/io/async/Request.h>
 #include <folly/io/async/ScopedEventBaseThread.h>
 #include <folly/portability/GTest.h>
 
@@ -1236,6 +1237,107 @@ TEST(FiberManager, fiberLocalDestructor) {
   EXPECT_FALSE(fm.hasTasks());
 }
 
+TEST(FiberManager, fiberRequestContext) {
+  folly::EventBase evb;
+  FiberManager fm(std::make_unique<EventBaseLoopController>());
+  dynamic_cast<EventBaseLoopController&>(fm.loopController())
+      .attachEventBase(evb);
+
+  struct TestContext : public folly::RequestData {
+    explicit TestContext(std::string s) : data(std::move(s)) {}
+    std::string data;
+  };
+
+  class AfterFibersCallback : public folly::EventBase::LoopCallback {
+   public:
+    AfterFibersCallback(
+        folly::EventBase& evb,
+        const bool& fibersDone,
+        folly::Function<void()> afterFibersFunc)
+        : evb_(evb),
+          fibersDone_(fibersDone),
+          afterFibersFunc_(std::move(afterFibersFunc)) {}
+
+    void runLoopCallback() noexcept override {
+      if (fibersDone_) {
+        afterFibersFunc_();
+        delete this;
+      } else {
+        evb_.runInLoop(this);
+      }
+    }
+
+   private:
+    folly::EventBase& evb_;
+    const bool& fibersDone_;
+    folly::Function<void()> afterFibersFunc_;
+  };
+
+  bool fibersDone = false;
+  size_t tasksRun = 0;
+  evb.runInEventBaseThread([&evb, &fm, &tasksRun, &fibersDone]() {
+    ++tasksRun;
+    auto* const evbCtx = folly::RequestContext::get();
+    EXPECT_NE(nullptr, evbCtx);
+    EXPECT_EQ(nullptr, evbCtx->getContextData("key"));
+    evbCtx->setContextData("key", std::make_unique<TestContext>("evb_value"));
+
+    // This callback allows us to check that FiberManager has restored the
+    // RequestContext provider as expected after a fiber loop.
+    auto* afterFibersCallback =
+        new AfterFibersCallback(evb, fibersDone, [&tasksRun, evbCtx]() {
+          ++tasksRun;
+          EXPECT_EQ(evbCtx, folly::RequestContext::get());
+          EXPECT_EQ(
+              "evb_value",
+              dynamic_cast<TestContext*>(evbCtx->getContextData("key"))->data);
+        });
+    evb.runInLoop(afterFibersCallback);
+
+    // Launching a fiber allows us to hit FiberManager RequestContext
+    // setup/teardown logic.
+    fm.addTask([&evb, &tasksRun, &fibersDone, evbCtx]() {
+      ++tasksRun;
+
+      // Initially, fiber starts with same RequestContext as its parent task.
+      EXPECT_EQ(evbCtx, folly::RequestContext::get());
+      EXPECT_NE(nullptr, evbCtx->getContextData("key"));
+      EXPECT_EQ(
+          "evb_value",
+          dynamic_cast<TestContext*>(evbCtx->getContextData("key"))->data);
+
+      // Create a new RequestContext for this fiber so we can distinguish from
+      // RequestContext first EventBase callback started with.
+      folly::RequestContext::create();
+      auto* const fiberCtx = folly::RequestContext::get();
+      EXPECT_NE(nullptr, fiberCtx);
+      EXPECT_EQ(nullptr, fiberCtx->getContextData("key"));
+      fiberCtx->setContextData(
+          "key", std::make_unique<TestContext>("fiber_value"));
+
+      // Task launched from within fiber should share current fiber's
+      // RequestContext
+      evb.runInEventBaseThread([&tasksRun, fiberCtx]() {
+        ++tasksRun;
+        auto* const evbCtx2 = folly::RequestContext::get();
+        EXPECT_EQ(fiberCtx, evbCtx2);
+        EXPECT_NE(nullptr, evbCtx2->getContextData("key"));
+        EXPECT_EQ(
+            "fiber_value",
+            dynamic_cast<TestContext*>(evbCtx2->getContextData("key"))->data);
+      });
+
+      fibersDone = true;
+    });
+  });
+
+  evb.loop();
+
+  EXPECT_EQ(4, tasksRun);
+  EXPECT_TRUE(fibersDone);
+  EXPECT_FALSE(fm.hasTasks());
+}
+
 TEST(FiberManager, yieldTest) {
   FiberManager manager(std::make_unique<SimpleLoopController>());
   auto& loopController =

folly/io/async/Request.cpp

@@ -13,14 +13,17 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 #include <folly/io/async/Request.h>
-#include <folly/tracing/StaticTracepoint.h>
+
+#include <algorithm>
+#include <stdexcept>
+#include <utility>
 
 #include <glog/logging.h>
 
 #include <folly/MapUtil.h>
 #include <folly/SingletonThreadLocal.h>
+#include <folly/tracing/StaticTracepoint.h>
 
 namespace folly {
 
@@ -115,19 +118,50 @@ std::shared_ptr<RequestContext> RequestContext::setContext(
   return ctx;
 }
 
-std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
-  using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
-  static SingletonT singleton;
+RequestContext::Provider& RequestContext::requestContextProvider() {
+  class DefaultProvider {
+   public:
+    constexpr DefaultProvider() = default;
+    DefaultProvider(const DefaultProvider&) = delete;
+    DefaultProvider& operator=(const DefaultProvider&) = delete;
+    DefaultProvider(DefaultProvider&&) = default;
+    DefaultProvider& operator=(DefaultProvider&&) = default;
+
+    std::shared_ptr<RequestContext>& operator()() {
+      return context;
+    }
+
+   private:
+    std::shared_ptr<RequestContext> context;
+  };
 
-  return singleton.get();
+  static SingletonThreadLocal<Provider> providerSingleton(
+      []() { return new Provider(DefaultProvider()); });
+  return providerSingleton.get();
+}
+
+std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
+  auto& provider = requestContextProvider();
+  return provider();
 }
 
 RequestContext* RequestContext::get() {
-  auto context = getStaticContext();
+  auto& context = getStaticContext();
   if (!context) {
     static RequestContext defaultContext;
     return std::addressof(defaultContext);
   }
   return context.get();
 }
+
+RequestContext::Provider RequestContext::setRequestContextProvider(
+    RequestContext::Provider newProvider) {
+  if (!newProvider) {
+    throw std::runtime_error("RequestContext provider must be non-empty");
+  }
+
+  auto& provider = requestContextProvider();
+  std::swap(provider, newProvider);
+  return newProvider;
+}
 }

folly/io/async/Request.h

@@ -19,6 +19,7 @@
 #include <map>
 #include <memory>
 
+#include <folly/Function.h>
 #include <folly/SharedMutex.h>
 #include <folly/Synchronized.h>
 
@@ -45,6 +46,8 @@ class RequestContext;
 // copied between threads.
 class RequestContext {
  public:
+  using Provider = folly::Function<std::shared_ptr<RequestContext>&()>;
+
   // Create a unique request context for this request.
   // It will be passed between queues / threads (where implemented),
   // so it should be valid for the lifetime of the request.
@@ -95,8 +98,22 @@ class RequestContext {
     return getStaticContext();
   }
 
+  // This API allows one to override the default behavior of getStaticContext()
+  // by providing a custom RequestContext provider. The old provider is
+  // returned, and the user must restore the old provider via a subsequent call
+  // to setRequestContextProvider() once the new provider is no longer needed.
+  //
+  // Using custom RequestContext providers can be more efficient than having to
+  // setContext() whenever context must be switched. This is especially true in
+  // applications that do not actually use RequestContext, but where library
+  // code must still support RequestContext for other use cases.  See
+  // FiberManager for an example of how a custom RequestContext provider can
+  // reduce calls to setContext().
+  static Provider setRequestContextProvider(Provider f);
+
  private:
   static std::shared_ptr<RequestContext>& getStaticContext();
+  static Provider& requestContextProvider();
 
   using Data = std::map<std::string, std::unique_ptr<RequestData>>;
   folly::Synchronized<Data, folly::SharedMutex> data_;

folly/io/async/test/RequestContextTest.cpp

@@ -77,6 +77,49 @@ TEST(RequestContext, SimpleTest) {
   EXPECT_TRUE(nullptr != RequestContext::get());
 }
 
+TEST(RequestContext, nonDefaultContextsAreThreadLocal) {
+  RequestContext* ctx1 = nullptr;
+  RequestContext* ctx2 = nullptr;
+
+  std::vector<std::thread> ts;
+  for (size_t i = 0; i < 2; ++i) {
+    auto*& ctx = (i == 0 ? ctx1 : ctx2);
+    ts.emplace_back([&ctx]() {
+      RequestContext::create();
+      ctx = RequestContext::get();
+    });
+  }
+  for (auto& t : ts) {
+    t.join();
+  }
+
+  EXPECT_NE(nullptr, ctx1);
+  EXPECT_NE(nullptr, ctx2);
+  EXPECT_NE(ctx1, ctx2);
+}
+
+TEST(RequestContext, customRequestContextProvider) {
+  auto customContext = std::make_shared<RequestContext>();
+  auto customProvider = [&customContext]() -> std::shared_ptr<RequestContext>& {
+    return customContext;
+  };
+
+  auto* const originalContext = RequestContext::get();
+  EXPECT_NE(nullptr, originalContext);
+
+  // Install new RequestContext provider
+  auto originalProvider =
+      RequestContext::setRequestContextProvider(std::move(customProvider));
+
+  auto* const newContext = RequestContext::get();
+  EXPECT_EQ(customContext.get(), newContext);
+  EXPECT_NE(originalContext, newContext);
+
+  // Restore original RequestContext provider
+  RequestContext::setRequestContextProvider(std::move(originalProvider));
+  EXPECT_EQ(originalContext, RequestContext::get());
+}
+
 TEST(RequestContext, setIfAbsentTest) {
   EXPECT_TRUE(RequestContext::get() != nullptr);