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