/*
 * 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.
 */

//
// Docs: https://fburl.com/fbcref_asyncscope
//

#pragma once

#include <folly/CancellationToken.h>
#include <folly/ExceptionWrapper.h>
#include <folly/experimental/coro/Coroutine.h>
#include <folly/experimental/coro/CurrentExecutor.h>
#include <folly/experimental/coro/Task.h>
#include <folly/experimental/coro/detail/Barrier.h>
#include <folly/experimental/coro/detail/BarrierTask.h>
#include <folly/futures/Future.h>
#include <folly/portability/SourceLocation.h>

#include <glog/logging.h>

#include <atomic>
#include <cassert>
#include <optional>

#if FOLLY_HAS_COROUTINES

namespace folly {
namespace coro {

/**
 * The AsyncScope class is used to allow you to start a dynamic, unbounded
 * number of tasks, which can all run concurrently, and then later wait for
 * completion of all of the tasks.
 *
 * Tasks added to an AsyncScope must have a void or folly::Unit result-type
 * and must handle any errors prior to completing.
 *
 * @refcode folly/docs/examples/folly/experimental/coro/AsyncScope.cpp
 * @class folly::coro::AsyncScope
 */
//
// Example:
//    folly::coro::Task<void> process(Event event) {
//      try {
//        co_await do_processing(event.data);
//      } catch (...) {
//        LOG(ERROR) << "Processing event failed";
//      }
//    }
//
//    folly::coro::AsyncScope scope;
//    scope.add(processEvent(ev1).scheduleOn(folly::getGlobalCPUExecutor()));
//    scope.add(processEvent(ev2).scheduleOn(folly::getGlobalCPUExecutor()));
//    scope.add(processEvent(ev3).scheduleOn(folly::getGlobalCPUExecutor()));
//    co_await scope.joinAsync();
//
class AsyncScope {
 public:
  AsyncScope() noexcept;

  // @param throwOnJoin If true, will throw last unhandled exception (if any)
  //                    on joinAsync. Default behavior is to ignore the
  //                    exception in opt builds and FATAL in dev builds
  explicit AsyncScope(bool throwOnJoin) noexcept;

  // Destroy the AsyncScope object.
  //
  // NOTE: If you have called add() on this scope then you _must_
  // call either cleanup() or joinAsync() and wait until the that operation
  // completes before calling the destructor.
  ~AsyncScope();

  /**
   * Query the number of tasks added to the scope that have not yet completed.
   */
  std::size_t remaining() const noexcept;

  /**
   * Start the specified task/awaitable by co_awaiting it.
   *
   * Exceptions
   * ----------
   * IMPORTANT: Tasks submitted to the AsyncScope by calling .add() must
   * ensure they do not complete with an exception. Exceptions propagating
   * from the 'co_await awaitable' expression are logged using DFATAL.
   *
   * To avoid this occurring you should make sure to catch and handle any
   * exceptions within the task being started here.
   *
   * Interaction with cleanup/joinAsync
   * ----------------------------------
   * It is invalid to call add() once the joinAsync() or cleanup()
   * operations have completed.
   *
   * This generally means that it is unsafe to call .add() once cleanup()
   * has started as it may be racing with completion of cleanup().
   *
   * The exception to this rule is for cases where you know you are running
   * within a task that has been started with .add() and thus you know that
   * cleanup() will not yet have completed.
   *
   * Passing folly::coro::Task
   * -------------------------
   * NOTE: You cannot pass a folly::coro::Task to this method.
   * You must first call .scheduleOn() to specify which executor the task
   * should run on.
   */
  // returnAddress customize entry point to async stack (useful if this is
  // called from async code already). If not set will default to
  // FOLLY_ASYNC_STACK_RETURN_ADDRESS()
  template <typename Awaitable>
  void add(Awaitable&& awaitable, void* returnAddress = nullptr);

  template <typename Awaitable>
  void addWithSourceLoc(
      Awaitable&& awaitable,
      void* returnAddress = nullptr,
      source_location sourceLocation = source_location::current());

  /**
   * Asynchronously wait for all started tasks to complete.
   *
   * Either call this method _or_ cleanup() to join the work.
   * It is invalid to call both of them.
   */
  Task<void> joinAsync() noexcept;

  /**
   * Asynchronously cleanup all started tasks.
   *
   * If you have previuosly called add() then you must call cleanup()
   * and wait for the retuned future to complete before the AsyncScope
   * object destructs.
   */
  SemiFuture<Unit> cleanup() noexcept;

