dispatcher.h

// Copyright 2025 The Pigweed Authors
//
// 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
//
//     https://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 <mutex>
#include <type_traits>
#include <utility>
 
#include "pw_allocator/allocator.h"
#include "pw_allocator/shared_ptr.h"
#include "pw_async2/func_task.h"
#include "pw_async2/future_task.h"
#include "pw_async2/internal/lock.h"
#include "pw_async2/task.h"
#include "pw_async2/waker.h"
#include "pw_containers/intrusive_queue.h"
#include "pw_sync/lock_annotations.h"
 
// Coroutines are supported if the build target depends on //pw_async2:coro.
#if defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
#include <functional>  // std::invoke
 
#include "pw_async2/coro.h"
#include "pw_async2/coro_task.h"
#include "pw_async2/fallible_coro_task.h"
#endif  // defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
 
namespace pw::async2 {
 
 
class Dispatcher {
 public:
  Dispatcher(const Dispatcher&) = delete;
  Dispatcher& operator=(const Dispatcher&) = delete;
 
  Dispatcher(Dispatcher&&) = delete;
  Dispatcher& operator=(Dispatcher&&) = delete;
 
  virtual ~Dispatcher() PW_LOCKS_EXCLUDED(internal::lock());
 
  void Terminate() PW_LOCKS_EXCLUDED(internal::lock());
 
  void Post(Task& task) PW_LOCKS_EXCLUDED(internal::lock());
 
  template <typename T>
  void PostShared(const SharedPtr<T>& task)
      PW_LOCKS_EXCLUDED(internal::lock()) {
    PW_ASSERT(PostAllocatedTask(task));
  }
 
  template <typename TaskType,
            typename... Args,
            typename = std::enable_if_t<std::is_base_of_v<Task, TaskType>>>
  [[nodiscard]] SharedPtr<TaskType> Post(Allocator& allocator, Args&&... args) {
    auto task = allocator.MakeShared<TaskType>(std::forward<Args>(args)...);
    if (!PostAllocatedTask(task)) {
      return nullptr;
    }
    return task;
  }
 
  template <
      typename Func = void,
      int&... kExplicitGuard,
      typename Arg,
      typename ActualFunc =
          std::conditional_t<std::is_void_v<Func>, std::decay_t<Arg>, Func>,
      typename = std::enable_if_t<!std::is_base_of_v<Task, ActualFunc>>>
  [[nodiscard]] SharedPtr<FuncTask<ActualFunc>> Post(Allocator& allocator,
                                                     Arg&& func) {
    return Post<FuncTask<ActualFunc>>(allocator, std::forward<Arg>(func));
  }
 
  template <typename Func>
  [[nodiscard]] SharedPtr<RunOnceTask<Func>> RunOnce(Allocator& allocator,
                                                     Func&& func) {
    return Post<RunOnceTask<Func>>(allocator, std::forward<Func>(func));
  }
 
  template <typename Fut>
  [[nodiscard]] SharedPtr<FutureTask<Fut>> PostFuture(Allocator& allocator,
                                                      Fut&& future) {
    return Post<FutureTask<Fut>>(allocator, std::forward<Fut>(future));
  }
 
#if defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
  template <typename T>
  [[nodiscard]] SharedPtr<CoroTask<T>> Post(Allocator& allocator,
                                            Coro<T>&& coro) {
    if (!coro.ok()) {
      return nullptr;
    }
    return Post<CoroTask<T>>(allocator, std::move(coro));
  }
 
  template <typename T,
            typename E = void,
            int&... kExplicitGuard,
            typename Arg,
            typename ErrorHandler =
                std::conditional_t<std::is_void_v<E>, std::decay_t<Arg>, E>>
  [[nodiscard]] SharedPtr<FallibleCoroTask<T, ErrorHandler>> Post(
      Allocator& allocator, Coro<T>&& coro, Arg&& error_handler) {
    if (!coro.ok()) {
      return nullptr;
    }
    return Post<FallibleCoroTask<T, ErrorHandler>>(
        allocator, std::move(coro), std::forward<Arg>(error_handler));
  }
 
  template <auto kCoroFunc, typename... Args>
  [[nodiscard]] auto Post(CoroContext coro_cx, Args&&... args) {
    if constexpr (std::is_member_function_pointer_v<decltype(kCoroFunc)>) {
      return PostSharedMemberCoro<kCoroFunc>(coro_cx,
                                             std::forward<Args>(args)...);
    } else {
      return Post(coro_cx.allocator(),
                  std::invoke(kCoroFunc, coro_cx, std::forward<Args>(args)...));
    }
  }
#endif  // defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
 
  void LogRegisteredTasks() PW_LOCKS_EXCLUDED(internal::lock());
 
