Request context: Add hazard pointer-based implementation

Summary:
Add hazard-pointer-based implementation to allow concurrent reading during updates. This is done by using lightweight concurrent structures protected by hazard pointers instead of sequential structures protected by read locks.

The implementation is gated by a gflag that is off by default.

Reviewed By: davidtgoldblatt

Differential Revision: D18622072

fbshipit-source-id: 120ba14a7a559883e64a3fdf81b35a26315312c1

folly/io/async/Request.cpp

@@ -15,12 +15,20 @@
  */
 
 #include <folly/io/async/Request.h>
+#include <folly/experimental/SingleWriterFixedHashMap.h>
+#include <folly/synchronization/Hazptr.h>
 #include <folly/tracing/StaticTracepoint.h>
 
 #include <glog/logging.h>
 
 #include <folly/MapUtil.h>
 #include <folly/SingletonThreadLocal.h>
+#include <folly/portability/GFlags.h>
+
+DEFINE_bool(
+    reqctx_use_hazptr,
+    false,
+    "RequestContext implementation using hazard pointers");
 
 namespace folly {
 
@@ -62,6 +70,19 @@ Synchronized<F14FastMap<std::string, uint32_t>>& RequestToken::getCache() {
   return *cache;
 }
 
+void RequestData::acquireRef() {
+  auto rc = keepAliveCounter_.fetch_add(1, std::memory_order_relaxed);
+  DCHECK_GE(rc, 0);
+}
+
+void RequestData::releaseRefDeleteIfNoRefs() {
+  auto rc = keepAliveCounter_.fetch_sub(1, std::memory_order_acq_rel);
+  DCHECK_GT(rc, 0);
+  if (rc == 1) {
+    delete this;
+  }
+}
+
 void RequestData::DestructPtr::operator()(RequestData* ptr) {
   if (ptr) {
     auto keepAliveCounter =
@@ -84,14 +105,368 @@ void RequestData::DestructPtr::operator()(RequestData* ptr) {
   return SharedPtr(ptr);
 }
 
-bool RequestContext::doSetContextData(
-    const RequestToken& val,
+// The Combined struct keeps the two structures for context data
+// and callbacks together, so that readers can protect consistent
+// versions of the two structures together using hazard pointers.
+struct RequestContext::StateHazptr::Combined : hazptr_obj_base<Combined> {
+  static constexpr size_t kInitialCapacity = 4;
+  static constexpr size_t kSlackReciprocal = 4; // unused >= 1/4 capacity
+
+  // This must be optimized for lookup, its hot path is getContextData
+  // Efficiency of copying the container also matters in setShallowCopyContext
+  SingleWriterFixedHashMap<RequestToken, RequestData*> requestData_;
+  // This must be optimized for iteration, its hot path is setContext
+  SingleWriterFixedHashMap<RequestData*, bool> callbackData_;
+
+  Combined()
+      : requestData_(kInitialCapacity), callbackData_(kInitialCapacity) {}
+
+  Combined(const Combined& o)
+      : Combined(o.requestData_.capacity(), o.callbackData_.capacity(), o) {}
+
+  Combined(size_t dataCapacity, size_t callbackCapacity, const Combined& o)
+      : requestData_(dataCapacity, o.requestData_),
+        callbackData_(callbackCapacity, o.callbackData_) {}
+
+  Combined(Combined&&) = delete;
+  Combined& operator=(const Combined&) = delete;
+  Combined& operator=(Combined&&) = delete;
+
+  ~Combined() {
+    releaseDataRefs();
+  }
+
+  /* acquireDataRefs - Called at most once per Combined instance. */
+  void acquireDataRefs() {
+    for (auto it = requestData_.begin(); it != requestData_.end(); ++it) {
+      auto p = it.value();
+      if (p) {
+        p->acquireRef();
+      }
+    }
+  }
+
+  /* releaseDataRefs - Called only once from ~Combined */
+  void releaseDataRefs() {
+    for (auto it = requestData_.begin(); it != requestData_.end(); ++it) {
+      auto p = it.value();
+      if (p) {
+        p->releaseRefDeleteIfNoRefs();
+      }
+    }
+  }
+
+  /* needExpand */
+  bool needExpand() {
+    return needExpandRequestData() || needExpandCallbackData();
+  }
+
+  /* needExpandRequestData */
+  bool needExpandRequestData() {
+    return kSlackReciprocal * (requestData_.available() - 1) <
+        requestData_.capacity();
+  }
+
+  /* needExpandCallbackData */
+  bool needExpandCallbackData() {
+    return kSlackReciprocal * (callbackData_.available() - 1) <
+        callbackData_.capacity();
+  }
+}; // Combined
+
+RequestContext::StateHazptr::StateHazptr() {}
+
+RequestContext::StateHazptr::StateHazptr(const StateHazptr& o) {
+  Combined* oc = o.combined();
+  if (oc) {
+    auto p = new Combined(*oc);
+    p->acquireDataRefs();
+    setCombined(p);
+  }
+}
+
+RequestContext::StateHazptr::~StateHazptr() {
+  batch_.shutdown_and_reclaim();
+  auto p = combined();
+  if (p) {
+    delete p;
+  }
+}
+
+FOLLY_ALWAYS_INLINE
+RequestContext::StateHazptr::Combined* RequestContext::StateHazptr::combined()
+    const {
+  return combined_.load(std::memory_order_acquire);
+}
+
+RequestContext::StateHazptr::Combined*
+RequestContext::StateHazptr::ensureCombined() {
+  auto c = combined();
+  if (!c) {
+    c = new Combined;
+    setCombined(c);
+  }
+  return c;
+}
+
+void RequestContext::StateHazptr::setCombined(Combined* p) {
+  p->set_batch_tag(&batch_);
+  combined_.store(p, std::memory_order_release);
+}
+
+bool RequestContext::StateHazptr::doSetContextData(
+    const RequestToken& token,
+    std::unique_ptr<RequestData>& data,
+    DoSetBehaviour behaviour,
+    bool safe) {
+  SetContextDataResult result;
+  if (safe) {
+    result = doSetContextDataHelper(token, data, behaviour, safe);
+  } else {
+    std::lock_guard<std::mutex> g(mutex_);
+    result = doSetContextDataHelper(token, data, behaviour, safe);
+  }
+  if (result.unexpected) {
+    LOG_FIRST_N(WARNING, 1)
+        << "Calling RequestContext::setContextData for "
+        << token.getDebugString() << " but it is already set";
+  }
+  if (result.replaced) {
+    result.replaced->retire(); // Retire to hazptr library
+  }
+  return result.changed;
+}
+
+RequestContext::StateHazptr::SetContextDataResult
+RequestContext::StateHazptr::doSetContextDataHelper(
+    const RequestToken& token,
+    std::unique_ptr<RequestData>& data,
+    DoSetBehaviour behaviour,
+    bool safe) {
+  bool unexpected = false;
+  Combined* cur = ensureCombined();
+  Combined* replaced = nullptr;
+  auto it = cur->requestData_.find(token);
+  bool found = it != cur->requestData_.end();
+  if (found) {
+    if (behaviour == DoSetBehaviour::SET_IF_ABSENT) {
+      return {false /* no changes made */,
+              false /* nothing unexpected */,
+              nullptr /* combined not replaced */};
+    }
+    RequestData* oldData = it.value();
+    if (oldData) {
+      // Always erase non-null old data (and run its onUnset callback,
+      // if any). Non-null old data will always be overwritten either
+      // by the new data (if behavior is OVERWRITE) or by nullptr (if
+      // behavior is SET).
+      Combined* newCombined = eraseOldData(cur, token, oldData, safe);
+      if (newCombined) {
+        replaced = cur;
+        cur = newCombined;
+      }
+    }
+    if (behaviour == DoSetBehaviour::SET) {
+      // The expected behavior for SET when found is to reset the
+      // pointer and warn, without updating to the new data.
+      if (oldData) {
+        cur->requestData_.insert(token, nullptr);
+      }
+      unexpected = true;
+    } else {
+      DCHECK(behaviour == DoSetBehaviour::OVERWRITE);
+    }
+  }
+  if (!unexpected) {
+    // Replace combined if needed, call onSet if any, insert new data.
+    Combined* newCombined = insertNewData(cur, token, data, found);
+    if (newCombined) {
+      replaced = cur;
+      cur = newCombined;
+    }
+  }
+  if (replaced) {
+    // Now the new Combined is consistent. Safe to publish.
+    setCombined(cur);
+  }
+  return {true, /* changes were made */
+          unexpected,
+          replaced};
+}
+
+RequestContext::StateHazptr::Combined* FOLLY_NULLABLE
+RequestContext::StateHazptr::eraseOldData(
+    RequestContext::StateHazptr::Combined* cur,
+    const RequestToken& token,
+    RequestData* olddata,
+    bool safe) {
+  Combined* newCombined = nullptr;
+  // Call onUnset, if any.
+  if (olddata->hasCallback()) {
+    olddata->onUnset();
+    bool erased = cur->callbackData_.erase(olddata);
+    DCHECK(erased);
+  }
+  if (safe) {
+    // If the caller guarantees thread-safety, then erase the
+    // entry in the current version.
+    cur->requestData_.erase(token);
+    olddata->releaseRefDeleteIfNoRefs();
+  } else {
+    // If there may be concurrent readers, then copy-on-erase.
+    // Update the data reference counts to account for the
+    // existence of the new copy.
+    newCombined = new Combined(*cur);
+    newCombined->requestData_.erase(token);
+    newCombined->acquireDataRefs();
+  }
+  return newCombined;
+}
+
+RequestContext::StateHazptr::Combined* FOLLY_NULLABLE
+RequestContext::StateHazptr::insertNewData(
+    RequestContext::StateHazptr::Combined* cur,
+    const RequestToken& token,
+    std::unique_ptr<RequestData>& data,
+    bool found) {
+  Combined* newCombined = nullptr;
+  // Update value to point to the new data.
+  if (!found && cur->needExpand()) {
+    // Replace the current Combined with an expanded one
+    newCombined = expand(cur);
+    cur = newCombined;
+    cur->acquireDataRefs();
+  }
+  if (data && data->hasCallback()) {
+    // If data has callback, insert in callback structure, call onSet
+    cur->callbackData_.insert(data.get(), true);
+    data->onSet();
+  }
+  if (data) {
+    data->acquireRef();
+  }
+  cur->requestData_.insert(token, data.release());
+  return newCombined;
+}
+
+FOLLY_ALWAYS_INLINE
+bool RequestContext::StateHazptr::hasContextData(
+    const RequestToken& token) const {
+  hazptr_local<1> h;
+  Combined* combined = h[0].get_protected(combined_);
+  return combined ? combined->requestData_.contains(token) : false;
+}
+
+FOLLY_ALWAYS_INLINE
+RequestData* FOLLY_NULLABLE
+RequestContext::StateHazptr::getContextData(const RequestToken& token) {
+  hazptr_local<1> h;
+  Combined* combined = h[0].get_protected(combined_);
+  if (!combined) {
+    return nullptr;
+  }
+  auto& reqData = combined->requestData_;
+  auto it = reqData.find(token);
+  return it == reqData.end() ? nullptr : it.value();
+}
+
+FOLLY_ALWAYS_INLINE
+const RequestData* FOLLY_NULLABLE
+RequestContext::StateHazptr::getContextData(const RequestToken& token) const {
+  hazptr_local<1> h;
+  Combined* combined = h[0].get_protected(combined_);
+  if (!combined) {
+    return nullptr;
+  }
+  auto& reqData = combined->requestData_;
+  auto it = reqData.find(token);
+  return it == reqData.end() ? nullptr : it.value();
+}
+
+FOLLY_ALWAYS_INLINE
+void RequestContext::StateHazptr::onSet() {
+  // Don't use hazptr_local because callback may use hazptr
+  hazptr_holder<> h;
+  Combined* combined = h.get_protected(combined_);
+  if (!combined) {
+    return;
+  }
+  auto& cb = combined->callbackData_;
+  for (auto it = cb.begin(); it != cb.end(); ++it) {
+    it.key()->onSet();
+  }
+}
+
+FOLLY_ALWAYS_INLINE
+void RequestContext::StateHazptr::onUnset() {
+  // Don't use hazptr_local because callback may use hazptr
+  hazptr_holder<> h;
+  Combined* combined = h.get_protected(combined_);
+  if (!combined) {
+    return;
+  }
+  auto& cb = combined->callbackData_;
+  for (auto it = cb.begin(); it != cb.end(); ++it) {
+    it.key()->onUnset();
+  }
+}
+
+void RequestContext::StateHazptr::clearContextData(const RequestToken& token) {
+  Combined* replaced = nullptr;
+  { // Lock mutex_
+    std::lock_guard<std::mutex> g(mutex_);
+    Combined* cur = combined();
+    if (!cur) {
+      return;
+    }
+    auto it = cur->requestData_.find(token);
+    if (it == cur->requestData_.end()) {
+      return;
+    }
+    RequestData* data = it.value();
+    if (!data) {
+      cur->requestData_.erase(token);
+      return;
+    }
+    if (data->hasCallback()) {
+      data->onUnset();
+      cur->callbackData_.erase(data);
+    }
+    replaced = cur;
+    cur = new Combined(*replaced);
+    cur->requestData_.erase(token);
+    cur->acquireDataRefs();
+    setCombined(cur);
+  } // Unlock mutex_
+  DCHECK(replaced);
+  replaced->retire();
+}
+
+RequestContext::StateHazptr::Combined* RequestContext::StateHazptr::expand(
+    RequestContext::StateHazptr::Combined* c) {
+  size_t dataCapacity = c->requestData_.capacity();
+  if (c->needExpandRequestData()) {
+    dataCapacity *= 2;
+  }
+  size_t callbackCapacity = c->callbackData_.capacity();
+  if (c->needExpandCallbackData()) {
+    callbackCapacity *= 2;
+  }
+  return new Combined(dataCapacity, callbackCapacity, *c);
+}
+
+RequestContext::RequestContext() {
+  useHazptr_ = FLAGS_reqctx_use_hazptr;
+}
+
+bool RequestContext::doSetContextDataLock(
+    const RequestToken& token,
     std::unique_ptr<RequestData>& data,
     DoSetBehaviour behaviour) {
   auto wlock = state_.wlock();
   auto& state = *wlock;
 
-  auto it = state.requestData_.find(val);
+  auto it = state.requestData_.find(token);
   if (it != state.requestData_.end()) {
     if (behaviour == DoSetBehaviour::SET_IF_ABSENT) {
       return false;
@@ -106,7 +481,7 @@ bool RequestContext::doSetContextData(
     if (behaviour == DoSetBehaviour::SET) {
       LOG_FIRST_N(WARNING, 1)
           << "Calling RequestContext::setContextData for "
-          << val.getDebugString() << " but it is already set";
+          << token.getDebugString() << " but it is already set";
       return true;
     }
     DCHECK(behaviour == DoSetBehaviour::OVERWRITE);
@@ -120,45 +495,74 @@ bool RequestContext::doSetContextData(
   if (it != state.requestData_.end()) {
     it->second = std::move(ptr);
   } else {
-    state.requestData_.insert(std::make_pair(val, std::move(ptr)));
+    state.requestData_.insert(std::make_pair(token, std::move(ptr)));
   }
   return true;
 }
 
 void RequestContext::setContextData(
-    const RequestToken& val,
+    const RequestToken& token,
     std::unique_ptr<RequestData> data) {
-  doSetContextData(val, data, DoSetBehaviour::SET);
+  if (useHazptr()) {
+    stateHazptr_.doSetContextData(token, data, DoSetBehaviour::SET, false);
+    return;
+  }
+  doSetContextDataLock(token, data, DoSetBehaviour::SET);
 }
 
 bool RequestContext::setContextDataIfAbsent(
-    const RequestToken& val,
+    const RequestToken& token,
     std::unique_ptr<RequestData> data) {
-  return doSetContextData(val, data, DoSetBehaviour::SET_IF_ABSENT);
+  if (useHazptr()) {
+    return stateHazptr_.doSetContextData(
+        token, data, DoSetBehaviour::SET_IF_ABSENT, false);
+  }
+  return doSetContextDataLock(token, data, DoSetBehaviour::SET_IF_ABSENT);
 }
 
-void RequestContext::overwriteContextData(
-    const RequestToken& val,
+void RequestContext::overwriteContextDataLock(
+    const RequestToken& token,
     std::unique_ptr<RequestData> data) {
-  doSetContextData(val, data, DoSetBehaviour::OVERWRITE);
+  doSetContextDataLock(token, data, DoSetBehaviour::OVERWRITE);
+}
+
+void RequestContext::overwriteContextDataHazptr(
+    const RequestToken& token,
+    std::unique_ptr<RequestData> data,
+    bool safe) {
+  stateHazptr_.doSetContextData(token, data, DoSetBehaviour::OVERWRITE, safe);
 }
 
 bool RequestContext::hasContextData(const RequestToken& val) const {
+  if (useHazptr()) {
+    return stateHazptr_.hasContextData(val);
+  }
   return state_.rlock()->requestData_.count(val);
 }
 
-RequestData* RequestContext::getContextData(const RequestToken& val) {
+RequestData* FOLLY_NULLABLE
+RequestContext::getContextData(const RequestToken& val) {
+  if (useHazptr()) {
+    return stateHazptr_.getContextData(val);
+  }
   const RequestData::SharedPtr dflt{nullptr};
   return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
 }
 
-const RequestData* RequestContext::getContextData(
-    const RequestToken& val) const {
+const RequestData* FOLLY_NULLABLE
+RequestContext::getContextData(const RequestToken& val) const {
+  if (useHazptr()) {
+    return stateHazptr_.getContextData(val);
+  }
   const RequestData::SharedPtr dflt{nullptr};
   return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
 }
 
 void RequestContext::onSet() {
+  if (useHazptr()) {
+    stateHazptr_.onSet();
+    return;
+  }
   auto rlock = state_.rlock();
   for (const auto& data : rlock->callbackData_) {
     data->onSet();
@@ -166,6 +570,10 @@ void RequestContext::onSet() {
 }
 
 void RequestContext::onUnset() {
+  if (useHazptr()) {
+    stateHazptr_.onUnset();
+    return;
+  }
   auto rlock = state_.rlock();
   for (const auto& data : rlock->callbackData_) {
     data->onUnset();
@@ -173,6 +581,10 @@ void RequestContext::onUnset() {
 }
 
 void RequestContext::clearContextData(const RequestToken& val) {
+  if (useHazptr()) {
+    stateHazptr_.clearContextData(val);
+    return;
+  }
   RequestData::SharedPtr requestData;
   // Delete the RequestData after giving up the wlock just in case one of the
   // RequestData destructors will try to grab the lock again.
@@ -225,12 +637,12 @@ void exec_set_difference(const TData& data, const TData& other, TExec&& exec) {
 }
 } // namespace
 
-std::shared_ptr<RequestContext> RequestContext::setContext(
+/* static */ std::shared_ptr<RequestContext> RequestContext::setContext(
     std::shared_ptr<RequestContext> const& newCtx) {
   return setContext(copy(newCtx));
 }
 
-std::shared_ptr<RequestContext> RequestContext::setContext(
+/* static */ std::shared_ptr<RequestContext> RequestContext::setContext(
     std::shared_ptr<RequestContext>&& newCtx_) {
   auto newCtx = std::move(newCtx_); // enforce that it is really moved-from
 
@@ -242,6 +654,20 @@ std::shared_ptr<RequestContext> RequestContext::setContext(
   FOLLY_SDT(
       folly, request_context_switch_before, staticCtx.get(), newCtx.get());
 
+  if ((newCtx.get() && newCtx->useHazptr()) ||
+      (staticCtx.get() && staticCtx->useHazptr())) {
+    DCHECK(!newCtx.get() || newCtx->useHazptr());
+    DCHECK(!staticCtx.get() || staticCtx->useHazptr());
+    return RequestContext::setContextHazptr(newCtx, staticCtx);
+  } else {
+    return RequestContext::setContextLock(newCtx, staticCtx);
+  }
+}
+
+FOLLY_ALWAYS_INLINE
+/* static */ std::shared_ptr<RequestContext> RequestContext::setContextLock(
+    std::shared_ptr<RequestContext>& newCtx,
+    std::shared_ptr<RequestContext>& staticCtx) {
   auto curCtx = staticCtx;
   if (newCtx && curCtx) {
     // Only call set/unset for all request data that differs
@@ -268,6 +694,46 @@ std::shared_ptr<RequestContext> RequestContext::setContext(
   return curCtx;
 }
 
+FOLLY_ALWAYS_INLINE
+/* static */ std::shared_ptr<RequestContext> RequestContext::setContextHazptr(
+    std::shared_ptr<RequestContext>& newCtx,
+    std::shared_ptr<RequestContext>& staticCtx) {
+  auto curCtx = std::move(staticCtx);
+  bool checkCur = curCtx && curCtx->stateHazptr_.combined();
+  bool checkNew = newCtx && newCtx->stateHazptr_.combined();
+  if (checkCur && checkNew) {
+    hazptr_array<2> h;
+    auto curc = h[0].get_protected(curCtx->stateHazptr_.combined_);
+    auto newc = h[1].get_protected(newCtx->stateHazptr_.combined_);
+    auto& curcb = curc->callbackData_;
+    auto& newcb = newc->callbackData_;
+    for (auto it = curcb.begin(); it != curcb.end(); ++it) {
+      DCHECK(it.key());
+      auto data = it.key();
+      if (!newcb.contains(data)) {
+        data->onUnset();
+      }
+    }
+    staticCtx = std::move(newCtx);
+    for (auto it = newcb.begin(); it != newcb.end(); ++it) {
+      DCHECK(it.key());
+      auto data = it.key();
+      if (!curcb.contains(data)) {
+        data->onSet();
+      }
+    }
+  } else {
+    if (curCtx) {
+      curCtx->stateHazptr_.onUnset();
+    }
+    staticCtx = std::move(newCtx);
+    if (staticCtx) {
+      staticCtx->stateHazptr_.onSet();
+    }
+  }
+  return curCtx;
+}
+
 std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
   using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
   return SingletonT::get();

folly/io/async/Request.h

@@ -16,12 +16,15 @@
 
 #pragma once
 
-#include <memory>
-#include <string>
-
 #include <folly/Synchronized.h>
 #include <folly/container/F14Map.h>
 #include <folly/sorted_vector_types.h>
+#include <folly/synchronization/Hazptr.h>
+
+#include <atomic>
+#include <memory>
+#include <mutex>
+#include <string>
 
 namespace folly {
 
@@ -32,6 +35,7 @@ namespace folly {
  */
 class RequestToken {
  public:
+  RequestToken() = default;
   explicit RequestToken(const std::string& str);
 
   bool operator==(const RequestToken& other) const {
@@ -62,6 +66,18 @@ struct hash<folly::RequestToken> {
 
 namespace folly {
 
+// - A runtime flag GFLAGS_reqctx_use_hazptr determines the
+//   implementation of RequestContext.
+// - The flag false implementation uses sequential data structures
+//   protected by a read-write lock.
+// - The flag true implementation uses single-writer multi-readers
+//   data structures protected by hazard pointers for readers and a
+//   lock for writers.
+// - Each RequestContext instances contains a bool member useHazptr_
+//   (readable by a public member function useHazptr()) that indicates
+//   the implementation of the instance depending on the value of the
+//   GFLAG at instance construction time..
+
 // Some request context that follows an async request through a process
 // Everything in the context must be thread safe
 
@@ -83,6 +99,10 @@ class RequestData {
   // instance overrides the hasCallback method to return true otherwise
   // the callback will not be executed
   virtual void onUnset() {}
+  // For debugging
+  int refCount() {
+    return keepAliveCounter_.load(std::memory_order_acquire);
+  }
 
  private:
   // Start shallow copy implementation details:
@@ -93,6 +113,12 @@ class RequestData {
 
   friend class RequestContext;
 
+  // Reference-counting functions used by the hazptr-based implementation.
+  // Increment the reference count
+  void acquireRef();
+  // Decrement the reference count and delete if zero
+  void releaseRefDeleteIfNoRefs();
+
   // Unique ptr with custom destructor, decrement the counter
   // and only free if 0
   struct DestructPtr {
@@ -121,6 +147,8 @@ class RequestData {
 // copied between threads.
 class RequestContext {
  public:
+  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.
@@ -142,7 +170,7 @@ class RequestContext {
   // used, will print a warning message for the first time, clear the existing
   // RequestData instance for "val", and **not** add "data".
   void setContextData(
-      const RequestToken& val,
+      const RequestToken& token,
       std::unique_ptr<RequestData> data);
   void setContextData(
       const std::string& val,
@@ -154,7 +182,7 @@ class RequestContext {
   // string identifier "val". If the same string identifier has already been
   // used, return false and do nothing. Otherwise add "data" and return true.
   bool setContextDataIfAbsent(
-      const RequestToken& val,
+      const RequestToken& token,
       std::unique_ptr<RequestData> data);
   bool setContextDataIfAbsent(
       const std::string& val,
@@ -163,22 +191,22 @@ class RequestContext {
   }
 
   // Remove the RequestData instance with string identifier "val", if it exists.
-  void clearContextData(const RequestToken& val);
+  void clearContextData(const RequestToken& token);
   void clearContextData(const std::string& val) {
     clearContextData(RequestToken(val));
   }
 
   // Returns true if and only if the RequestData instance with string identifier
   // "val" exists in this RequestContext instnace.
-  bool hasContextData(const RequestToken& val) const;
+  bool hasContextData(const RequestToken& token) const;
   bool hasContextData(const std::string& val) const {
     return hasContextData(RequestToken(val));
   }
 
   // Get (constant) raw pointer of the RequestData instance with string
   // identifier "val" if it exists, otherwise returns null pointer.
-  RequestData* getContextData(const RequestToken& val);
-  const RequestData* getContextData(const RequestToken& val) const;
+  RequestData* getContextData(const RequestToken& token);
+  const RequestData* getContextData(const RequestToken& token) const;
   RequestData* getContextData(const std::string& val) {
     return getContextData(RequestToken(val));
   }
@@ -189,6 +217,11 @@ class RequestContext {
   void onSet();
   void onUnset();
 
+  // useHazptr
+  FOLLY_ALWAYS_INLINE bool useHazptr() const {
+    return useHazptr_;
+  }
+
   // The following API is used to pass the context through queues / threads.
   // saveContext is called to get a shared_ptr to the context, and
   // setContext is used to reset it on the other side of the queue.
@@ -211,6 +244,13 @@ class RequestContext {
  private:
   static std::shared_ptr<RequestContext>& getStaticContext();
 
+  static std::shared_ptr<RequestContext> setContextLock(
+      std::shared_ptr<RequestContext>& newCtx,
+      std::shared_ptr<RequestContext>& staticCtx);
+  static std::shared_ptr<RequestContext> setContextHazptr(
+      std::shared_ptr<RequestContext>& newCtx,
+      std::shared_ptr<RequestContext>& staticCtx);
+
   // Start shallow copy guard implementation details:
   // All methods are private to encourage proper use
   friend struct ShallowCopyRequestContextScopeGuard;
@@ -221,33 +261,62 @@ class RequestContext {
 
   // Similar to setContextData, except it overwrites the data
   // if already set (instead of warn + reset ptr).
-  void overwriteContextData(
-      const RequestToken& val,
+  void overwriteContextDataLock(
+      const RequestToken& token,
       std::unique_ptr<RequestData> data);
-  void overwriteContextData(
+  void overwriteContextDataLock(
       const std::string& val,
       std::unique_ptr<RequestData> data) {
-    overwriteContextData(RequestToken(val), std::move(data));
+    overwriteContextDataLock(RequestToken(val), std::move(data));
   }
   // End shallow copy guard
 
+  // For functions with a parameter safe, if safe is true then the
+  // caller guarantees that there are no concurrent readers or writers
+  // accessing the structure.
+  void overwriteContextDataHazptr(
+      const RequestToken& token,
+      std::unique_ptr<RequestData> data,
+      bool safe = false);
+  void overwriteContextDataHazptr(
+      const std::string& val,
+      std::unique_ptr<RequestData> data,
+      bool safe = false) {
+    overwriteContextDataHazptr(RequestToken(val), std::move(data), safe);
+  }
+
   enum class DoSetBehaviour {
     SET,
     SET_IF_ABSENT,
     OVERWRITE,
   };
 
-  bool doSetContextData(
-      const RequestToken& val,
+  bool doSetContextDataLock(
+      const RequestToken& token,
       std::unique_ptr<RequestData>& data,
       DoSetBehaviour behaviour);
-  bool doSetContextData(
+  bool doSetContextDataLock(
       const std::string& val,
       std::unique_ptr<RequestData>& data,
       DoSetBehaviour behaviour) {
-    return doSetContextData(RequestToken(val), data, behaviour);
+    return doSetContextDataLock(RequestToken(val), data, behaviour);
+  }
+
+  bool doSetContextDataHazptr(
+      const RequestToken& token,
+      std::unique_ptr<RequestData>& data,
+      DoSetBehaviour behaviour,
+      bool safe = false);
+  bool doSetContextDataHazptr(
+      const std::string& val,
+      std::unique_ptr<RequestData>& data,
+      DoSetBehaviour behaviour,
+      bool safe = false) {
+    return doSetContextDataHazptr(RequestToken(val), data, behaviour, safe);
   }
 
+  // State immplementation with sequential data structures protected by a
+  // read-write locks.
   struct State {
     // This must be optimized for lookup, its hot path is getContextData
     // Efficiency of copying the container also matters in setShallowCopyContext
@@ -258,6 +327,67 @@ class RequestContext {
     sorted_vector_set<RequestData*> callbackData_;
   };
   folly::Synchronized<State> state_;
+
+  // State implementation with single-writer multi-reader data
+  // structures protected by hazard pointers for readers and a lock
+  // for writers.
+  struct StateHazptr {
+    // Hazard pointer-protected combined structure for request data
+    // and callbacks.
+    struct Combined;
+    hazptr_obj_batch<> batch_; // For destruction order
+    std::atomic<Combined*> combined_{nullptr};
+    std::mutex mutex_;
+
+    StateHazptr();
+    StateHazptr(const StateHazptr& o);
+    StateHazptr(StateHazptr&&) = delete;
+    StateHazptr& operator=(const StateHazptr&) = delete;
+    StateHazptr& operator=(StateHazptr&&) = delete;
+    ~StateHazptr();
+
+    Combined* combined() const;
+    Combined* ensureCombined(); // Lazy allocation if needed
+    void setCombined(Combined* combined);
+    Combined* expand(Combined* combined);
+
+    bool doSetContextData(
+        const RequestToken& token,
+        std::unique_ptr<RequestData>& data,
+        DoSetBehaviour behaviour,
+        bool safe);
+    bool hasContextData(const RequestToken& token) const;
+    RequestData* getContextData(const RequestToken& token);
+    const RequestData* getContextData(const RequestToken& token) const;
+    void onSet();
+    void onUnset();
+    void clearContextData(const RequestToken& token);
+
+   private:
+    struct SetContextDataResult {
+      bool changed; // Changes were made
+      bool unexpected; // Update was unexpected
+      Combined* replaced; // The combined structure was replaced
+    };
+
+    SetContextDataResult doSetContextDataHelper(
+        const RequestToken& token,
+        std::unique_ptr<RequestData>& data,
+        DoSetBehaviour behaviour,
+        bool safe);
+    Combined* eraseOldData(
+        Combined* combined,
+        const RequestToken& token,
+        RequestData* oldData,
+        bool safe);
+    Combined* insertNewData(
+        Combined* combined,
+        const RequestToken& token,
+        std::unique_ptr<RequestData>& data,
+        bool found);
+  }; // StateHazptr
+  StateHazptr stateHazptr_;
+  bool useHazptr_;
 };
 
 /**
@@ -310,16 +440,26 @@ struct ShallowCopyRequestContextScopeGuard {
    * "clearRequestData" then "setRequestData" after the guard.
    */
   ShallowCopyRequestContextScopeGuard(
-      const RequestToken& val,
+      const RequestToken& token,
       std::unique_ptr<RequestData> data)
       : ShallowCopyRequestContextScopeGuard() {
-    RequestContext::get()->overwriteContextData(val, std::move(data));
+    auto ctx = RequestContext::get();
+    if (ctx->useHazptr()) {
+      ctx->overwriteContextDataHazptr(token, std::move(data), true);
+    } else {
+      ctx->overwriteContextDataLock(token, std::move(data));
+    }
   }
   ShallowCopyRequestContextScopeGuard(
       const std::string& val,
       std::unique_ptr<RequestData> data)
       : ShallowCopyRequestContextScopeGuard() {
-    RequestContext::get()->overwriteContextData(val, std::move(data));
+    auto ctx = RequestContext::get();
+    if (ctx->useHazptr()) {
+      ctx->overwriteContextDataHazptr(val, std::move(data), true);
+    } else {
+      ctx->overwriteContextDataLock(val, std::move(data));
+    }
   }
 
   ~ShallowCopyRequestContextScopeGuard() {