Improve folly::RequestContext onSet and onUnset efficiency

[folly.git] / folly / io / async / Request.h

diff --git a/folly/io/async/Request.h b/folly/io/async/Request.h
index ce733301345303e319790e29d1af0be031f57d7a..7398acb2c51161bd4743c62d504735dd53675217 100644 (file)
--- a/folly/io/async/Request.h
+++ b/folly/io/async/Request.h
@@ -18,8 +18,8 @@
 
 #include <map>
 #include <memory>
+#include <set>
 
-#include <folly/SharedMutex.h>
 #include <folly/Synchronized.h>
 
 namespace folly {
@@ -30,16 +30,23 @@ namespace folly {
 class RequestData {
  public:
   virtual ~RequestData() = default;
+
   // Avoid calling RequestContext::setContextData, setContextDataIfAbsent, or
   // clearContextData from these callbacks. Doing so will cause deadlock. We
   // could fix these deadlocks, but only at significant performance penalty, so
   // just don't do it!
+
+  virtual bool hasCallback() = 0;
+  // Callback executed when setting RequestContext. Make sure your RequestData
+  // instance overrides the hasCallback method to return true otherwise
+  // the callback will not be executed
   virtual void onSet() {}
+  // Callback executed when unsetting RequestContext. Make sure your RequestData
+  // instance overrides the hasCallback method to return true otherwise
+  // the callback will not be executed
   virtual void onUnset() {}
 };
 
-class RequestContext;
-
 // If you do not call create() to create a unique request context,
 // this default request context will always be returned, and is never
 // copied between threads.
@@ -55,29 +62,42 @@ class RequestContext {
   // Get the current context.
   static RequestContext* get();
 
-  // The following API may be used to set per-request data in a thread-safe way.
-  // This access is still performance sensitive, so please ask if you need help
-  // profiling any use of these functions.
+  // The following APIs are used to add, remove and access RequestData instance
+  // in the RequestContext instance, normally used for per-RequestContext
+  // tracking or callback on set and unset. These APIs are Thread-safe.
+  // These APIs are performance sensitive, so please ask if you need help
+  // profiling any use of these APIs.
+
+  // Add RequestData instance "data" to this RequestContext instance, with
+  // string identifier "val". If the same string identifier has already been
+  // 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,
       std::unique_ptr<RequestData> data);
 
-  // Unlike setContextData, this method does not panic if the key is already
-  // present. Returns true iff the new value has been inserted.
+  // 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,
       std::unique_ptr<RequestData> data);
 
+  // Remove the RequestData instance with string identifier "val", if it exists.
+  void clearContextData(const std::string& 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;
 
+  // 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;
 
   void onSet();
   void onUnset();
 
-  void clearContextData(const std::string& val);
-
   // 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.
@@ -98,8 +118,16 @@ class RequestContext {
  private:
   static std::shared_ptr<RequestContext>& getStaticContext();
 
-  using Data = std::map<std::string, std::unique_ptr<RequestData>>;
-  folly::Synchronized<Data, folly::SharedMutex> data_;
+  bool doSetContextData(
+      const std::string& val,
+      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_;
 };
 
 class RequestContextScopeGuard {
@@ -121,8 +149,7 @@ class RequestContextScopeGuard {
   // Set a RequestContext that was previously captured by saveContext(). It will
   // be automatically reset to the original value when this goes out of scope.
   explicit RequestContextScopeGuard(std::shared_ptr<RequestContext> ctx)
-      : prev_(RequestContext::setContext(std::move(ctx))) {
-  }
+      : prev_(RequestContext::setContext(std::move(ctx))) {}
 
   ~RequestContextScopeGuard() {
     RequestContext::setContext(std::move(prev_));