// Copyright 2023 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef LIB_ASYNC_PATTERNS_CPP_TASK_SCOPE_H_
#define LIB_ASYNC_PATTERNS_CPP_TASK_SCOPE_H_

#include <lib/async/dispatcher.h>
#include <lib/async_patterns/cpp/internal/task_queue.h>
#include <lib/fit/function.h>
#include <lib/zx/time.h>
#include <zircon/compiler.h>

namespace async_patterns {

// |TaskScope| lets you post asynchronous tasks that are silently discarded when
// the |TaskScope| object is destroyed. In contrast, tasks posted using
// |async::PostTask|, |async::PostDelayedTask|, and so on will always run unless
// the dispatcher is shutdown -- there is no separate mechanism to cancel them.
//
// The task scope is usually a member of some bigger async business logic
// object. By associating the lifetime of the posted async tasks with this
// object, one can conveniently capture or borrow other members from the async
// tasks without the need for reference counting to achieve memory safety.
// Example:
//
//     class AsyncCounter {
//      public:
//       explicit AsyncCounter(async_dispatcher_t* dispatcher) : tasks_(dispatcher) {
//         // Post some tasks to asynchronously count upwards.
//         //
//         // It is okay if |AsyncCounter| is destroyed before some of these tasks
//         // come due. Those tasks will be discarded so they will not end up accessing
//         // a destroyed object.
//         tasks_.Post([this] { count_++ });
//         tasks_.PostDelayed([this] { count_++ }, zx::sec(1));
//         tasks_.PostDelayed(fit::bind_member<&AsyncCounter::CheckCount>(this), zx::sec(5));
//       }
//
//       void CheckCount() {
//         assert(count_ == 2);
//       }
//
//      private:
//       TaskScope tasks_;
//       int count_ = 0;
//     };
//
// |TaskScope| is thread-unsafe, and must be used and managed from a
// [synchronized dispatcher][synchronized-dispatcher].
//
// [synchronized-dispatcher]:
// https://fuchsia.dev/fuchsia-src/development/languages/c-cpp/thread-safe-async#synchronized-dispatcher
class TaskScope {
 private:
  // |F| must take zero arguments ard return void.
  template <typename F>
  using require_nullary_fn = std::enable_if_t<std::is_void_v<std::invoke_result_t<F>>>;

 public:
  // Creates a |TaskScope| that will post all tasks to the provided |dispatcher|.
  explicit TaskScope(async_dispatcher_t* dispatcher);

  // Destroying the |TaskScope| synchronously destroys all pending tasks. If a
  // task attempts to reentrantly post more tasks into the |TaskScope| within
  // its destructor, those tasks will be synchronously destroyed too.
  ~TaskScope();

  // Schedules to invoke |handler| with a deadline of now.
  //
  // |handler| should be a |void()| callable object.
  //
  // The handler will not run if |TaskScope| is destroyed before it comes due.
  // The handler will not run if the dispatcher shuts down before it comes due.
  //
  // This is a drop-in replacement for |async::PostTask|.
  template <typename Closure>
  require_nullary_fn<Closure> Post(Closure&& handler) {
    PostImpl(internal::Task::Box(std::forward<Closure>(handler)));
  }

  // Schedules to invoke |handler| with a deadline expressed as a |delay| from now.
  //
  // |handler| should be a |void()| callable object.
  //
  // The handler will not run if |TaskScope| is destroyed before it comes due.
  // The handler will not run if the dispatcher shuts down before it comes due.
  //
  // This is a drop-in replacement for |async::PostDelayedTask|.
  template <typename Closure>
  require_nullary_fn<Closure> PostDelayed(Closure&& handler, zx::duration delay) {
    PostImpl(DelayedTask::Box(std::forward<Closure>(handler), this), delay);
  }

  // Schedules to invoke |handler| with the specified |deadline|.
  //
  // |handler| should be a |void()| callable object.
  //
  // The handler will not run if |TaskScope| is destroyed before it comes due.
  // The handler will not run if the dispatcher shuts down before it comes due.
  //
  // This is a drop-in replacement for |async::PostTaskForTime|.
  template <typename Closure>
  require_nullary_fn<Closure> PostForTime(Closure&& handler, zx::time deadline) {
    PostImpl(DelayedTask::Box(std::forward<Closure>(handler), this), deadline);
  }

 private:
  class DelayedTask : public list_node_t, public async_task_t {
   public:
    explicit DelayedTask(TaskScope* owner);
    virtual ~DelayedTask() = default;

    bool Post(async_dispatcher_t* dispatcher, zx::time deadline);
    virtual void Run() = 0;

    template <typename Callable>
    static std::unique_ptr<DelayedTask> Box(Callable&& callable, TaskScope* owner) {
      class TaskImpl final : public DelayedTask {
       public:
        explicit TaskImpl(Callable&& callable, TaskScope* owner)
            : DelayedTask(owner), callable_(std::forward<Callable>(callable)) {}
        void Run() final { callable_(); }

       private:
        Callable callable_;
      };
      return std::make_unique<TaskImpl>(std::forward<Callable>(callable), owner);
    }

   private:
    static void Handler(async_dispatcher_t* dispatcher, async_task_t* task, zx_status_t status);

    TaskScope* owner_;
  };

  void PostImpl(std::unique_ptr<internal::Task> task);
  void PostImpl(std::unique_ptr<DelayedTask> task, zx::duration delay);
  void PostImpl(std::unique_ptr<DelayedTask> task, zx::time deadline);
  void RemoveDelayedTask(DelayedTask* task) __TA_EXCLUDES(checker_);
  void RemoveDelayedTaskLocked(DelayedTask* task) __TA_REQUIRES(checker_);

  async_dispatcher_t* dispatcher_;
  async::synchronization_checker checker_;
  internal::TaskQueue queue_ __TA_GUARDED(checker_);

  // A list of |DelayedTask|.
  list_node_t delayed_task_list_ __TA_GUARDED(checker_) = LIST_INITIAL_VALUE(delayed_task_list_);

  // If true, tasks are discarded instead of posted.
  bool stopped_ __TA_GUARDED(checker_) = false;
};

}  // namespace async_patterns

#endif  // LIB_ASYNC_PATTERNS_CPP_TASK_SCOPE_H_