folly/fibers/Baton-inl.h

   1 /*
   2  * Copyright 2017-present 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 <folly/fibers/Fiber.h>
  17 #include <folly/fibers/FiberManagerInternal.h>
  18
  19 namespace folly {
  20 namespace fibers {
  21
  22 inline Baton::Baton() : Baton(NO_WAITER) {
  23   assert(Baton(NO_WAITER).futex_.futex == static_cast<uint32_t>(NO_WAITER));
  24   assert(Baton(POSTED).futex_.futex == static_cast<uint32_t>(POSTED));
  25   assert(Baton(TIMEOUT).futex_.futex == static_cast<uint32_t>(TIMEOUT));
  26   assert(
  27       Baton(THREAD_WAITING).futex_.futex ==
  28       static_cast<uint32_t>(THREAD_WAITING));
  29
  30   assert(futex_.futex.is_lock_free());
  31   assert(waitingFiber_.is_lock_free());
  32 }
  33
  34 template <typename F>
  35 void Baton::wait(F&& mainContextFunc) {
  36   auto fm = FiberManager::getFiberManagerUnsafe();
  37   if (!fm || !fm->activeFiber_) {
  38     mainContextFunc();
  39     return waitThread();
  40   }
  41
  42   return waitFiber(*fm, std::forward<F>(mainContextFunc));
  43 }
  44
  45 template <typename F>
  46 void Baton::waitFiber(FiberManager& fm, F&& mainContextFunc) {
  47   auto& waitingFiber = waitingFiber_;
  48   auto f = [&mainContextFunc, &waitingFiber](Fiber& fiber) mutable {
  49     auto baton_fiber = waitingFiber.load();
  50     do {
  51       if (LIKELY(baton_fiber == NO_WAITER)) {
  52         continue;
  53       } else if (baton_fiber == POSTED || baton_fiber == TIMEOUT) {
  54         fiber.resume();
  55         break;
  56       } else {
  57         throw std::logic_error("Some Fiber is already waiting on this Baton.");
  58       }
  59     } while (!waitingFiber.compare_exchange_weak(
  60         baton_fiber, reinterpret_cast<intptr_t>(&fiber)));
  61
  62     mainContextFunc();
  63   };
  64
  65   fm.awaitFunc_ = std::ref(f);
  66   fm.activeFiber_->preempt(Fiber::AWAITING);
  67 }
  68
  69 template <typename Rep, typename Period, typename F>
  70 bool Baton::try_wait_for(
  71     const std::chrono::duration<Rep, Period>& timeout,
  72     F&& mainContextFunc) {
  73   auto fm = FiberManager::getFiberManagerUnsafe();
  74
  75   if (!fm || !fm->activeFiber_) {
  76     mainContextFunc();
  77     return timedWaitThread(timeout);
  78   }
  79
  80   auto& baton = *this;
  81   bool canceled = false;
  82   auto timeoutFunc = [&baton, &canceled]() mutable {
  83     baton.postHelper(TIMEOUT);
  84     canceled = true;
  85   };
  86
  87   auto id =
  88       fm->timeoutManager_->registerTimeout(std::ref(timeoutFunc), timeout);
  89
  90   waitFiber(*fm, static_cast<F&&>(mainContextFunc));
  91
  92   auto posted = waitingFiber_ == POSTED;
  93
  94   if (!canceled) {
  95     fm->timeoutManager_->cancel(id);
  96   }
  97
  98   return posted;
  99 }
 100
 101 template <typename Clock, typename Duration, typename F>
 102 bool Baton::try_wait_until(
 103     const std::chrono::time_point<Clock, Duration>& deadline,
 104     F&& mainContextFunc) {
 105   auto now = Clock::now();
 106
 107   if (LIKELY(now <= deadline)) {
 108     return try_wait_for(deadline - now, static_cast<F&&>(mainContextFunc));
 109   } else {
 110     return try_wait_for(Duration{}, static_cast<F&&>(mainContextFunc));
 111   }
 112 }
 113 } // namespace fibers
 114 } // namespace folly