folly/io/async/EventBaseLocal.h

   1 /*
   2  * Copyright 2017 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 <boost/noncopyable.hpp>
  20 #include <folly/Synchronized.h>
  21 #include <folly/io/async/EventBase.h>
  22 #include <memory>
  23 #include <mutex>
  24 #include <unordered_set>
  25 #include <utility>
  26
  27 namespace folly {
  28
  29 namespace detail {
  30
  31 class EventBaseLocalBase : public EventBaseLocalBaseBase, boost::noncopyable {
  32  public:
  33   EventBaseLocalBase() {}
  34   virtual ~EventBaseLocalBase();
  35   void erase(EventBase& evb);
  36   void onEventBaseDestruction(EventBase& evb) override;
  37
  38  protected:
  39   void setVoid(EventBase& evb, std::shared_ptr<void>&& ptr);
  40   void setVoidUnlocked(EventBase& evb, std::shared_ptr<void>&& ptr);
  41   void* getVoid(EventBase& evb);
  42
  43   folly::Synchronized<std::unordered_set<EventBase*>> eventBases_;
  44   static std::atomic<uint64_t> keyCounter_;
  45   uint64_t key_{keyCounter_++};
  46 };
  47
  48 }
  49
  50 /**
  51  * A storage abstraction for data that should be tied to an EventBase.
  52  *
  53  *   struct Foo { Foo(int a, int b); };
  54  *   EventBaseLocal<Foo> myFoo;
  55  *   ...
  56  *   EventBase evb;
  57  *   myFoo.set(evb, new Foo(1, 2));
  58  *   myFoo.set(evb, 1, 2);
  59  *   Foo* foo = myFoo.get(evb);
  60  *   myFoo.erase(evb);
  61  *   Foo& foo = myFoo.getOrCreate(evb, 1, 2); // ctor
  62  *   Foo& foo = myFoo.getOrCreate(evb, 1, 2); // no ctor
  63  *   myFoo.erase(evb);
  64  *   Foo& foo = myFoo.getOrCreateFn(evb, [] () { return new Foo(3, 4); })
  65  *
  66  * The objects will be deleted when the EventBaseLocal or the EventBase is
  67  * destructed (whichever comes first).  All methods are thread-safe.
  68  *
  69  * The user is responsible for throwing away invalid references/ptrs returned
  70  * by the get() method after set/erase is called.  If shared ownership is
  71  * needed, use a EventBaseLocal<shared_ptr<...>>.
  72  */
  73 template<typename T>
  74 class EventBaseLocal : public detail::EventBaseLocalBase {
  75  public:
  76   EventBaseLocal(): EventBaseLocalBase() {}
  77
  78   T* get(EventBase& evb) {
  79     return static_cast<T*>(getVoid(evb));
  80   }
  81
  82   void emplace(EventBase& evb, T* ptr) {
  83     std::shared_ptr<T> smartPtr(ptr);
  84     setVoid(evb, std::move(smartPtr));
  85   }
  86
  87   template <typename... Args>
  88   void emplace(EventBase& evb, Args&&... args) {
  89     auto smartPtr = std::make_shared<T>(std::forward<Args>(args)...);
  90     setVoid(evb, smartPtr);
  91   }
  92
  93   template <typename... Args>
  94   T& getOrCreate(EventBase& evb, Args&&... args) {
  95     std::lock_guard<std::mutex> lg(evb.localStorageMutex_);
  96
  97     auto it2 = evb.localStorage_.find(key_);
  98     if (LIKELY(it2 != evb.localStorage_.end())) {
  99       return *static_cast<T*>(it2->second.get());
 100     } else {
 101       auto smartPtr = std::make_shared<T>(std::forward<Args>(args)...);
 102       auto ptr = smartPtr.get();
 103       setVoidUnlocked(evb, std::move(smartPtr));
 104       return *ptr;
 105     }
 106   }
 107
 108   template <typename Func>
 109   T& getOrCreateFn(EventBase& evb, Func& fn) {
 110     // If this looks like it's copy/pasted from above, that's because it is.
 111     // gcc has a bug (fixed in 4.9) that doesn't allow capturing variadic
 112     // params in a lambda.
 113     std::lock_guard<std::mutex> lg(evb.localStorageMutex_);
 114
 115     auto it2 = evb.localStorage_.find(key_);
 116     if (LIKELY(it2 != evb.localStorage_.end())) {
 117       return *static_cast<T*>(it2->second.get());
 118     } else {
 119       std::shared_ptr<T> smartPtr(fn());
 120       auto ptr = smartPtr.get();
 121       setVoidUnlocked(evb, std::move(smartPtr));
 122       return *ptr;
 123     }
 124   }
 125 };
 126
 127
 128 }