folly/fibers/FiberManagerMap.cpp

   1 /*
   2  * Copyright 2016 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 #include "FiberManagerMap.h"
  17
  18 #include <memory>
  19 #include <unordered_map>
  20
  21 #include <folly/Synchronized.h>
  22 #include <folly/ThreadLocal.h>
  23
  24 namespace folly {
  25 namespace fibers {
  26
  27 namespace {
  28
  29 class EventBaseOnDestructionCallback : public EventBase::LoopCallback {
  30  public:
  31   explicit EventBaseOnDestructionCallback(EventBase& evb) : evb_(evb) {}
  32   void runLoopCallback() noexcept override;
  33
  34  private:
  35   EventBase& evb_;
  36 };
  37
  38 class GlobalCache {
  39  public:
  40   static FiberManager& get(EventBase& evb, const FiberManager::Options& opts) {
  41     return instance().getImpl(evb, opts);
  42   }
  43
  44   static std::unique_ptr<FiberManager> erase(EventBase& evb) {
  45     return instance().eraseImpl(evb);
  46   }
  47
  48  private:
  49   GlobalCache() {}
  50
  51   // Leak this intentionally. During shutdown, we may call getFiberManager,
  52   // and want access to the fiber managers during that time.
  53   static GlobalCache& instance() {
  54     static auto ret = new GlobalCache();
  55     return *ret;
  56   }
  57
  58   FiberManager& getImpl(EventBase& evb, const FiberManager::Options& opts) {
  59     std::lock_guard<std::mutex> lg(mutex_);
  60
  61     auto& fmPtrRef = map_[&evb];
  62
  63     if (!fmPtrRef) {
  64       auto loopController = make_unique<EventBaseLoopController>();
  65       loopController->attachEventBase(evb);
  66       evb.runOnDestruction(new EventBaseOnDestructionCallback(evb));
  67
  68       fmPtrRef = make_unique<FiberManager>(std::move(loopController), opts);
  69     }
  70
  71     return *fmPtrRef;
  72   }
  73
  74   std::unique_ptr<FiberManager> eraseImpl(EventBase& evb) {
  75     std::lock_guard<std::mutex> lg(mutex_);
  76
  77     DCHECK_EQ(1, map_.count(&evb));
  78
  79     auto ret = std::move(map_[&evb]);
  80     map_.erase(&evb);
  81     return ret;
  82   }
  83
  84   std::mutex mutex_;
  85   std::unordered_map<EventBase*, std::unique_ptr<FiberManager>> map_;
  86 };
  87
  88 constexpr size_t kEraseListMaxSize = 64;
  89
  90 class ThreadLocalCache {
  91  public:
  92   static FiberManager& get(EventBase& evb, const FiberManager::Options& opts) {
  93     return instance()->getImpl(evb, opts);
  94   }
  95
  96   static void erase(EventBase& evb) {
  97     for (auto& localInstance : instance().accessAllThreads()) {
  98       SYNCHRONIZED(info, localInstance.eraseInfo_) {
  99         if (info.eraseList.size() >= kEraseListMaxSize) {
 100           info.eraseAll = true;
 101         } else {
 102           info.eraseList.push_back(&evb);
 103         }
 104         localInstance.eraseRequested_ = true;
 105       }
 106     }
 107   }
 108
 109  private:
 110   ThreadLocalCache() {}
 111
 112   struct ThreadLocalCacheTag {};
 113   using ThreadThreadLocalCache =
 114       ThreadLocal<ThreadLocalCache, ThreadLocalCacheTag>;
 115
 116   // Leak this intentionally. During shutdown, we may call getFiberManager,
 117   // and want access to the fiber managers during that time.
 118   static ThreadThreadLocalCache& instance() {
 119     static auto ret =
 120         new ThreadThreadLocalCache([]() { return new ThreadLocalCache(); });
 121     return *ret;
 122   }
 123
 124   FiberManager& getImpl(EventBase& evb, const FiberManager::Options& opts) {
 125     eraseImpl();
 126
 127     auto& fmPtrRef = map_[&evb];
 128     if (!fmPtrRef) {
 129       fmPtrRef = &GlobalCache::get(evb, opts);
 130     }
 131
 132     DCHECK(fmPtrRef != nullptr);
 133
 134     return *fmPtrRef;
 135   }
 136
 137   void eraseImpl() {
 138     if (!eraseRequested_.load()) {
 139       return;
 140     }
 141
 142     SYNCHRONIZED(info, eraseInfo_) {
 143       if (info.eraseAll) {
 144         map_.clear();
 145       } else {
 146         for (auto evbPtr : info.eraseList) {
 147           map_.erase(evbPtr);
 148         }
 149       }
 150
 151       info.eraseList.clear();
 152       info.eraseAll = false;
 153       eraseRequested_ = false;
 154     }
 155   }
 156
 157   std::unordered_map<EventBase*, FiberManager*> map_;
 158   std::atomic<bool> eraseRequested_{false};
 159
 160   struct EraseInfo {
 161     bool eraseAll{false};
 162     std::vector<EventBase*> eraseList;
 163   };
 164
 165   folly::Synchronized<EraseInfo> eraseInfo_;
 166 };
 167
 168 void EventBaseOnDestructionCallback::runLoopCallback() noexcept {
 169   auto fm = GlobalCache::erase(evb_);
 170   DCHECK(fm.get() != nullptr);
 171   ThreadLocalCache::erase(evb_);
 172
 173   delete this;
 174 }
 175
 176 } // namespace
 177
 178 FiberManager& getFiberManager(
 179     EventBase& evb,
 180     const FiberManager::Options& opts) {
 181   return ThreadLocalCache::get(evb, opts);
 182 }
 183 }
 184 }