 protected:
  constexpr Dispatcher() = default;
 
  bool PopAndRunAllReadyTasks() PW_LOCKS_EXCLUDED(internal::lock());
 
  Task* PopTaskToRun() PW_LOCKS_EXCLUDED(internal::lock()) {
    std::lock_guard lock(internal::lock());
    return PopTaskToRunLocked();
  }
 
  Task* PopTaskToRun(bool& has_posted_tasks)
      PW_LOCKS_EXCLUDED(internal::lock()) {
    std::lock_guard lock(internal::lock());
    Task* task = PopTaskToRunLocked();
    has_posted_tasks = task != nullptr || !sleeping_.empty();
    return task;
  }
 
  Task* PopSingleTaskForThisWake() PW_LOCKS_EXCLUDED(internal::lock()) {
    std::lock_guard lock(internal::lock());
    wants_wake_ = true;
    return PopTaskToRunLocked();
  }
 
  RunTaskResult RunTask(Task& task) PW_LOCKS_EXCLUDED(internal::lock()) {
    return task.RunInDispatcher();
  }
 
 private:
  friend class Task;
  friend class Waker;
 
  // Allow DispatcherForTestFacade to wrap another dispatcher (call Do*).
  template <typename>
  friend class DispatcherForTestFacade;
 
  virtual void DoWake() PW_LOCKS_EXCLUDED(internal::lock()) = 0;
 
  void Wake(Task* task_to_release = nullptr)
      PW_UNLOCK_FUNCTION(internal::lock());
 
  // Removes a task, waking the dispatcher if it is the last task.
  void DeregisterTask(Task& task) PW_UNLOCK_FUNCTION(internal::lock()) {
    if (has_tasks()) {
      task.UnpostAndReleaseRef();
    } else {
      Wake(&task);
    }
  }
 
  Task* PopTaskToRunLocked() PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  static void UnpostTaskList(IntrusiveQueue<Task>& list)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  void RemoveWokenTaskLocked(Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    woken_.remove(task);
  }
  void RemoveSleepingTaskLocked(Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    sleeping_.remove(task);
  }
  void AddSleepingTaskLocked(Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    sleeping_.push_front(task);
  }
  void AddWokenTaskLocked(Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    woken_.push_back(task);
  }
 
  // True if any tasks are registered with the dispatcher.
  bool has_tasks() const PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock()) {
    return !woken_.empty() || !sleeping_.empty();
  }
 
  void LogTaskWakers(const Task& task)
      PW_EXCLUSIVE_LOCKS_REQUIRED(internal::lock());
 
  // Checks if `task` contains a value and posts it to the dispatcher as an
  // allocated task. Returns true if the task posted successfully.
  template <typename T>
  bool PostAllocatedTask(const SharedPtr<T>& task)
      PW_LOCKS_EXCLUDED(internal::lock()) {
    return PostAllocatedTask(
        task.get(),
        task.GetControlBlock(
            allocator::internal::ControlBlockHandle::GetInstance_DO_NOT_USE()));
  }
 
  bool PostAllocatedTask(Task* task,
                         allocator::internal::ControlBlock* control_block)
      PW_LOCKS_EXCLUDED(internal::lock());
 
#if defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
  template <auto kCoroMemberFunc, typename Receiver, typename... Args>
  [[nodiscard]] auto PostSharedMemberCoro(CoroContext coro_cx,
                                          Receiver&& receiver,
                                          Args&&... args)
      PW_LOCKS_EXCLUDED(internal::lock()) {
    return Post(coro_cx.allocator(),
                std::invoke(kCoroMemberFunc,
                            std::forward<Receiver>(receiver),
                            coro_cx,
                            std::forward<Args>(args)...));
  }
#endif  // defined(__cpp_impl_coroutine) && __has_include("pw_async2/coro.h")
 
  IntrusiveQueue<Task> woken_ PW_GUARDED_BY(internal::lock());
  IntrusiveQueue<Task> sleeping_ PW_GUARDED_BY(internal::lock());
 
  // Counts pending DoWake() calls (really should only ever be 1). This is
  // necessary to prevent the Dispatcher from being destroyed while DoWake() is
  // running, since the lock is not held at that time.
  std::atomic<uint8_t> wakes_pending_ = 0;
 
  // Indicates that this Dispatcher should be woken when Wake() is called. This
  // prevents unnecessary wakes when, for example, multiple wakers wake the same
  // task or multiple tasks are posted before the dipsatcher runs.
  bool wants_wake_ PW_GUARDED_BY(internal::lock()) = false;
  bool terminated_ PW_GUARDED_BY(internal::lock()) = false;
};
 
 
}  // namespace pw::async2