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 void RequestContext::setContextData(
  28     const std::string& val,
  29     std::unique_ptr<RequestData> data) {
  30   auto wlock = data_.wlock();
  31   if (wlock->count(val)) {
  32     LOG_FIRST_N(WARNING, 1)
  33         << "Called RequestContext::setContextData with data already set";
  34
  35     (*wlock)[val] = nullptr;
  36   } else {
  37     (*wlock)[val] = std::move(data);
  38   }
  39 }
  40
  41 bool RequestContext::setContextDataIfAbsent(
  42     const std::string& val,
  43     std::unique_ptr<RequestData> data) {
  44   auto ulock = data_.ulock();
  45   if (ulock->count(val)) {
  46     return false;
  47   }
  48
  49   auto wlock = ulock.moveFromUpgradeToWrite();
  50   (*wlock)[val] = std::move(data);
  51   return true;
  52 }
  53
  54 bool RequestContext::hasContextData(const std::string& val) const {
  55   return data_.rlock()->count(val);
  56 }
  57
  58 RequestData* RequestContext::getContextData(const std::string& val) {
  59   const std::unique_ptr<RequestData> dflt{nullptr};
  60   return get_ref_default(*data_.rlock(), val, dflt).get();
  61 }
  62
  63 const RequestData* RequestContext::getContextData(
  64     const std::string& val) const {
  65   const std::unique_ptr<RequestData> dflt{nullptr};
  66   return get_ref_default(*data_.rlock(), val, dflt).get();
  67 }
  68
  69 void RequestContext::onSet() {
  70   auto rlock = data_.rlock();
  71   for (auto const& ent : *rlock) {
  72     if (auto& data = ent.second) {
  73       data->onSet();
  74     }
  75   }
  76 }
  77
  78 void RequestContext::onUnset() {
  79   auto rlock = data_.rlock();
  80   for (auto const& ent : *rlock) {
  81     if (auto& data = ent.second) {
  82       data->onUnset();
  83     }
  84   }
  85 }
  86
  87 void RequestContext::clearContextData(const std::string& val) {
  88   std::unique_ptr<RequestData> requestData;
  89   // Delete the RequestData after giving up the wlock just in case one of the
  90   // RequestData destructors will try to grab the lock again.
  91   {
  92     auto wlock = data_.wlock();
  93     auto it = wlock->find(val);
  94     if (it != wlock->end()) {
  95       requestData = std::move(it->second);
  96       wlock->erase(it);
  97     }
  98   }
  99 }
 100
 101 std::shared_ptr<RequestContext> RequestContext::setContext(
 102     std::shared_ptr<RequestContext> ctx) {
 103   auto& curCtx = getStaticContext();
 104   if (ctx != curCtx) {
 105     FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
 106     using std::swap;
 107     if (curCtx) {
 108       curCtx->onUnset();
 109     }
 110     swap(ctx, curCtx);
 111     if (curCtx) {
 112       curCtx->onSet();
 113     }
 114   }
 115   return ctx;
 116 }
 117
 118 std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
 119   using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
 120   static SingletonT singleton;
 121
 122   return singleton.get();
 123 }
 124
 125 RequestContext* RequestContext::get() {
 126   auto context = getStaticContext();
 127   if (!context) {
 128     static RequestContext defaultContext;
 129     return std::addressof(defaultContext);
 130   }
 131   return context.get();
 132 }
 133 }