Request.h 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329
  1. /*
  2. * Copyright 2014-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. #pragma once
  17. #include <memory>
  18. #include <string>
  19. #include <folly/Synchronized.h>
  20. #include <folly/container/F14Map.h>
  21. #include <folly/sorted_vector_types.h>
  22. namespace folly {
  23. /*
  24. * A token to be used to fetch data from RequestContext.
  25. * Generally you will want this to be a static, created only once using a
  26. * string, and then only copied. The string constructor is expensive.
  27. */
  28. class RequestToken {
  29. public:
  30. explicit RequestToken(const std::string& str);
  31. bool operator==(const RequestToken& other) const {
  32. return token_ == other.token_;
  33. }
  34. // Slow, use only for debug log messages.
  35. std::string getDebugString() const;
  36. friend struct std::hash<folly::RequestToken>;
  37. private:
  38. static Synchronized<std::unordered_map<std::string, uint32_t>>& getCache();
  39. uint32_t token_;
  40. };
  41. } // namespace folly
  42. namespace std {
  43. template <>
  44. struct hash<folly::RequestToken> {
  45. size_t operator()(const folly::RequestToken& token) const {
  46. return hash<uint32_t>()(token.token_);
  47. }
  48. };
  49. } // namespace std
  50. namespace folly {
  51. // Some request context that follows an async request through a process
  52. // Everything in the context must be thread safe
  53. class RequestData {
  54. public:
  55. virtual ~RequestData() = default;
  56. // Avoid calling RequestContext::setContextData, setContextDataIfAbsent, or
  57. // clearContextData from these callbacks. Doing so will cause deadlock. We
  58. // could fix these deadlocks, but only at significant performance penalty, so
  59. // just don't do it!
  60. virtual bool hasCallback() = 0;
  61. // Callback executed when setting RequestContext. Make sure your RequestData
  62. // instance overrides the hasCallback method to return true otherwise
  63. // the callback will not be executed
  64. virtual void onSet() {}
  65. // Callback executed when unsetting RequestContext. Make sure your RequestData
  66. // instance overrides the hasCallback method to return true otherwise
  67. // the callback will not be executed
  68. virtual void onUnset() {}
  69. private:
  70. // Start shallow copy implementation details:
  71. // For efficiency, RequestContext provides a raw ptr interface.
  72. // To support shallow copy, we need a shared ptr.
  73. // To keep it as safe as possible (even if a raw ptr is passed back),
  74. // the counter lives directly in RequestData.
  75. friend class RequestContext;
  76. // Unique ptr with custom destructor, decrement the counter
  77. // and only free if 0
  78. struct DestructPtr {
  79. void operator()(RequestData* ptr);
  80. };
  81. using SharedPtr = std::unique_ptr<RequestData, DestructPtr>;
  82. // Initialize the pseudo-shared ptr, increment the counter
  83. static SharedPtr constructPtr(RequestData* ptr);
  84. std::atomic<int> keepAliveCounter_{0};
  85. // End shallow copy
  86. };
  87. // If you do not call create() to create a unique request context,
  88. // this default request context will always be returned, and is never
  89. // copied between threads.
  90. class RequestContext {
  91. public:
  92. // Create a unique request context for this request.
  93. // It will be passed between queues / threads (where implemented),
  94. // so it should be valid for the lifetime of the request.
  95. static void create() {
  96. setContext(std::make_shared<RequestContext>());
  97. }
  98. // Get the current context.
  99. static RequestContext* get();
  100. // The following APIs are used to add, remove and access RequestData instance
  101. // in the RequestContext instance, normally used for per-RequestContext
  102. // tracking or callback on set and unset. These APIs are Thread-safe.
  103. // These APIs are performance sensitive, so please ask if you need help
  104. // profiling any use of these APIs.
  105. // Add RequestData instance "data" to this RequestContext instance, with
  106. // string identifier "val". If the same string identifier has already been
  107. // used, will print a warning message for the first time, clear the existing
  108. // RequestData instance for "val", and **not** add "data".
  109. void setContextData(
  110. const RequestToken& val,
  111. std::unique_ptr<RequestData> data);
  112. void setContextData(
  113. const std::string& val,
  114. std::unique_ptr<RequestData> data) {
  115. setContextData(RequestToken(val), std::move(data));
  116. }
  117. // Add RequestData instance "data" to this RequestContext instance, with
  118. // string identifier "val". If the same string identifier has already been
  119. // used, return false and do nothing. Otherwise add "data" and return true.
  120. bool setContextDataIfAbsent(
  121. const RequestToken& val,
  122. std::unique_ptr<RequestData> data);
  123. bool setContextDataIfAbsent(
  124. const std::string& val,
  125. std::unique_ptr<RequestData> data) {
  126. return setContextDataIfAbsent(RequestToken(val), std::move(data));
  127. }
  128. // Remove the RequestData instance with string identifier "val", if it exists.
  129. void clearContextData(const RequestToken& val);
  130. void clearContextData(const std::string& val) {
  131. clearContextData(RequestToken(val));
  132. }
  133. // Returns true if and only if the RequestData instance with string identifier
  134. // "val" exists in this RequestContext instnace.
  135. bool hasContextData(const RequestToken& val) const;
  136. bool hasContextData(const std::string& val) const {
  137. return hasContextData(RequestToken(val));
  138. }
  139. // Get (constant) raw pointer of the RequestData instance with string
  140. // identifier "val" if it exists, otherwise returns null pointer.
  141. RequestData* getContextData(const RequestToken& val);
  142. const RequestData* getContextData(const RequestToken& val) const;
  143. RequestData* getContextData(const std::string& val) {
  144. return getContextData(RequestToken(val));
  145. }
  146. const RequestData* getContextData(const std::string& val) const {
  147. return getContextData(RequestToken(val));
  148. }
  149. void onSet();
  150. void onUnset();
  151. // The following API is used to pass the context through queues / threads.
  152. // saveContext is called to get a shared_ptr to the context, and
  153. // setContext is used to reset it on the other side of the queue.
  154. //
  155. // Whenever possible, use RequestContextScopeGuard instead of setContext
  156. // to make sure that RequestContext is reset to the original value when
  157. // we exit the scope.
  158. //
  159. // A shared_ptr is used, because many request may fan out across
  160. // multiple threads, or do post-send processing, etc.
  161. static std::shared_ptr<RequestContext> setContext(
  162. std::shared_ptr<RequestContext> ctx);
  163. static std::shared_ptr<RequestContext> saveContext() {
  164. return getStaticContext();
  165. }
  166. private:
  167. static std::shared_ptr<RequestContext>& getStaticContext();
  168. // Start shallow copy guard implementation details:
  169. // All methods are private to encourage proper use
  170. friend struct ShallowCopyRequestContextScopeGuard;
  171. // This sets a shallow copy of the current context as current,
  172. // then return the previous context (so it can be reset later).
  173. static std::shared_ptr<RequestContext> setShallowCopyContext();
  174. // Similar to setContextData, except it overwrites the data
  175. // if already set (instead of warn + reset ptr).
  176. void overwriteContextData(
  177. const RequestToken& val,
  178. std::unique_ptr<RequestData> data);
  179. void overwriteContextData(
  180. const std::string& val,
  181. std::unique_ptr<RequestData> data) {
  182. overwriteContextData(RequestToken(val), std::move(data));
  183. }
  184. // End shallow copy guard
  185. enum class DoSetBehaviour {
  186. SET,
  187. SET_IF_ABSENT,
  188. OVERWRITE,
  189. };
  190. bool doSetContextData(
  191. const RequestToken& val,
  192. std::unique_ptr<RequestData>& data,
  193. DoSetBehaviour behaviour);
  194. bool doSetContextData(
  195. const std::string& val,
  196. std::unique_ptr<RequestData>& data,
  197. DoSetBehaviour behaviour) {
  198. return doSetContextData(RequestToken(val), data, behaviour);
  199. }
  200. struct State {
  201. // This must be optimized for lookup, its hot path is getContextData
  202. // Efficiency of copying the container also matters in setShallowCopyContext
  203. F14FastMap<RequestToken, RequestData::SharedPtr> requestData_;
  204. // This must be optimized for iteration, its hot path is setContext
  205. // We also use the fact that it's ordered to efficiently compute
  206. // the difference with previous context
  207. sorted_vector_set<RequestData*> callbackData_;
  208. };
  209. folly::Synchronized<State> state_;
  210. };
  211. /**
  212. * Note: you probably want to use ShallowCopyRequestContextScopeGuard
  213. * This resets all other RequestData for the duration of the scope!
  214. */
  215. class RequestContextScopeGuard {
  216. private:
  217. std::shared_ptr<RequestContext> prev_;
  218. public:
  219. RequestContextScopeGuard(const RequestContextScopeGuard&) = delete;
  220. RequestContextScopeGuard& operator=(const RequestContextScopeGuard&) = delete;
  221. RequestContextScopeGuard(RequestContextScopeGuard&&) = delete;
  222. RequestContextScopeGuard& operator=(RequestContextScopeGuard&&) = delete;
  223. // Create a new RequestContext and reset to the original value when
  224. // this goes out of scope.
  225. RequestContextScopeGuard() : prev_(RequestContext::saveContext()) {
  226. RequestContext::create();
  227. }
  228. // Set a RequestContext that was previously captured by saveContext(). It will
  229. // be automatically reset to the original value when this goes out of scope.
  230. explicit RequestContextScopeGuard(std::shared_ptr<RequestContext> ctx)
  231. : prev_(RequestContext::setContext(std::move(ctx))) {}
  232. ~RequestContextScopeGuard() {
  233. RequestContext::setContext(std::move(prev_));
  234. }
  235. };
  236. /**
  237. * This guard maintains all the RequestData pointers of the parent.
  238. * This allows to overwrite a specific RequestData pointer for the
  239. * scope's duration, without breaking others.
  240. *
  241. * Only modified pointers will have their set/onset methods called
  242. */
  243. struct ShallowCopyRequestContextScopeGuard {
  244. ShallowCopyRequestContextScopeGuard()
  245. : prev_(RequestContext::setShallowCopyContext()) {}
  246. /**
  247. * Shallow copy then overwrite one specific RequestData
  248. *
  249. * Helper constructor which is a more efficient equivalent to
  250. * "clearRequestData" then "setRequestData" after the guard.
  251. */
  252. ShallowCopyRequestContextScopeGuard(
  253. const RequestToken& val,
  254. std::unique_ptr<RequestData> data)
  255. : ShallowCopyRequestContextScopeGuard() {
  256. RequestContext::get()->overwriteContextData(val, std::move(data));
  257. }
  258. ShallowCopyRequestContextScopeGuard(
  259. const std::string& val,
  260. std::unique_ptr<RequestData> data)
  261. : ShallowCopyRequestContextScopeGuard() {
  262. RequestContext::get()->overwriteContextData(val, std::move(data));
  263. }
  264. ~ShallowCopyRequestContextScopeGuard() {
  265. RequestContext::setContext(std::move(prev_));
  266. }
  267. ShallowCopyRequestContextScopeGuard(
  268. const ShallowCopyRequestContextScopeGuard&) = delete;
  269. ShallowCopyRequestContextScopeGuard& operator=(
  270. const ShallowCopyRequestContextScopeGuard&) = delete;
  271. ShallowCopyRequestContextScopeGuard(ShallowCopyRequestContextScopeGuard&&) =
  272. delete;
  273. ShallowCopyRequestContextScopeGuard& operator=(
  274. ShallowCopyRequestContextScopeGuard&&) = delete;
  275. private:
  276. std::shared_ptr<RequestContext> prev_;
  277. };
  278. } // namespace folly