folly/io/async/Request.cpp

   1 /*
   2  * Copyright 2016 Facebook, Inc.
   3  *
   4  * Licensed to the Apache Software Foundation (ASF) under one
   5  * or more contributor license agreements. See the NOTICE file
   6  * distributed with this work for additional information
   7  * regarding copyright ownership. The ASF licenses this file
   8  * to you under the Apache License, Version 2.0 (the
   9  * "License"); you may not use this file except in compliance
  10  * with the License. You may obtain a copy of the License at
  11  *
  12  *   http://www.apache.org/licenses/LICENSE-2.0
  13  *
  14  * Unless required by applicable law or agreed to in writing,
  15  * software distributed under the License is distributed on an
  16  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  17  * KIND, either express or implied. See the License for the
  18  * specific language governing permissions and limitations
  19  * under the License.
  20  */
  21
  22 #include <folly/io/async/Request.h>
  23 #include <folly/tracing/StaticTracepoint.h>
  24
  25 #include <glog/logging.h>
  26
  27 #include <folly/MapUtil.h>
  28 #include <folly/SingletonThreadLocal.h>
  29
  30 namespace folly {
  31
  32 void RequestContext::setContextData(
  33     const std::string& val,
  34     std::unique_ptr<RequestData> data) {
  35   auto wlock = data_.wlock();
  36   if (wlock->count(val)) {
  37     LOG_FIRST_N(WARNING, 1)
  38         << "Called RequestContext::setContextData with data already set";
  39
  40     (*wlock)[val] = nullptr;
  41   } else {
  42     (*wlock)[val] = std::move(data);
  43   }
  44 }
  45
  46 bool RequestContext::setContextDataIfAbsent(
  47     const std::string& val,
  48     std::unique_ptr<RequestData> data) {
  49   auto ulock = data_.ulock();
  50   if (ulock->count(val)) {
  51     return false;
  52   }
  53
  54   auto wlock = ulock.moveFromUpgradeToWrite();
  55   (*wlock)[val] = std::move(data);
  56   return true;
  57 }
  58
  59 bool RequestContext::hasContextData(const std::string& val) const {
  60   return data_.rlock()->count(val);
  61 }
  62
  63 RequestData* RequestContext::getContextData(const std::string& val) {
  64   return get_ref_default(*data_.rlock(), val, nullptr).get();
  65 }
  66
  67 const RequestData* RequestContext::getContextData(
  68     const std::string& val) const {
  69   return get_ref_default(*data_.rlock(), val, nullptr).get();
  70 }
  71
  72 void RequestContext::onSet() {
  73   auto rlock = data_.rlock();
  74   for (auto const& ent : *rlock) {
  75     if (auto& data = ent.second) {
  76       data->onSet();
  77     }
  78   }
  79 }
  80
  81 void RequestContext::onUnset() {
  82   auto rlock = data_.rlock();
  83   for (auto const& ent : *rlock) {
  84     if (auto& data = ent.second) {
  85       data->onUnset();
  86     }
  87   }
  88 }
  89
  90 void RequestContext::clearContextData(const std::string& val) {
  91   data_.wlock()->erase(val);
  92 }
  93
  94 std::shared_ptr<RequestContext> RequestContext::setContext(
  95     std::shared_ptr<RequestContext> ctx) {
  96   auto& curCtx = getStaticContext();
  97   if (ctx != curCtx) {
  98     FOLLY_SDT(folly, request_context_switch_before, curCtx.get(), ctx.get());
  99     using std::swap;
 100     if (curCtx) {
 101       curCtx->onUnset();
 102     }
 103     swap(ctx, curCtx);
 104     if (curCtx) {
 105       curCtx->onSet();
 106     }
 107   }
 108   return ctx;
 109 }
 110
 111 std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
 112   using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
 113   static SingletonT singleton;
 114
 115   return singleton.get();
 116 }
 117
 118 RequestContext* RequestContext::get() {
 119   auto context = getStaticContext();
 120   if (!context) {
 121     static RequestContext defaultContext;
 122     return std::addressof(defaultContext);
 123   }
 124   return context.get();
 125 }
 126 }