folly/folly/fibers/Baton.h

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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/experimental/coro/Coroutine.h>
#include <folly/io/async/HHWheelTimer.h>

namespace folly {
namespace fibers {

class Fiber;
class FiberManager;

/**
 * @class Baton
 *
 * Primitive which allows one to put current Fiber to sleep and wake it from
 * another Fiber/thread.
 */
class Baton { … };

#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(coro::coroutine_handle<> h) {
    assert(!h_);
    h_ = std::move(h);
    baton_.setWaiter(*this);
  }

 private:
  coro::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>