 private:
  template <typename Awaitable>
  static detail::DetachedBarrierTask addImpl(
      Awaitable awaitable,
      bool throwOnJoin,
      folly::exception_wrapper& maybeException,
      std::optional<source_location> source,
      std::atomic<bool>& exceptionRaised) {
    static_assert(
        std::is_void_v<await_result_t<Awaitable>> ||
            std::is_same_v<await_result_t<Awaitable>, folly::Unit>,
        "Result of the task would be discarded. Make sure task result is either void or folly::Unit.");

    try {
      co_await std::move(awaitable);
    } catch (const OperationCancelled&) {
    } catch (...) {
      if (throwOnJoin) {
        LOG(ERROR)
            << (source.has_value() ? sourceLocationToString(source.value())
                                   : "")
            << "Unhandled exception thrown from task added to AsyncScope: "
            << folly::exceptionStr(current_exception());
        if (!exceptionRaised.exchange(true)) {
          maybeException = folly::exception_wrapper(current_exception());
        }
      } else {
        LOG(DFATAL)
            << (source.has_value() ? sourceLocationToString(source.value())
                                   : "")
            << "Unhandled exception thrown from task added to AsyncScope: "
            << folly::exceptionStr(current_exception());
      }
    }
  }

  detail::Barrier barrier_{1};
  std::atomic<bool> anyTasksStarted_{false};
  std::atomic<bool> joinStarted_{false};
  bool joined_{false};
  bool throwOnJoin_{false};
  folly::exception_wrapper maybeException_;
  std::atomic<bool> exceptionRaised_{false};
};

inline AsyncScope::AsyncScope() noexcept : throwOnJoin_(false) {}

inline AsyncScope::AsyncScope(bool throwOnJoin) noexcept
    : throwOnJoin_(throwOnJoin) {}

inline AsyncScope::~AsyncScope() {
  CHECK(!anyTasksStarted_.load(std::memory_order_relaxed) || joined_)
      << "AsyncScope::cleanup() not yet complete";
}

inline std::size_t AsyncScope::remaining() const noexcept {
  const std::size_t count = barrier_.remaining();
  return joinStarted_.load(std::memory_order_relaxed)
      ? count
      : (count > 1 ? count - 1 : 0);
}

template <typename Awaitable>
FOLLY_NOINLINE inline void AsyncScope::add(
    Awaitable&& awaitable, void* returnAddress) {
  assert(
      !joined_ &&
      "It is invalid to add() more work after work has been joined");
  anyTasksStarted_.store(true, std::memory_order_relaxed);
  addImpl(
      static_cast<Awaitable&&>(awaitable),
      throwOnJoin_,
      maybeException_,
      std::nullopt,
      exceptionRaised_)
      .start(
          &barrier_,
          returnAddress ? returnAddress : FOLLY_ASYNC_STACK_RETURN_ADDRESS());
}

template <typename Awaitable>
FOLLY_NOINLINE inline void AsyncScope::addWithSourceLoc(
    Awaitable&& awaitable,
    void* returnAddress,
    source_location sourceLocation) {
  assert(
      !joined_ &&
      "It is invalid to add() more work after work has been joined");
  anyTasksStarted_.store(true, std::memory_order_relaxed);
  addImpl(
      static_cast<Awaitable&&>(awaitable),
      throwOnJoin_,
      maybeException_,
      sourceLocation,
      exceptionRaised_)
      .start(
          &barrier_,
          returnAddress ? returnAddress : FOLLY_ASYNC_STACK_RETURN_ADDRESS());
}

inline Task<void> AsyncScope::joinAsync() noexcept {
  assert(
      !joinStarted_.load(std::memory_order_relaxed) &&
      "It is invalid to join a scope multiple times");
  joinStarted_.store(true, std::memory_order_relaxed);
  co_await barrier_.arriveAndWait();
  joined_ = true;
  if (maybeException_.has_exception_ptr()) {
    maybeException_.throw_exception();
  }
}

inline folly::SemiFuture<folly::Unit> AsyncScope::cleanup() noexcept {
  return joinAsync().semi();
}

/**
 * A cancellable version of AsyncScope. Work added to this scope will be
 * provided a cancellation token for cancelling during join.
 *
 * See add() and cancelAndJoinAsync() for more information.
 *
 * Note: Task and AsyncGenerator will ignore the internal cancellation
 * signal if they already have a cancellation token (i.e. if someone has already
 * called co_withCancellation on them.)
 * If you need an external cancellation signal as well, pass that token to this
 * constructor or to add() instead of attaching it to the Awaitable.
 *
 * @refcode
 * folly/docs/examples/folly/experimental/coro/CancellableAsyncScope.cpp
 * @class folly::coro::CancellableAsyncScope
 */
class CancellableAsyncScope {
 public:
  CancellableAsyncScope() noexcept
      : cancellationToken_(cancellationSource_.getToken()) {}
  explicit CancellableAsyncScope(bool throwOnJoin) noexcept
      : cancellationToken_(cancellationSource_.getToken()),
        scope_(throwOnJoin) {}
  explicit CancellableAsyncScope(CancellationToken&& token)
      : cancellationToken_(CancellationToken::merge(
            cancellationSource_.getToken(), std::move(token))) {}
  CancellableAsyncScope(CancellationToken&& token, bool throwOnJoin)
      : cancellationToken_(CancellationToken::merge(
            cancellationSource_.getToken(), std::move(token))),
        scope_(throwOnJoin) {}

