folly/io/async/Request.h

   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 #pragma once
  18
  19 #include <map>
  20 #include <memory>
  21
  22 #include <folly/SharedMutex.h>
  23 #include <folly/Synchronized.h>
  24
  25 namespace folly {
  26
  27 // Some request context that follows an async request through a process
  28 // Everything in the context must be thread safe
  29
  30 class RequestData {
  31  public:
  32   virtual ~RequestData() = default;
  33   // Avoid calling RequestContext::setContextData, setContextDataIfAbsent, or
  34   // clearContextData from these callbacks. Doing so will cause deadlock. We
  35   // could fix these deadlocks, but only at significant performance penalty, so
  36   // just don't do it!
  37   virtual void onSet() {}
  38   virtual void onUnset() {}
  39 };
  40
  41 class RequestContext;
  42
  43 // If you do not call create() to create a unique request context,
  44 // this default request context will always be returned, and is never
  45 // copied between threads.
  46 class RequestContext {
  47  public:
  48   // Create a unique request context for this request.
  49   // It will be passed between queues / threads (where implemented),
  50   // so it should be valid for the lifetime of the request.
  51   static void create() {
  52     setContext(std::make_shared<RequestContext>());
  53   }
  54
  55   // Get the current context.
  56   static RequestContext* get();
  57
  58   // The following API may be used to set per-request data in a thread-safe way.
  59   // This access is still performance sensitive, so please ask if you need help
  60   // profiling any use of these functions.
  61   void setContextData(
  62       const std::string& val,
  63       std::unique_ptr<RequestData> data);
  64
  65   // Unlike setContextData, this method does not panic if the key is already
  66   // present. Returns true iff the new value has been inserted.
  67   bool setContextDataIfAbsent(
  68       const std::string& val,
  69       std::unique_ptr<RequestData> data);
  70
  71   bool hasContextData(const std::string& val) const;
  72
  73   RequestData* getContextData(const std::string& val);
  74   const RequestData* getContextData(const std::string& val) const;
  75
  76   void onSet();
  77   void onUnset();
  78
  79   void clearContextData(const std::string& val);
  80
  81   // The following API is used to pass the context through queues / threads.
  82   // saveContext is called to get a shared_ptr to the context, and
  83   // setContext is used to reset it on the other side of the queue.
  84   //
  85   // Whenever possible, use RequestContextScopeGuard instead of setContext
  86   // to make sure that RequestContext is reset to the original value when
  87   // we exit the scope.
  88   //
  89   // A shared_ptr is used, because many request may fan out across
  90   // multiple threads, or do post-send processing, etc.
  91   static std::shared_ptr<RequestContext> setContext(
  92       std::shared_ptr<RequestContext> ctx);
  93
  94   static std::shared_ptr<RequestContext> saveContext() {
  95     return getStaticContext();
  96   }
  97
  98  private:
  99   static std::shared_ptr<RequestContext>& getStaticContext();
 100
 101   using Data = std::map<std::string, std::unique_ptr<RequestData>>;
 102   folly::Synchronized<Data, folly::SharedMutex> data_;
 103 };
 104
 105 class RequestContextScopeGuard {
 106  private:
 107   std::shared_ptr<RequestContext> prev_;
 108
 109  public:
 110   RequestContextScopeGuard(const RequestContextScopeGuard&) = delete;
 111   RequestContextScopeGuard& operator=(const RequestContextScopeGuard&) = delete;
 112   RequestContextScopeGuard(RequestContextScopeGuard&&) = delete;
 113   RequestContextScopeGuard& operator=(RequestContextScopeGuard&&) = delete;
 114
 115   // Create a new RequestContext and reset to the original value when
 116   // this goes out of scope.
 117   RequestContextScopeGuard() : prev_(RequestContext::saveContext()) {
 118     RequestContext::create();
 119   }
 120
 121   // Set a RequestContext that was previously captured by saveContext(). It will
 122   // be automatically reset to the original value when this goes out of scope.
 123   explicit RequestContextScopeGuard(std::shared_ptr<RequestContext> ctx)
 124       : prev_(RequestContext::setContext(std::move(ctx))) {
 125   }
 126
 127   ~RequestContextScopeGuard() {
 128     RequestContext::setContext(std::move(prev_));
 129   }
 130 };
 131 }