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