  /**
   * Query the number of tasks added to the scope that have not yet completed.
   */
  std::size_t remaining() const noexcept { return scope_.remaining(); }

  /**
   * Start the specified task/awaitable by co_awaiting it. The awaitable will be
   * provided a cancellation token to respond to cancelAndJoinAsync() in the
   * future.
   *
   * An additional cancellation token may be passed in to apply to the
   * awaitable; it will be merged with the internal token.
   *
   * Note that cancellation is cooperative, your task must handle cancellation
   * in order to have any effect.
   *
   * See the documentation on AsyncScope::add.
   */
  template <typename Awaitable>
  FOLLY_NOINLINE void add(
      Awaitable&& awaitable,
      std::optional<CancellationToken> token = std::nullopt,
      void* returnAddress = nullptr) {
    scope_.add(
        co_withCancellation(
            token ? CancellationToken::merge(*token, cancellationToken_)
                  : cancellationToken_,
            static_cast<Awaitable&&>(awaitable)),
        returnAddress ? returnAddress : FOLLY_ASYNC_STACK_RETURN_ADDRESS());
  }

  template <typename Awaitable>
  FOLLY_NOINLINE void addWithSourceLoc(
      Awaitable&& awaitable,
      std::optional<CancellationToken> token,
      void* returnAddress = nullptr,
      source_location sourceLocation = source_location::current()) {
    scope_.addWithSourceLoc(
        co_withCancellation(
            token ? CancellationToken::merge(*token, cancellationToken_)
                  : cancellationToken_,
            static_cast<Awaitable&&>(awaitable)),
        returnAddress ? returnAddress : FOLLY_ASYNC_STACK_RETURN_ADDRESS(),
        sourceLocation);
  }

  template <typename Awaitable>
  void addWithSourceLoc(
      Awaitable&& awaitable,
      source_location sourceLocation = source_location::current()) {
    addWithSourceLoc(
        std::forward<Awaitable>(awaitable),
        std::nullopt,
        nullptr,
        std::move(sourceLocation));
  }

  /**
   * Schedules the given task on the current executor and adds it to the
   * AsyncScope. The task will be provided a cancellation token to respond to
   * cancelAndJoinAsync() in the future.
   *
   * Note that cancellation is cooperative, your task must handle cancellation
   * in order to have any effect.
   */
  template <class T>
  folly::coro::Task<void> co_schedule(folly::coro::Task<T>&& task) {
    add(std::move(task).scheduleOn(co_await co_current_executor));
    co_return;
  }

  /**
   * Request cancellation for all started tasks that accepted a
   * CancellationToken in add().
   */
  void requestCancellation() const noexcept {
    cancellationSource_.requestCancellation();
  }

  /**
   * Query if cancellation was requested on the tasks added to this AsyncScope.
   *
   * This will return true if either cancellation was requested using
   * `requestCancellation()` method of this scope OR if cancellation is
   * requested on the token passed into the constructor.
   */
  bool isScopeCancellationRequested() const noexcept {
    return cancellationToken_.isCancellationRequested();
  }

  /**
   * Request cancellation and asynchronously wait for all started tasks to
   * complete.
   *
   * Either call this method, _or_ joinAsync() to join the work. It is invalid
   * to call both of them.
   */
  Task<void> cancelAndJoinAsync() noexcept {
    requestCancellation();
    co_await joinAsync();
  }

  /**
   * Asynchronously wait for all started tasks to complete without requesting
   * cancellation.
   *
   * Either call this method _or_ cancelAndJoinAsync() to join the
   * work. It is invalid to call both of them.
   */
  Task<void> joinAsync() noexcept { co_await scope_.joinAsync(); }

 private:
  folly::CancellationSource cancellationSource_;
  CancellationToken cancellationToken_;
  AsyncScope scope_;
};

} // namespace coro
} // namespace folly

#endif // FOLLY_HAS_COROUTINES
