folly/io/async/Request.h

   1 /*
   2  * Copyright 2004-present Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #pragma once
  18
  19 #include <map>
  20 #include <memory>
  21 #include <set>
  22
  23 #include <folly/Synchronized.h>
  24
  25 namespace folly {
  26
  27 // Some request context that follows an async request through a process
  28 // Everything in the context must be thread safe
  29
  30 class RequestData {
  31  public:
  32   virtual ~RequestData() = default;
  33
  34   // Avoid calling RequestContext::setContextData, setContextDataIfAbsent, or
  35   // clearContextData from these callbacks. Doing so will cause deadlock. We
  36   // could fix these deadlocks, but only at significant performance penalty, so
  37   // just don't do it!
  38
  39   virtual bool hasCallback() = 0;
  40   // Callback executed when setting RequestContext. Make sure your RequestData
  41   // instance overrides the hasCallback method to return true otherwise
  42   // the callback will not be executed
  43   virtual void onSet() {}
  44   // Callback executed when unsetting RequestContext. Make sure your RequestData
  45   // instance overrides the hasCallback method to return true otherwise
  46   // the callback will not be executed
  47   virtual void onUnset() {}
  48 };
  49
  50 // If you do not call create() to create a unique request context,
  51 // this default request context will always be returned, and is never
  52 // copied between threads.
  53 class RequestContext {
  54  public:
  55   // Create a unique request context for this request.
  56   // It will be passed between queues / threads (where implemented),
  57   // so it should be valid for the lifetime of the request.
  58   static void create() {
  59     setContext(std::make_shared<RequestContext>());
  60   }
  61
  62   // Get the current context.
  63   static RequestContext* get();
  64
  65   // The following APIs are used to add, remove and access RequestData instance
  66   // in the RequestContext instance, normally used for per-RequestContext
  67   // tracking or callback on set and unset. These APIs are Thread-safe.
  68   // These APIs are performance sensitive, so please ask if you need help
  69   // profiling any use of these APIs.
  70
  71   // Add RequestData instance "data" to this RequestContext instance, with
  72   // string identifier "val". If the same string identifier has already been
  73   // used, will print a warning message for the first time, clear the existing
  74   // RequestData instance for "val", and **not** add "data".
  75   void setContextData(
  76       const std::string& val,
  77       std::unique_ptr<RequestData> data);
  78
  79   // Add RequestData instance "data" to this RequestContext instance, with
  80   // string identifier "val". If the same string identifier has already been
  81   // used, return false and do nothing. Otherwise add "data" and return true.
  82   bool setContextDataIfAbsent(
  83       const std::string& val,
  84       std::unique_ptr<RequestData> data);
  85
  86   // Remove the RequestData instance with string identifier "val", if it exists.
  87   void clearContextData(const std::string& val);
  88
  89   // Returns true if and only if the RequestData instance with string identifier
  90   // "val" exists in this RequestContext instnace.
  91   bool hasContextData(const std::string& val) const;
  92
  93   // Get (constant) raw pointer of the RequestData instance with string
  94   // identifier "val" if it exists, otherwise returns null pointer.
  95   RequestData* getContextData(const std::string& val);
  96   const RequestData* getContextData(const std::string& val) const;
  97
  98   void onSet();
  99   void onUnset();
 100
 101   // The following API is used to pass the context through queues / threads.
 102   // saveContext is called to get a shared_ptr to the context, and
 103   // setContext is used to reset it on the other side of the queue.
 104   //
 105   // Whenever possible, use RequestContextScopeGuard instead of setContext
 106   // to make sure that RequestContext is reset to the original value when
 107   // we exit the scope.
 108   //
 109   // A shared_ptr is used, because many request may fan out across
 110   // multiple threads, or do post-send processing, etc.
 111   static std::shared_ptr<RequestContext> setContext(
 112       std::shared_ptr<RequestContext> ctx);
 113
 114   static std::shared_ptr<RequestContext> saveContext() {
 115     return getStaticContext();
 116   }
 117
 118  private:
 119   static std::shared_ptr<RequestContext>& getStaticContext();
 120
 121   bool doSetContextData(
 122       const std::string& val,
 123       std::unique_ptr<RequestData>& data,
 124       bool strict);
 125
 126   struct State {
 127     std::map<std::string, std::unique_ptr<RequestData>> requestData_;
 128     std::set<RequestData*> callbackData_;
 129   };
 130   folly::Synchronized<State> state_;
 131 };
 132
 133 class RequestContextScopeGuard {
 134  private:
 135   std::shared_ptr<RequestContext> prev_;
 136
 137  public:
 138   RequestContextScopeGuard(const RequestContextScopeGuard&) = delete;
 139   RequestContextScopeGuard& operator=(const RequestContextScopeGuard&) = delete;
 140   RequestContextScopeGuard(RequestContextScopeGuard&&) = delete;
 141   RequestContextScopeGuard& operator=(RequestContextScopeGuard&&) = delete;
 142
 143   // Create a new RequestContext and reset to the original value when
 144   // this goes out of scope.
 145   RequestContextScopeGuard() : prev_(RequestContext::saveContext()) {
 146     RequestContext::create();
 147   }
 148
 149   // Set a RequestContext that was previously captured by saveContext(). It will
 150   // be automatically reset to the original value when this goes out of scope.
 151   explicit RequestContextScopeGuard(std::shared_ptr<RequestContext> ctx)
 152       : prev_(RequestContext::setContext(std::move(ctx))) {}
 153
 154   ~RequestContextScopeGuard() {
 155     RequestContext::setContext(std::move(prev_));
 156   }
 157 };
 158 } // namespace folly