Request.cpp 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300
  1. /*
  2. * Copyright 2016-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. #include <folly/io/async/Request.h>
  17. #include <folly/tracing/StaticTracepoint.h>
  18. #include <glog/logging.h>
  19. #include <folly/MapUtil.h>
  20. #include <folly/SingletonThreadLocal.h>
  21. namespace folly {
  22. RequestToken::RequestToken(const std::string& str) {
  23. auto& cache = getCache();
  24. {
  25. auto c = cache.rlock();
  26. auto res = c->find(str);
  27. if (res != c->end()) {
  28. token_ = res->second;
  29. return;
  30. }
  31. }
  32. auto c = cache.wlock();
  33. auto res = c->find(str);
  34. if (res != c->end()) {
  35. token_ = res->second;
  36. return;
  37. }
  38. static uint32_t nextToken{1};
  39. token_ = nextToken++;
  40. (*c)[str] = token_;
  41. }
  42. std::string RequestToken::getDebugString() const {
  43. auto& cache = getCache();
  44. auto c = cache.rlock();
  45. for (auto& v : *c) {
  46. if (v.second == token_) {
  47. return v.first;
  48. }
  49. }
  50. throw std::logic_error("Could not find debug string in RequestToken");
  51. }
  52. Synchronized<std::unordered_map<std::string, uint32_t>>&
  53. RequestToken::getCache() {
  54. static Indestructible<Synchronized<std::unordered_map<std::string, uint32_t>>>
  55. cache;
  56. return *cache;
  57. }
  58. void RequestData::DestructPtr::operator()(RequestData* ptr) {
  59. if (ptr) {
  60. auto keepAliveCounter =
  61. ptr->keepAliveCounter_.fetch_sub(1, std::memory_order_acq_rel);
  62. // Note: this is the value before decrement, hence == 1 check
  63. DCHECK(keepAliveCounter > 0);
  64. if (keepAliveCounter == 1) {
  65. delete ptr;
  66. }
  67. }
  68. }
  69. /* static */ RequestData::SharedPtr RequestData::constructPtr(
  70. RequestData* ptr) {
  71. if (ptr) {
  72. auto keepAliveCounter =
  73. ptr->keepAliveCounter_.fetch_add(1, std::memory_order_relaxed);
  74. DCHECK(keepAliveCounter >= 0);
  75. }
  76. return SharedPtr(ptr);
  77. }
  78. bool RequestContext::doSetContextData(
  79. const RequestToken& val,
  80. std::unique_ptr<RequestData>& data,
  81. DoSetBehaviour behaviour) {
  82. auto ulock = state_.ulock();
  83. bool conflict = false;
  84. auto it = ulock->requestData_.find(val);
  85. if (it != ulock->requestData_.end()) {
  86. if (behaviour == DoSetBehaviour::SET_IF_ABSENT) {
  87. return false;
  88. } else if (behaviour == DoSetBehaviour::SET) {
  89. LOG_FIRST_N(WARNING, 1)
  90. << "Calling RequestContext::setContextData for "
  91. << val.getDebugString() << " but it is already set";
  92. }
  93. conflict = true;
  94. }
  95. auto wlock = ulock.moveFromUpgradeToWrite();
  96. if (conflict) {
  97. if (it->second) {
  98. if (it->second->hasCallback()) {
  99. it->second->onUnset();
  100. wlock->callbackData_.erase(it->second.get());
  101. }
  102. it->second.reset(nullptr);
  103. }
  104. if (behaviour == DoSetBehaviour::SET) {
  105. return true;
  106. }
  107. }
  108. if (data && data->hasCallback()) {
  109. wlock->callbackData_.insert(data.get());
  110. data->onSet();
  111. }
  112. wlock->requestData_[val] = RequestData::constructPtr(data.release());
  113. return true;
  114. }
  115. void RequestContext::setContextData(
  116. const RequestToken& val,
  117. std::unique_ptr<RequestData> data) {
  118. doSetContextData(val, data, DoSetBehaviour::SET);
  119. }
  120. bool RequestContext::setContextDataIfAbsent(
  121. const RequestToken& val,
  122. std::unique_ptr<RequestData> data) {
  123. return doSetContextData(val, data, DoSetBehaviour::SET_IF_ABSENT);
  124. }
  125. void RequestContext::overwriteContextData(
  126. const RequestToken& val,
  127. std::unique_ptr<RequestData> data) {
  128. doSetContextData(val, data, DoSetBehaviour::OVERWRITE);
  129. }
  130. bool RequestContext::hasContextData(const RequestToken& val) const {
  131. return state_.rlock()->requestData_.count(val);
  132. }
  133. RequestData* RequestContext::getContextData(const RequestToken& val) {
  134. const RequestData::SharedPtr dflt{nullptr};
  135. return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
  136. }
  137. const RequestData* RequestContext::getContextData(
  138. const RequestToken& val) const {
  139. const RequestData::SharedPtr dflt{nullptr};
  140. return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
  141. }
  142. void RequestContext::onSet() {
  143. auto rlock = state_.rlock();
  144. for (const auto& data : rlock->callbackData_) {
  145. data->onSet();
  146. }
  147. }
  148. void RequestContext::onUnset() {
  149. auto rlock = state_.rlock();
  150. for (const auto& data : rlock->callbackData_) {
  151. data->onUnset();
  152. }
  153. }
  154. void RequestContext::clearContextData(const RequestToken& val) {
  155. RequestData::SharedPtr requestData;
  156. // Delete the RequestData after giving up the wlock just in case one of the
  157. // RequestData destructors will try to grab the lock again.
  158. {
  159. auto ulock = state_.ulock();
  160. auto it = ulock->requestData_.find(val);
  161. if (it == ulock->requestData_.end()) {
  162. return;
  163. }
  164. auto wlock = ulock.moveFromUpgradeToWrite();
  165. if (it->second && it->second->hasCallback()) {
  166. it->second->onUnset();
  167. wlock->callbackData_.erase(it->second.get());
  168. }
  169. requestData = std::move(it->second);
  170. wlock->requestData_.erase(it);
  171. }
  172. }
  173. namespace {
  174. // Execute functor exec for all RequestData in data, which are not in other
  175. // Similar to std::set_difference but avoid intermediate data structure
  176. template <typename TData, typename TExec>
  177. void exec_set_difference(const TData& data, const TData& other, TExec&& exec) {
  178. auto diter = data.begin();
  179. auto dend = data.end();
  180. auto oiter = other.begin();
  181. auto oend = other.end();
  182. while (diter != dend) {
  183. // Order of "if" optimizes for the 2 common cases:
  184. // 1) empty other, switching to default context
  185. // 2) identical other, switching to similar context with same callbacks
  186. if (oiter == oend) {
  187. exec(*diter);
  188. ++diter;
  189. } else if (*diter == *oiter) {
  190. ++diter;
  191. ++oiter;
  192. } else if (*diter < *oiter) {
  193. exec(*diter);
  194. ++diter;
  195. } else {
  196. ++oiter;
  197. }
  198. }
  199. }
  200. } // namespace
  201. std::shared_ptr<RequestContext> RequestContext::setContext(
  202. std::shared_ptr<RequestContext> newCtx) {
  203. auto& staticCtx = getStaticContext();
  204. if (newCtx == staticCtx) {
  205. return newCtx;
  206. }
  207. FOLLY_SDT(
  208. folly, request_context_switch_before, staticCtx.get(), newCtx.get());
  209. auto curCtx = staticCtx;
  210. if (newCtx && curCtx) {
  211. // Only call set/unset for all request data that differs
  212. auto ret = folly::acquireLocked(
  213. as_const(newCtx->state_), as_const(curCtx->state_));
  214. auto& newLock = std::get<0>(ret);
  215. auto& curLock = std::get<1>(ret);
  216. auto& newData = newLock->callbackData_;
  217. auto& curData = curLock->callbackData_;
  218. exec_set_difference(
  219. curData, newData, [](RequestData* data) { data->onUnset(); });
  220. staticCtx = newCtx;
  221. exec_set_difference(
  222. newData, curData, [](RequestData* data) { data->onSet(); });
  223. } else {
  224. if (curCtx) {
  225. curCtx->onUnset();
  226. }
  227. staticCtx = newCtx;
  228. if (newCtx) {
  229. newCtx->onSet();
  230. }
  231. }
  232. return curCtx;
  233. }
  234. std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
  235. using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
  236. return SingletonT::get();
  237. }
  238. /* static */ std::shared_ptr<RequestContext>
  239. RequestContext::setShallowCopyContext() {
  240. auto& parent = getStaticContext();
  241. auto child = std::make_shared<RequestContext>();
  242. if (parent) {
  243. auto ret = folly::acquireLocked(as_const(parent->state_), child->state_);
  244. auto& parentLock = std::get<0>(ret);
  245. auto& childLock = std::get<1>(ret);
  246. childLock->callbackData_ = parentLock->callbackData_;
  247. childLock->requestData_.reserve(parentLock->requestData_.size());
  248. for (const auto& entry : parentLock->requestData_) {
  249. childLock->requestData_.insert(std::make_pair(
  250. entry.first, RequestData::constructPtr(entry.second.get())));
  251. }
  252. }
  253. // Do not use setContext to avoid global set/unset
  254. std::swap(child, parent);
  255. return child;
  256. }
  257. RequestContext* RequestContext::get() {
  258. auto& context = getStaticContext();
  259. if (!context) {
  260. static RequestContext defaultContext;
  261. return std::addressof(defaultContext);
  262. }
  263. return context.get();
  264. }
  265. } // namespace folly