Baton.h

Go to the documentation of this file.

/*
 * Copyright 2014-present Facebook, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <atomic>

#include <folly/Portability.h>
#include <folly/detail/Futex.h>
#include <folly/fibers/TimeoutController.h>

#if FOLLY_HAS_COROUTINES
#include <experimental/coroutine>
#endif

namespace folly {
namespace fibers {

class Fiber;
class FiberManager;

class Baton {
 public:
  class TimeoutHandler;

  class Waiter {
   public:
    virtual void post() = 0;

    virtual ~Waiter() {}
  };

  Baton() noexcept;

  ~Baton() noexcept = default;

  bool ready() const {
    auto state = waiter_.load();
    return state == POSTED;
  }

  void setWaiter(Waiter& waiter);

  void wait();

  void wait(TimeoutHandler& timeoutHandler);

  template <typename F>
  void wait(F&& mainContextFunc);

  bool try_wait();

  template <typename Rep, typename Period>
  bool try_wait_for(const std::chrono::duration<Rep, Period>& timeout) {
    return try_wait_for(timeout, [] {});
  }

  template <typename Rep, typename Period, typename F>
  bool try_wait_for(
      const std::chrono::duration<Rep, Period>& timeout,
      F&& mainContextFunc);

  template <typename Clock, typename Duration>
  bool try_wait_until(
      const std::chrono::time_point<Clock, Duration>& deadline) {
    return try_wait_until(deadline, [] {});
  }

  template <typename Clock, typename Duration, typename F>
  bool try_wait_until(
      const std::chrono::time_point<Clock, Duration>& deadline,
      F&& mainContextFunc);

  template <typename Clock, typename Duration, typename F>
  bool try_wait_for(
      const std::chrono::time_point<Clock, Duration>& deadline,
      F&& mainContextFunc);

  template <typename Rep, typename Period>
  bool timed_wait(const std::chrono::duration<Rep, Period>& timeout) {
    return try_wait_for(timeout);
  }

  template <typename Rep, typename Period, typename F>
  bool timed_wait(
      const std::chrono::duration<Rep, Period>& timeout,
      F&& mainContextFunc) {
    return try_wait_for(timeout, static_cast<F&&>(mainContextFunc));
  }

  template <typename Clock, typename Duration>
  bool timed_wait(const std::chrono::time_point<Clock, Duration>& deadline) {
    return try_wait_until(deadline);
  }

  template <typename Clock, typename Duration, typename F>
  bool timed_wait(
      const std::chrono::time_point<Clock, Duration>& deadline,
      F&& mainContextFunc) {
    return try_wait_until(deadline, static_cast<F&&>(mainContextFunc));
  }

  void post();

  void reset();

  class TimeoutHandler {
   public:
    void scheduleTimeout(TimeoutController::Duration timeoutMs);

   private:
    friend class Baton;

    void cancelTimeout();

    std::function<void()> timeoutFunc_{nullptr};
    FiberManager* fiberManager_{nullptr};

    intptr_t timeoutPtr_{0};
  };

 private:
  class FiberWaiter;

  enum {
    PreBlockAttempts = 300,
  };

  explicit Baton(intptr_t state) : waiter_(state) {}

  void postHelper(intptr_t new_value);
  void postThread();
  void waitThread();

  template <typename F>
  inline void waitFiber(FiberManager& fm, F&& mainContextFunc);
  bool spinWaitForEarlyPost();

  bool timedWaitThread(TimeoutController::Duration timeout);

  static constexpr intptr_t NO_WAITER = 0;
  static constexpr intptr_t POSTED = -1;
  static constexpr intptr_t TIMEOUT = -2;
  static constexpr intptr_t THREAD_WAITING = -3;

  union {
    std::atomic<intptr_t> waiter_;
    struct {
      folly::detail::Futex<> futex{};
      int32_t _unused_packing;
    } futex_;
  };
};

#if FOLLY_HAS_COROUTINES
namespace detail {
class BatonAwaitableWaiter : public Baton::Waiter {
 public:
  explicit BatonAwaitableWaiter(Baton& baton) : baton_(baton) {}

  void post() override {
    assert(h_);
    h_();
  }

  bool await_ready() const {
    return baton_.ready();
  }

  void await_resume() {}

  void await_suspend(std::experimental::coroutine_handle<> h) {
    assert(!h_);
    h_ = std::move(h);
    baton_.setWaiter(*this);
  }

 private:
  std::experimental::coroutine_handle<> h_;
  Baton& baton_;
};
} // namespace detail

inline detail::BatonAwaitableWaiter /* implicit */ operator co_await(
    Baton& baton) {
  return detail::BatonAwaitableWaiter(baton);
}
#endif
} // namespace fibers
} // namespace folly

#include <folly/fibers/Baton-inl.h>