folly/io/async/Request.h

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one
   3  * or more contributor license agreements. See the NOTICE file
   4  * distributed with this work for additional information
   5  * regarding copyright ownership. The ASF licenses this file
   6  * to you under the Apache License, Version 2.0 (the
   7  * "License"); you may not use this file except in compliance
   8  * with the License. You may obtain a copy of the License at
   9  *
  10  *   http://www.apache.org/licenses/LICENSE-2.0
  11  *
  12  * Unless required by applicable law or agreed to in writing,
  13  * software distributed under the License is distributed on an
  14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15  * KIND, either express or implied. See the License for the
  16  * specific language governing permissions and limitations
  17  * under the License.
  18  */
  19 #pragma once
  20
  21 #include <map>
  22 #include <memory>
  23 #include <glog/logging.h>
  24 #include "folly/ThreadLocal.h"
  25 #include "folly/RWSpinLock.h"
  26
  27 /**
  28  * In many cases this header is included as a
  29  * dependency to libraries which do not need
  30  * command line flags. GFLAGS is a large binary
  31  * and thus we do this so that a library which
  32  * is size sensitive doesn't have to pull in
  33  * GFLAGS if it doesn't want to.
  34  */
  35 #ifndef NO_LIB_GFLAGS
  36   #include <gflags/gflags.h>
  37   DECLARE_bool(enable_request_context);
  38 #endif
  39
  40 namespace folly {
  41
  42 #ifdef NO_LIB_GFLAGS
  43   extern bool FLAGS_enable_request_context;
  44 #endif
  45
  46 // Some request context that follows an async request through a process
  47 // Everything in the context must be thread safe
  48
  49 class RequestData {
  50  public:
  51   virtual ~RequestData() {}
  52 };
  53
  54 class RequestContext;
  55
  56 // If you do not call create() to create a unique request context,
  57 // this default request context will always be returned, and is never
  58 // copied between threads.
  59 extern RequestContext* defaultContext;
  60
  61 class RequestContext {
  62  public:
  63   // Create a unique requext context for this request.
  64   // It will be passed between queues / threads (where implemented),
  65   // so it should be valid for the lifetime of the request.
  66   static bool create() {
  67     if(!FLAGS_enable_request_context) {
  68       return false;
  69     }
  70     bool prev = getStaticContext().get() != nullptr;
  71     getStaticContext().reset(new std::shared_ptr<RequestContext>(
  72                      std::make_shared<RequestContext>()));
  73     return prev;
  74   }
  75
  76   // Get the current context.
  77   static RequestContext* get() {
  78     if (!FLAGS_enable_request_context ||
  79         getStaticContext().get() == nullptr) {
  80       if (defaultContext == nullptr) {
  81         defaultContext = new RequestContext;
  82       }
  83       return defaultContext;
  84     }
  85     return getStaticContext().get()->get();
  86   }
  87
  88   // The following API may be used to set per-request data in a thread-safe way.
  89   // This access is still performance sensitive, so please ask if you need help
  90   // profiling any use of these functions.
  91   void setContextData(
  92     const std::string& val, std::unique_ptr<RequestData> data) {
  93     if (!FLAGS_enable_request_context) {
  94       return;
  95     }
  96
  97     folly::RWSpinLock::WriteHolder guard(lock);
  98     if (data_.find(val) != data_.end()) {
  99       LOG_FIRST_N(WARNING, 1) <<
 100         "Called RequestContext::setContextData with data already set";
 101
 102       data_[val] = nullptr;
 103     } else {
 104       data_[val] = std::move(data);
 105     }
 106   }
 107
 108   bool hasContextData(const std::string& val) {
 109     folly::RWSpinLock::ReadHolder guard(lock);
 110     return data_.find(val) != data_.end();
 111   }
 112
 113   RequestData* getContextData(const std::string& val) {
 114     folly::RWSpinLock::ReadHolder guard(lock);
 115     auto r = data_.find(val);
 116     if (r == data_.end()) {
 117       return nullptr;
 118     } else {
 119       return r->second.get();
 120     }
 121   }
 122
 123   void clearContextData(const std::string& val) {
 124     folly::RWSpinLock::WriteHolder guard(lock);
 125     data_.erase(val);
 126   }
 127
 128   // The following API is used to pass the context through queues / threads.
 129   // saveContext is called to geta shared_ptr to the context, and
 130   // setContext is used to reset it on the other side of the queue.
 131   //
 132   // A shared_ptr is used, because many request may fan out across
 133   // multiple threads, or do post-send processing, etc.
 134
 135   static std::shared_ptr<RequestContext>
 136   setContext(std::shared_ptr<RequestContext> ctx) {
 137     if (FLAGS_enable_request_context) {
 138       std::shared_ptr<RequestContext> old_ctx;
 139       if (getStaticContext().get()) {
 140         old_ctx = *getStaticContext().get();
 141       }
 142       if (ctx == nullptr) {
 143         getStaticContext().reset(nullptr);
 144       } else {
 145         getStaticContext().reset(new std::shared_ptr<RequestContext>(ctx));
 146       }
 147       return old_ctx;
 148     }
 149     return std::shared_ptr<RequestContext>();
 150   }
 151
 152   static std::shared_ptr<RequestContext> saveContext() {
 153     if (!FLAGS_enable_request_context) {
 154       return std::shared_ptr<RequestContext>();
 155     }
 156     if (getStaticContext().get() == nullptr) {
 157       return std::shared_ptr<RequestContext>();
 158     } else {
 159       return *getStaticContext().get();
 160     }
 161   }
 162
 163   // Used to solve static destruction ordering issue.  Any static object
 164   // that uses RequestContext must call this function in its constructor.
 165   //
 166   // See below link for more details.
 167   // http://stackoverflow.com/questions/335369/
 168   // finding-c-static-initialization-order-problems#335746
 169   static folly::ThreadLocalPtr<std::shared_ptr<RequestContext>>&
 170   getStaticContext() {
 171     static folly::ThreadLocalPtr<std::shared_ptr<RequestContext> > context;
 172     return context;
 173   }
 174
 175  private:
 176   folly::RWSpinLock lock;
 177   std::map<std::string, std::unique_ptr<RequestData>> data_;
 178 };
 179
 180 /**
 181  * Set the request context for a specific scope. For example,
 182  * if you ran a part of a request in another thread you could
 183  * use RequestContextGuard to copy apply the request context
 184  * inside the other therad.
 185  */
 186 class RequestContextGuard {
 187  public:
 188   explicit RequestContextGuard(std::shared_ptr<RequestContext> ctx) {
 189     oldctx_ = RequestContext::setContext(std::move(ctx));
 190   }
 191
 192   ~RequestContextGuard() {
 193     RequestContext::setContext(std::move(oldctx_));
 194   }
 195
 196  private:
 197   std::shared_ptr<RequestContext> oldctx_;
 198 };
 199
 200 }