folly/io/async/Request.h

   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 #pragma once
  22
  23 #include <map>
  24 #include <memory>
  25
  26 #include <folly/RWSpinLock.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   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::RWSpinLock> data_;
 103 };
 104
 105 class RequestContextScopeGuard {
 106  private:
 107   std::shared_ptr<RequestContext> prev_;
 108
 109  public:
 110   // Create a new RequestContext and reset to the original value when
 111   // this goes out of scope.
 112   RequestContextScopeGuard() : prev_(RequestContext::saveContext()) {
 113     RequestContext::create();
 114   }
 115
 116   // Set a RequestContext that was previously captured by saveContext(). It will
 117   // be automatically reset to the original value when this goes out of scope.
 118   explicit RequestContextScopeGuard(std::shared_ptr<RequestContext> ctx)
 119       : prev_(RequestContext::setContext(std::move(ctx))) {}
 120
 121   ~RequestContextScopeGuard() {
 122     RequestContext::setContext(std::move(prev_));
 123   }
 124 };
 125 }