folly/experimental/fibers/FiberManagerMap.cpp

   1 /*
   2  * Copyright 2015 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 <cassert>
  19 #include <memory>
  20 #include <unordered_map>
  21
  22 #include <folly/ThreadLocal.h>
  23
  24 namespace folly { namespace fibers {
  25
  26 namespace {
  27
  28 // Leak these intentionally.  During shutdown, we may call getFiberManager, and
  29 // want access to the fiber managers during that time.
  30 class LocalFiberManagerMapTag;
  31 typedef folly::ThreadLocal<
  32     std::unordered_map<folly::EventBase*, FiberManager*>,
  33     LocalFiberManagerMapTag>
  34   LocalMapType;
  35 LocalMapType* localFiberManagerMap() {
  36   static auto ret = new LocalMapType();
  37   return ret;
  38 }
  39
  40 typedef
  41   std::unordered_map<folly::EventBase*, std::unique_ptr<FiberManager>>
  42   MapType;
  43 MapType* fiberManagerMap() {
  44   static auto ret = new MapType();
  45   return ret;
  46 }
  47
  48 std::mutex* fiberManagerMapMutex() {
  49   static auto ret = new std::mutex();
  50   return ret;
  51 }
  52
  53
  54 class OnEventBaseDestructionCallback : public folly::EventBase::LoopCallback {
  55  public:
  56   explicit OnEventBaseDestructionCallback(folly::EventBase& evb)
  57            : evb_(&evb) {}
  58   void runLoopCallback() noexcept override {
  59     for (auto& localMap : localFiberManagerMap()->accessAllThreads()) {
  60       localMap.erase(evb_);
  61     }
  62     std::unique_ptr<FiberManager> fm;
  63     {
  64       std::lock_guard<std::mutex> lg(*fiberManagerMapMutex());
  65       auto it = fiberManagerMap()->find(evb_);
  66       assert(it != fiberManagerMap()->end());
  67       fm = std::move(it->second);
  68       fiberManagerMap()->erase(it);
  69     }
  70     assert(fm.get() != nullptr);
  71     fm->loopUntilNoReady();
  72     delete this;
  73   }
  74  private:
  75   folly::EventBase* evb_;
  76 };
  77
  78 FiberManager* getFiberManagerThreadSafe(folly::EventBase& evb,
  79                                         const FiberManager::Options& opts) {
  80   std::lock_guard<std::mutex> lg(*fiberManagerMapMutex());
  81
  82   auto it = fiberManagerMap()->find(&evb);
  83   if (LIKELY(it != fiberManagerMap()->end())) {
  84     return it->second.get();
  85   }
  86
  87   auto loopController = folly::make_unique<EventBaseLoopController>();
  88   loopController->attachEventBase(evb);
  89   auto fiberManager =
  90       folly::make_unique<FiberManager>(std::move(loopController), opts);
  91   auto result = fiberManagerMap()->emplace(&evb, std::move(fiberManager));
  92   evb.runOnDestruction(new OnEventBaseDestructionCallback(evb));
  93   return result.first->second.get();
  94 }
  95
  96 } // namespace
  97
  98 FiberManager& getFiberManager(folly::EventBase& evb,
  99                               const FiberManager::Options& opts) {
 100   auto it = (*localFiberManagerMap())->find(&evb);
 101   if (LIKELY(it != (*localFiberManagerMap())->end())) {
 102     return *(it->second);
 103   }
 104
 105   auto fm = getFiberManagerThreadSafe(evb, opts);
 106   (*localFiberManagerMap())->emplace(&evb, fm);
 107   return *fm;
 108 }
 109
 110 }}