RequestToken

Summary:
Introduce a RequestToken class to RequestContext.  Map indicies are now uint32_t.

New interfaces are exposed using RequestToken, previous std::string interface forwards
to new interface.  This is a regression for the previous std::string interface since it
must do two hash lookups now, however, the new interface is faster.

Reviewed By: A5he

Differential Revision: D9197415

fbshipit-source-id: ffb8be51c609233fecfdff205a3ea4d4bb68de4b

folly/io/async/Request.cpp

@@ -24,6 +24,46 @@
 
 namespace folly {
 
+RequestToken::RequestToken(const std::string& str) {
+  auto& cache = getCache();
+  {
+    auto c = cache.rlock();
+    auto res = c->find(str);
+    if (res != c->end()) {
+      token_ = res->second;
+      return;
+    }
+  }
+  auto c = cache.wlock();
+  auto res = c->find(str);
+  if (res != c->end()) {
+    token_ = res->second;
+    return;
+  }
+  static uint32_t nextToken{1};
+
+  token_ = nextToken++;
+  (*c)[str] = token_;
+}
+
+std::string RequestToken::getDebugString() const {
+  auto& cache = getCache();
+  auto c = cache.rlock();
+  for (auto& v : *c) {
+    if (v.second == token_) {
+      return v.first;
+    }
+  }
+  throw std::logic_error("Could not find debug string in RequestToken");
+}
+
+Synchronized<std::unordered_map<std::string, uint32_t>>&
+RequestToken::getCache() {
+  static Indestructible<Synchronized<std::unordered_map<std::string, uint32_t>>>
+      cache;
+  return *cache;
+}
+
 void RequestData::DestructPtr::operator()(RequestData* ptr) {
   if (ptr) {
     auto keepAliveCounter =
@@ -47,7 +87,7 @@ void RequestData::DestructPtr::operator()(RequestData* ptr) {
 }
 
 bool RequestContext::doSetContextData(
-    const std::string& val,
+    const RequestToken& val,
     std::unique_ptr<RequestData>& data,
     DoSetBehaviour behaviour) {
   auto ulock = state_.ulock();
@@ -58,8 +98,9 @@ bool RequestContext::doSetContextData(
     if (behaviour == DoSetBehaviour::SET_IF_ABSENT) {
       return false;
     } else if (behaviour == DoSetBehaviour::SET) {
-      LOG_FIRST_N(WARNING, 1) << "Calling RequestContext::setContextData for "
-                              << val << " but it is already set";
+      LOG_FIRST_N(WARNING, 1)
+          << "Calling RequestContext::setContextData for "
+          << val.getDebugString() << " but it is already set";
     }
     conflict = true;
   }
@@ -88,34 +129,34 @@ bool RequestContext::doSetContextData(
 }
 
 void RequestContext::setContextData(
-    const std::string& val,
+    const RequestToken& val,
     std::unique_ptr<RequestData> data) {
   doSetContextData(val, data, DoSetBehaviour::SET);
 }
 
 bool RequestContext::setContextDataIfAbsent(
-    const std::string& val,
+    const RequestToken& val,
     std::unique_ptr<RequestData> data) {
   return doSetContextData(val, data, DoSetBehaviour::SET_IF_ABSENT);
 }
 
 void RequestContext::overwriteContextData(
-    const std::string& val,
+    const RequestToken& val,
     std::unique_ptr<RequestData> data) {
   doSetContextData(val, data, DoSetBehaviour::OVERWRITE);
 }
 
-bool RequestContext::hasContextData(const std::string& val) const {
+bool RequestContext::hasContextData(const RequestToken& val) const {
   return state_.rlock()->requestData_.count(val);
 }
 
-RequestData* RequestContext::getContextData(const std::string& val) {
+RequestData* RequestContext::getContextData(const RequestToken& val) {
   const RequestData::SharedPtr dflt{nullptr};
   return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
 }
 
 const RequestData* RequestContext::getContextData(
-    const std::string& val) const {
+    const RequestToken& val) const {
   const RequestData::SharedPtr dflt{nullptr};
   return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
 }
@@ -134,7 +175,7 @@ void RequestContext::onUnset() {
   }
 }
 
-void RequestContext::clearContextData(const std::string& val) {
+void RequestContext::clearContextData(const RequestToken& val) {
   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.

folly/io/async/Request.h

@@ -25,6 +25,43 @@
 
 namespace folly {
 
+/*
+ * A token to be used to fetch data from RequestContext.
+ * Generally you will want this to be a static, created only once using a
+ * string, and then only copied. The string constructor is expensive.
+ */
+class RequestToken {
+ public:
+  explicit RequestToken(const std::string& str);
+
+  bool operator==(const RequestToken& other) const {
+    return token_ == other.token_;
+  }
+
+  // Slow, use only for debug log messages.
+  std::string getDebugString() const;
+
+  friend struct std::hash<folly::RequestToken>;
+
+ private:
+  static Synchronized<std::unordered_map<std::string, uint32_t>>& getCache();
+
+  uint32_t token_;
+};
+
+} // namespace folly
+
+namespace std {
+template <>
+struct hash<folly::RequestToken> {
+  size_t operator()(const folly::RequestToken& token) const {
+    return hash<uint32_t>()(token.token_);
+  }
+};
+} // namespace std
+
+namespace folly {
+
 // Some request context that follows an async request through a process
 // Everything in the context must be thread safe
 
@@ -96,27 +133,49 @@ 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 std::string& val,
+      const RequestToken& val,
       std::unique_ptr<RequestData> data);
+  void setContextData(
+      const std::string& val,
+      std::unique_ptr<RequestData> data) {
+    setContextData(RequestToken(val), std::move(data));
+  }
 
   // Add RequestData instance "data" to this RequestContext instance, with
   // 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 std::string& val,
+      const RequestToken& val,
       std::unique_ptr<RequestData> data);
+  bool setContextDataIfAbsent(
+      const std::string& val,
+      std::unique_ptr<RequestData> data) {
+    return setContextDataIfAbsent(RequestToken(val), std::move(data));
+  }
 
   // Remove the RequestData instance with string identifier "val", if it exists.
-  void clearContextData(const std::string& val);
+  void clearContextData(const RequestToken& val);
+  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 std::string& val) const;
+  bool hasContextData(const RequestToken& val) 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 std::string& val);
-  const RequestData* getContextData(const std::string& val) const;
+  RequestData* getContextData(const RequestToken& val);
+  const RequestData* getContextData(const RequestToken& val) const;
+  RequestData* getContextData(const std::string& val) {
+    return getContextData(RequestToken(val));
+  }
+  const RequestData* getContextData(const std::string& val) const {
+    return getContextData(RequestToken(val));
+  }
 
   void onSet();
   void onUnset();
@@ -152,8 +211,13 @@ class RequestContext {
   // Similar to setContextData, except it overwrites the data
   // if already set (instead of warn + reset ptr).
   void overwriteContextData(
-      const std::string& val,
+      const RequestToken& val,
       std::unique_ptr<RequestData> data);
+  void overwriteContextData(
+      const std::string& val,
+      std::unique_ptr<RequestData> data) {
+    overwriteContextData(RequestToken(val), std::move(data));
+  }
   // End shallow copy guard
 
   enum class DoSetBehaviour {
@@ -163,14 +227,20 @@ class RequestContext {
   };
 
   bool doSetContextData(
-      const std::string& val,
+      const RequestToken& val,
       std::unique_ptr<RequestData>& data,
       DoSetBehaviour behaviour);
+  bool doSetContextData(
+      const std::string& val,
+      std::unique_ptr<RequestData>& data,
+      DoSetBehaviour behaviour) {
+    return doSetContextData(RequestToken(val), data, behaviour);
+  }
 
   struct State {
     // This must be optimized for lookup, its hot path is getContextData
     // Efficiency of copying the container also matters in setShallowCopyContext
-    F14FastMap<std::string, RequestData::SharedPtr> requestData_;
+    F14FastMap<RequestToken, RequestData::SharedPtr> requestData_;
     // This must be optimized for iteration, its hot path is setContext
     // We also use the fact that it's ordered to efficiently compute
     // the difference with previous context
@@ -226,6 +296,12 @@ struct ShallowCopyRequestContextScopeGuard {
    * Helper constructor which is a more efficient equivalent to
    * "clearRequestData" then "setRequestData" after the guard.
    */
+  ShallowCopyRequestContextScopeGuard(
+      const RequestToken& val,
+      std::unique_ptr<RequestData> data)
+      : ShallowCopyRequestContextScopeGuard() {
+    RequestContext::get()->overwriteContextData(val, std::move(data));
+  }
   ShallowCopyRequestContextScopeGuard(
       const std::string& val,
       std::unique_ptr<RequestData> data)

folly/io/async/test/RequestContextTest.cpp

@@ -23,6 +23,8 @@
 
 using namespace folly;
 
+RequestToken testtoken("test");
+
 class TestData : public RequestData {
  public:
   explicit TestData(int data) : data_(data) {}
@@ -108,9 +110,9 @@ TEST_F(RequestContextTest, SimpleTest) {
   RequestContext::get()->setContextData("test", std::make_unique<TestData>(10));
   base.runInEventBaseThread([&]() {
     EXPECT_TRUE(RequestContext::get() != nullptr);
-    auto data =
-        dynamic_cast<TestData*>(RequestContext::get()->getContextData("test"))
-            ->data_;
+    auto data = dynamic_cast<TestData*>(
+                    RequestContext::get()->getContextData(testtoken))
+                    ->data_;
     EXPECT_EQ(10, data);
     base.terminateLoopSoon();
   });
@@ -163,7 +165,7 @@ TEST_F(RequestContextTest, setIfAbsentTest) {
       "test", std::make_unique<TestData>(20)));
   EXPECT_EQ(
       10,
-      dynamic_cast<TestData*>(RequestContext::get()->getContextData("test"))
+      dynamic_cast<TestData*>(RequestContext::get()->getContextData(testtoken))
           ->data_);
 
   EXPECT_TRUE(RequestContext::get()->setContextDataIfAbsent(
@@ -189,8 +191,8 @@ TEST_F(RequestContextTest, testSetUnset) {
   // Override RequestContext
   RequestContext::create();
   auto ctx2 = RequestContext::saveContext();
-  ctx2->setContextData("test", std::make_unique<TestData>(20));
-  auto testData2 = dynamic_cast<TestData*>(ctx2->getContextData("test"));
+  ctx2->setContextData(testtoken, std::make_unique<TestData>(20));
+  auto testData2 = dynamic_cast<TestData*>(ctx2->getContextData(testtoken));
 
   // onSet called in setContextData
   EXPECT_EQ(1, testData2->set_);
@@ -229,7 +231,7 @@ TEST_F(RequestContextTest, deadlockTest) {
 
   RequestContext::get()->setContextData(
       "test", std::make_unique<DeadlockTestData>("test2"));
-  RequestContext::get()->clearContextData("test");
+  RequestContext::get()->clearContextData(testtoken);
 }
 
 // A common use case is to use set/unset to maintain a thread global