folly/io/async/Request.cpp

   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 #include <folly/io/async/Request.h>
  18 #include <folly/tracing/StaticTracepoint.h>
  19
  20 #include <glog/logging.h>
  21
  22 #include <folly/MapUtil.h>
  23 #include <folly/SingletonThreadLocal.h>
  24
  25 namespace folly {
  26
  27 bool RequestContext::doSetContextData(
  28     const std::string& val,
  29     std::unique_ptr<RequestData>& data,
  30     bool strict) {
  31   auto ulock = state_.ulock();
  32
  33   bool conflict = false;
  34   auto it = ulock->requestData_.find(val);
  35   if (it != ulock->requestData_.end()) {
  36     if (strict) {
  37       return false;
  38     } else {
  39       LOG_FIRST_N(WARNING, 1) << "Calling RequestContext::setContextData for "
  40                               << val << " but it is already set";
  41       conflict = true;
  42     }
  43   }
  44
  45   auto wlock = ulock.moveFromUpgradeToWrite();
  46   if (conflict) {
  47     if (it->second) {
  48       if (it->second->hasCallback()) {
  49         wlock->callbackData_.erase(it->second.get());
  50       }
  51       it->second.reset(nullptr);
  52     }
  53     return true;
  54   }
  55
  56   if (data && data->hasCallback()) {
  57     wlock->callbackData_.insert(data.get());
  58   }
  59   wlock->requestData_[val] = std::move(data);
  60
  61   return true;
  62 }
  63
  64 void RequestContext::setContextData(
  65     const std::string& val,
  66     std::unique_ptr<RequestData> data) {
  67   doSetContextData(val, data, false /* strict */);
  68 }
  69
  70 bool RequestContext::setContextDataIfAbsent(
  71     const std::string& val,
  72     std::unique_ptr<RequestData> data) {
  73   return doSetContextData(val, data, true /* strict */);
  74 }
  75
  76 bool RequestContext::hasContextData(const std::string& val) const {
  77   return state_.rlock()->requestData_.count(val);
  78 }
  79
  80 RequestData* RequestContext::getContextData(const std::string& val) {
  81   const std::unique_ptr<RequestData> dflt{nullptr};
  82   return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
  83 }
  84
  85 const RequestData* RequestContext::getContextData(
  86     const std::string& val) const {
  87   const std::unique_ptr<RequestData> dflt{nullptr};
  88   return get_ref_default(state_.rlock()->requestData_, val, dflt).get();
  89 }
  90
  91 void RequestContext::onSet() {
  92   auto rlock = state_.rlock();
  93   for (const auto& data : rlock->callbackData_) {
  94     data->onSet();
  95   }
  96 }
  97
  98 void RequestContext::onUnset() {
  99   auto rlock = state_.rlock();
 100   for (const auto& data : rlock->callbackData_) {
 101     data->onUnset();
 102   }
 103 }
 104
 105 void RequestContext::clearContextData(const std::string& val) {
 106   std::unique_ptr<RequestData> requestData;
 107   // Delete the RequestData after giving up the wlock just in case one of the
 108   // RequestData destructors will try to grab the lock again.
 109   {
 110     auto ulock = state_.ulock();
 111     auto it = ulock->requestData_.find(val);
 112     if (it == ulock->requestData_.end()) {
 113       return;
 114     }
 115
 116     auto wlock = ulock.moveFromUpgradeToWrite();
 117     if (it->second && it->second->hasCallback()) {
 118       wlock->callbackData_.erase(it->second.get());
 119     }
 120
 121     requestData = std::move(it->second);
 122     wlock->requestData_.erase(it);
 123   }
 124 }
 125
 126 std::shared_ptr<RequestContext> RequestContext::setContext(
 127     std::shared_ptr<RequestContext> ctx) {
 128   auto& curCtx = getStaticContext();
 129   if (ctx != curCtx) {
 130     FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
 131     using std::swap;
 132     if (curCtx) {
 133       curCtx->onUnset();
 134     }
 135     swap(ctx, curCtx);
 136     if (curCtx) {
 137       curCtx->onSet();
 138     }
 139   }
 140   return ctx;
 141 }
 142
 143 std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
 144   using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
 145   static SingletonT singleton;
 146
 147   return singleton.get();
 148 }
 149
 150 RequestContext* RequestContext::get() {
 151   auto& context = getStaticContext();
 152   if (!context) {
 153     static RequestContext defaultContext;
 154     return std::addressof(defaultContext);
 155   }
 156   return context.get();
 157 }
 158 } // namespace folly