folly/io/async/Request.h

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