Use hazptr in RequestContext

Summary:
Instead of a reader-writer lock in RequestContext data_, use an allocation + hazptr.
This improves setContext/saveContext perf at the expense of setData.

Differential Revision: D8397670

fbshipit-source-id: 105fd715f1eb09499c88d1a205f019ae01a54f13

folly/futures/detail/Core.h

@@ -31,6 +31,7 @@
 #include <folly/lang/Assume.h>
 #include <folly/lang/Exception.h>
 #include <folly/synchronization/MicroSpinLock.h>
+#include <glog/logging.h>
 
 #include <folly/io/async/Request.h>
 

folly/io/async/Request.cpp

@@ -21,44 +21,68 @@
 
 #include <folly/MapUtil.h>
 #include <folly/SingletonThreadLocal.h>
+#include <folly/synchronization/Hazptr.h>
 
 namespace folly {
 
+struct RequestContext::State : hazptr_obj_base<State> {
+  std::map<std::string, std::shared_ptr<RequestData>> requestData_;
+  std::set<RequestData*> callbackData_;
+};
+
+RequestContext::RequestContext() : state_(new State) {}
+
+RequestContext::~RequestContext() {
+  auto p = state_.load(std::memory_order_relaxed);
+  delete p;
+}
+
 bool RequestContext::doSetContextData(
     const std::string& val,
     std::unique_ptr<RequestData>& data,
     bool strict) {
-  auto ulock = state_.ulock();
-
-  bool conflict = false;
-  auto it = ulock->requestData_.find(val);
-  if (it != ulock->requestData_.end()) {
-    if (strict) {
-      return false;
-    } else {
-      LOG_FIRST_N(WARNING, 1) << "Calling RequestContext::setContextData for "
-                              << val << " but it is already set";
-      conflict = true;
+  State* p{nullptr};
+
+  {
+    std::lock_guard<std::mutex> g(m_);
+    p = state_.load(std::memory_order_acquire);
+
+    bool conflict = false;
+    auto it = p->requestData_.find(val);
+    if (it != p->requestData_.end()) {
+      if (strict) {
+        return false;
+      } else {
+        LOG_FIRST_N(WARNING, 1) << "Calling RequestContext::setContextData for "
+                                << val << " but it is already set";
+        conflict = true;
+      }
     }
-  }
 
-  auto wlock = ulock.moveFromUpgradeToWrite();
-  if (conflict) {
-    if (it->second) {
-      if (it->second->hasCallback()) {
-        it->second->onUnset();
-        wlock->callbackData_.erase(it->second.get());
+    auto newstate = new State(*p);
+    it = newstate->requestData_.find(val);
+
+    if (conflict) {
+      if (it->second) {
+        if (it->second->hasCallback()) {
+          it->second->onUnset();
+          newstate->callbackData_.erase(it->second.get());
+        }
+        it->second.reset();
       }
-      it->second.reset(nullptr);
     }
-    return true;
+
+    if (data && data->hasCallback()) {
+      newstate->callbackData_.insert(data.get());
+      data->onSet();
+    }
+    newstate->requestData_[val] = std::move(data);
+    state_.store(newstate, std::memory_order_release);
   }
 
-  if (data && data->hasCallback()) {
-    wlock->callbackData_.insert(data.get());
-    data->onSet();
+  if (p) {
+    p->retire();
   }
-  wlock->requestData_[val] = std::move(data);
 
   return true;
 }
@@ -76,38 +100,49 @@ bool RequestContext::setContextDataIfAbsent(
 }
 
 bool RequestContext::hasContextData(const std::string& val) const {
-  return state_.rlock()->requestData_.count(val);
+  hazptr_local<1> hptr;
+  auto p = hptr[0].get_protected(state_);
+  return p->requestData_.count(val) != 0;
 }
 
 RequestData* RequestContext::getContextData(const std::string& val) {
-  const std::unique_ptr<RequestData> dflt{nullptr};
-  return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
+  const std::shared_ptr<RequestData> dflt{nullptr};
+  hazptr_local<1> hptr;
+  auto p = hptr[0].get_protected(state_);
+  p->requestData_.count(val);
+  return get_ref_default(p->requestData_, val, dflt).get();
 }
 
 const RequestData* RequestContext::getContextData(
     const std::string& val) const {
-  const std::unique_ptr<RequestData> dflt{nullptr};
-  return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
+  const std::shared_ptr<RequestData> dflt{nullptr};
+  hazptr_local<1> hptr;
+  auto p = hptr[0].get_protected(state_);
+  p->requestData_.count(val);
+  return get_ref_default(p->requestData_, val, dflt).get();
 }
 
 void RequestContext::onSet() {
-  auto rlock = state_.rlock();
-  for (const auto& data : rlock->callbackData_) {
+  hazptr_holder<> hptr;
+  auto p = hptr.get_protected(state_);
+  for (const auto& data : p->callbackData_) {
     data->onSet();
   }
 }
 
 void RequestContext::onUnset() {
-  auto rlock = state_.rlock();
-  for (const auto& data : rlock->callbackData_) {
+  hazptr_holder<> hptr;
+  auto p = hptr.get_protected(state_);
+  for (const auto& data : p->callbackData_) {
     data->onUnset();
   }
 }
 
 std::shared_ptr<RequestContext> RequestContext::createChild() {
+  hazptr_local<1> hptr;
+  auto p = hptr[0].get_protected(state_);
   auto child = std::make_shared<RequestContext>();
-  auto rlock = state_.rlock();
-  for (const auto& entry : rlock->requestData_) {
+  for (const auto& entry : p->requestData_) {
     auto& key = entry.first;
     auto childData = entry.second->createChild();
     if (childData) {
@@ -118,24 +153,37 @@ std::shared_ptr<RequestContext> RequestContext::createChild() {
 }
 
 void RequestContext::clearContextData(const std::string& val) {
-  std::unique_ptr<RequestData> requestData;
-  // Delete the RequestData after giving up the wlock just in case one of the
-  // RequestData destructors will try to grab the lock again.
+  std::shared_ptr<RequestData> requestData;
+  // requestData is deleted while lock is not held, in case
+  // destructor contains calls to RequestContext
+  State* p{nullptr};
+
   {
-    auto ulock = state_.ulock();
-    auto it = ulock->requestData_.find(val);
-    if (it == ulock->requestData_.end()) {
+    std::lock_guard<std::mutex> g(m_);
+    p = state_.load(std::memory_order_acquire);
+
+    auto it = p->requestData_.find(val);
+    if (it == p->requestData_.end()) {
       return;
     }
 
-    auto wlock = ulock.moveFromUpgradeToWrite();
+    auto newstate = new State(*p);
+
+    it = newstate->requestData_.find(val);
+    CHECK(it != newstate->requestData_.end());
+
     if (it->second && it->second->hasCallback()) {
       it->second->onUnset();
-      wlock->callbackData_.erase(it->second.get());
+      newstate->callbackData_.erase(it->second.get());
     }
 
-    requestData = std::move(it->second);
-    wlock->requestData_.erase(it);
+    requestData = it->second;
+    newstate->requestData_.erase(it);
+    state_.store(newstate, std::memory_order_release);
+  }
+
+  if (p) {
+    p->retire();
   }
 }
 

folly/io/async/Request.h

@@ -16,12 +16,12 @@
 
 #pragma once
 
+#include <atomic>
 #include <map>
 #include <memory>
+#include <mutex>
 #include <set>
 
-#include <folly/Synchronized.h>
-
 namespace folly {
 
 // Some request context that follows an async request through a process
@@ -58,6 +58,8 @@ class RequestData {
 // copied between threads.
 class RequestContext {
  public:
+  RequestContext();
+  ~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.
@@ -131,11 +133,9 @@ class RequestContext {
       std::unique_ptr<RequestData>& data,
       bool strict);
 
-  struct State {
-    std::map<std::string, std::unique_ptr<RequestData>> requestData_;
-    std::set<RequestData*> callbackData_;
-  };
-  folly::Synchronized<State> state_;
+  struct State;
+  std::atomic<State*> state_;
+  std::mutex m_;
 };
 
 class RequestContextScopeGuard {

folly/io/async/VirtualEventBase.h

@@ -19,6 +19,7 @@
 #include <future>
 
 #include <folly/Executor.h>
+#include <folly/Synchronized.h>
 #include <folly/io/async/EventBase.h>
 #include <folly/synchronization/Baton.h>